import java.io.PrintWriter;
import java.math.BigInteger;
import java.util.Scanner;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.Base64;

import static javax.crypto.Cipher.getInstance;

/**
 * RSA算法
 */
public class Main {

    private static final BigInteger a = BigInteger.valueOf('a');
    private static final Scanner sc = new Scanner(System.in);
    private static final PrintWriter print = new PrintWriter(System.out,true);
    // 公钥
    public static BigInteger n;
    public static BigInteger e;
    // 私钥
    public static BigInteger d;

    // 帮助消息1
    private static final String help1 = "1.自己书写的 2.javaAPI生成";
    // 帮助信息2
    private static final String help = "1. 生成密钥  2. RSA加密  3. RSA解密 4.重置数据 5.输入密钥 -1. 退出";

    public static void main(String[] args) {
        while(true){
            print.println(help1);
            switch(sc.nextInt()){
                case 1 -> MyRSA();
                case 2 -> javaRSA();
            }
        }
    }
    public static void MyRSA(){
        while (true) {
            print.println("当前数据" + (n == null ? "" : " n = "+n) + (e == null ? "" : " e = "+e) + (d == null ? "" : " d = "+d));
            print.println(help);
            int i = sc.nextInt();
            switch (i) {
                case 1 -> createN();
                case 2 -> getCipher();
                case 3 -> getPlain();
                case 4 -> init();
                case 5 -> {
                    print.println("请输入密钥 n,e,d");
                    n = sc.nextBigInteger();
                    e = sc.nextBigInteger();
                    d = sc.nextBigInteger();
                }
                case -1 -> {
                    return;
                }
            }
        }
    }
    public static void javaRSA(){
        print.println("1.输入指定的公钥和私钥 2.随机生成器生成");
        RSA rsa = null;
        switch(sc.nextInt()){
            case 1->{
                print.println("输入密钥n、e和d");
                rsa = new RSA(sc.nextBigInteger(),sc.nextBigInteger(),sc.nextBigInteger());
            }
            case 2->{
                print.println("输入密钥长度,长度不小于512");
                rsa = new RSA(sc.nextInt());
            }
        }
        while(true){
            print.println("1.加密 2.解密 -1 退出 然后一行输入文本");
            switch(sc.nextInt()){
                case 1->print.println(rsa.encode(sc.next()));
                case 2 -> print.println(rsa.decode(sc.next()));
                case -1 ->{
                    return;
                }
                default -> print.println("输入不正确");
            }
        }
    }

    private static void init() {
        n = null;
        e = null;
        d = null;
    }

    private static void getPlain() {
        if (d == null) {
            print.println("请输入密钥");
            d = sc.nextBigInteger();
        }
        print.println("请输入密文");
        getModPow(d);
    }

    private static void getModPow(BigInteger d) {
        sc.next().chars().mapToObj(c -> BigInteger.valueOf(c - 'a')).forEach(i -> {
            BigInteger j = i.modPow(d, n);
            j = j.add(a);
            print.print((char) j.intValue());
        });
        print.println();
    }

    private static void createN() {
        print.println("请输入两个质数p,q");
        BigInteger p = sc.nextBigInteger();
        BigInteger q = sc.nextBigInteger();
        n = p.multiply(q);
        BigInteger f = p.subtract(BigInteger.ONE).multiply(q.subtract(BigInteger.ONE));
        e = BigInteger.TWO;
        while (!f.gcd(e).equals(BigInteger.ONE)) {
            e = e.add(BigInteger.ONE);
        }
        d = e.modInverse(f);
        print.println("公钥: (" + n + "," + e + ")");
        print.println("私钥: (" + n + "," + d + ")");
    }

    private static void getCipher() {
        if (n == null || e == null) {
            print.println("请输入公钥n，e");
            n = sc.nextBigInteger();
            e = sc.nextBigInteger();
        }
        print.println("请输入明文");
        getModPow(e);
    }
}


class RSA {
        private final PublicKey publicKey;
        private final PrivateKey privateKey;
    
        /**
         * 自定义公钥和私钥
         * @param n 公钥和私钥的n
         * @param e 公钥的e
         * @param d 私钥的d
         */
        public RSA(BigInteger n,BigInteger e,BigInteger d){
            RSAPublicKeySpec publicKeySpec = new RSAPublicKeySpec(n, e);
            RSAPrivateKeySpec privateKeySpec = new RSAPrivateKeySpec(n, d);
            try {
                KeyFactory rsa = KeyFactory.getInstance("RSA");
                publicKey = rsa.generatePublic(publicKeySpec);
                privateKey = rsa.generatePrivate(privateKeySpec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException ex) {
                throw new RuntimeException(ex);
            }
        }
    
        /**
         * 指定密钥大小生成公钥和私钥
         * @param keySize 密钥大小
         */
        public RSA(int keySize){
            if(keySize < 512){
                System.err.println("都说了不要小于512，会报错的,重新启动程序吧");
            }
            try {
                KeyPairGenerator keyGenerator = KeyPairGenerator.getInstance("RSA");
                keyGenerator.initialize(keySize);
                KeyPair keyPair = keyGenerator.generateKeyPair();
                publicKey = keyPair.getPublic();
                privateKey = keyPair.getPrivate();
                System.out.println("publicKey = " + Base64.getEncoder().encodeToString(publicKey.getEncoded()));
                System.out.println("privateKey = " + Base64.getEncoder().encodeToString(privateKey.getEncoded()));
            } catch (NoSuchAlgorithmException ex) {
                throw new RuntimeException(ex);
            }
        }
    
        /**
         * 获取密文
         * @param plainText 明文
         * @return 密文
         */
        public String encode(String plainText){
            try {
                Cipher cipher = getInstance("RSA");
                cipher.init(Cipher.ENCRYPT_MODE, publicKey);
                byte[] bytes = cipher.doFinal(plainText.getBytes());
                return Base64.getEncoder().encodeToString(bytes);
            } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException |
                     BadPaddingException ex) {
                throw new RuntimeException(ex);
            }
        }
    
        /**
         * 获取明文
         * @param cipherText 密文
         * @return 明文
         */
        public String decode(String cipherText){
            try {
                Cipher cipher = getInstance("RSA");
                cipher.init(Cipher.DECRYPT_MODE, privateKey);
                byte[] bytes = cipher.doFinal(Base64.getDecoder().decode(cipherText));
                return new String(bytes);
            } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException |
                     BadPaddingException ex) {
                throw new RuntimeException(ex);
            }
        }
}