import org.apache.tomcat.util.codec.binary.Base64;
import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

public class RSAEncrypt {
    private static final int KEY_SIZE = 1024;
    private static final String RSA_ALGORITHM = "RSA";
    // RSA最大加密明文大小
    private static final int MAX_ENCRYPT_LENGTH = 117;
    // RSA最大解密明文大小
    private static final int MAX_DECRYPT_LENGTH = 128;

    /**
     * 随机生成密钥对
     * @return 密钥对map
     */
    public static Map<String, String> getKeyPair() {
        HashMap<String, String> keyPairMap = new HashMap<>();
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA_ALGORITHM);
            keyPairGenerator.initialize(KEY_SIZE);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
            keyPairMap.put("publicKey", encodeBase64(publicKey.getEncoded()));
            keyPairMap.put("privateKey", encodeBase64(privateKey.getEncoded()));
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return keyPairMap;
    }

    /**
     * 分段加密
     * @param str 需要加密的字符串
     * @param publicKey 加密的公钥
     */
    public static String encrypt(String str, String publicKey) {
        String outStr = null;
        try {
            byte[] publicKeyBytes = decodeBase46(publicKey);
            byte[] data = str.getBytes(StandardCharsets.UTF_8);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey pubKey = keyFactory.generatePublic(new X509EncodedKeySpec(publicKeyBytes));
            Cipher cipher = Cipher.getInstance(RSA_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, pubKey);
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            int inputLen = data.length;
            int offset = 0;
            int i = 0;
            byte[] cache;
            while (inputLen > offset) {
                int len = Math.min(inputLen - offset, MAX_ENCRYPT_LENGTH);
                cache = cipher.doFinal(data, offset, len);
                out.write(cache, 0, cache.length);
                i++;
                offset = i * MAX_ENCRYPT_LENGTH;
            }
            byte[] outBytes = out.toByteArray();
            outStr = encodeBase64(outBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return outStr;
    }

    /**
     * 分段解密
     * @param str 需要解码的字符串(是经过base64编码的)
     * @param privateKey 解码用的私钥
     */
    public static String decrypt(String str, String privateKey) {
        byte[] data = decodeBase46(str);
        byte[] privateKeyBytes = decodeBase46(privateKey);
        String outStr = null;
        try {
            KeyFactory instance = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey priKey = instance.generatePrivate(new PKCS8EncodedKeySpec(privateKeyBytes));
            Cipher cipher = Cipher.getInstance(RSA_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, priKey);
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            int inputLen = data.length;
            int offset = 0;
            int i = 0;
            byte[] cache;
            while (inputLen > offset) {
                int len = Math.min(inputLen - offset, MAX_DECRYPT_LENGTH);
                cache = cipher.doFinal(data, offset, len);
                out.write(cache, 0, cache.length);
                i++;
                offset = i * MAX_DECRYPT_LENGTH;
            }
            byte[] outBytes = out.toByteArray();
            outStr = new String(outBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return outStr;
    }

    /**
     * RSA加密
     * @param str 需要加密的字符串
     * @param publicKey 加密的公钥
     */
    public static String encryptNoSegment(String str, String publicKey) {
        String outStr = null;
        try {
            byte[] publicKeyBytes = decodeBase46(publicKey);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey pubKey = keyFactory.generatePublic(new X509EncodedKeySpec(publicKeyBytes));
            Cipher cipher = Cipher.getInstance(RSA_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, pubKey);
            byte[] bytes = cipher.doFinal(str.getBytes(StandardCharsets.UTF_8));
            outStr = encodeBase64(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return outStr;
    }

    /**
     *
     * @param str 需要解码的字符串(是经过base64编码的)
     * @param privateKey 解码用的私钥
     * @return
     */
    public static String decryptNoSegment(String str, String privateKey) {
        byte[] bytes = decodeBase46(str);
        byte[] privateKeyBytes = decodeBase46(privateKey);
        String outStr = null;
        try {
            KeyFactory instance = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey priKey = instance.generatePrivate(new PKCS8EncodedKeySpec(privateKeyBytes));
            Cipher cipher = Cipher.getInstance(RSA_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, priKey);
            outStr = new String(cipher.doFinal(bytes));
        } catch (Exception e) {
            e.printStackTrace();
        }
        return outStr;
    }

    /**
     * base64编码
     * @param encoded 需要编码的字节数组
     * @return base64字符串
     */
    private static String encodeBase64(byte[] encoded) {
        return Base64.encodeBase64String(encoded);
    }

    /**
     * base64解码成字节数组
     * @param base64String base64字符串
     * @return 解码出的字节数组
     */
    private static byte[] decodeBase46(String base64String) {
        return Base64.decodeBase64(base64String);
    }
}