package com.nova.sankuai.infra.utils.rongbei;

import lombok.SneakyThrows;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.security.KeyFactory;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA 支持类
 */
public class RsaSupport {
    /**
     * 默认值-密钥key算法
     */
    public static final String DEFAULT_KEY_ALGORITHM = "RSA";
    /**
     * 默认值-签名算法
     */
    public static final String DEFAULT_SIGNATURE_ALGORITHM = "SHA256WithRSA";
    /**
     * 默认值-加解密算法
     */
    public static final String DEFAULT_CIPHER_ALGORITHM = "RSA/ECB/PKCS1Padding";
    /**
     * 默认值-默认字符编码
     */
    public static final String DEFAULT_CHARSET_NAME = "utf-8";

    /**
     * 默认值-加密块长度
     */
    private final int DEFAULT_MAX_ENCRYPT_BLOCK = 117;

    private final String KEY_ALGORITHM;
    private final String SIGNATURE_ALGORITHM;
    private final String CIPHER_ALGORITHM;
    private final String CHARSET_NAME;

    /**
     * rsa-公钥
     */
    private String pukStr;
    /**
     * rsa-私钥
     */
    private String prkStr;

    public RsaSupport() {
        this.KEY_ALGORITHM = DEFAULT_KEY_ALGORITHM;
        this.SIGNATURE_ALGORITHM = DEFAULT_SIGNATURE_ALGORITHM;
        this.CIPHER_ALGORITHM = DEFAULT_CIPHER_ALGORITHM;
        this.CHARSET_NAME = DEFAULT_CHARSET_NAME;
    }

    public RsaSupport(String keyAlgorithm, String signatureAlgorithm, String cipherAlgorithm, String charsetName) {
        this.KEY_ALGORITHM = keyAlgorithm;
        this.SIGNATURE_ALGORITHM = signatureAlgorithm;
        this.CIPHER_ALGORITHM = cipherAlgorithm;
        this.CHARSET_NAME = charsetName;
    }

    public RsaSupport(String keyAlgorithm, String signatureAlgorithm, String cipherAlgorithm, String charsetName,
                      String pukStr, String prkStr) {
        this.KEY_ALGORITHM = keyAlgorithm;
        this.SIGNATURE_ALGORITHM = signatureAlgorithm;
        this.CIPHER_ALGORITHM = cipherAlgorithm;
        this.CHARSET_NAME = charsetName;
        this.pukStr = pukStr;
        this.prkStr = prkStr;
    }

    /**
     * 使用私钥对数据进行签名
     *
     * @param data
     *            需要签名的数据
     * @return 返回加签名后的BASE64编码的字符串
     * @throws Exception
     */
    public String signByPrivateKey(String data) {
        return signByPrivateKey(data, prkStr);
    }

    /**
     * 使用私钥对数据进行签名
     *
     * @param data
     *            需要签名的数据
     * @param prkStr
     *            私钥（使用BASE64进行编码）
     * @return 返回加签名后的BASE64编码的字符串
     * @throws Exception
     */
    @SneakyThrows
    public String signByPrivateKey(String data, String prkStr) {
        if (prkStr != null) {
            prkStr = prkStr.replaceAll(" ", "");
        }
        byte[] privateKeyByte = java.util.Base64.getDecoder().decode(prkStr);
        // 私钥格式，从字节数组中构建的私钥内容
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyByte);
        // 密钥工厂
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        // 通过密钥工厂，从指定的密钥中生成RSA格式的私钥对象
        RSAPrivateKey privateKey = (RSAPrivateKey)keyFactory.generatePrivate(keySpec);
        // 签名算法
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
        // 签名算法采用私钥初始化
        signature.initSign(privateKey);
        // 将字符串编码为指定字符集字符数组，此处传入被签名的内容
        signature.update(data.getBytes(CHARSET_NAME));
        // 签名，并将签名后的内容编码为base64的字符串返回
        return java.util.Base64.getEncoder().encodeToString(signature.sign());
    }

    /**
     * 使用公钥进行数字签名的校验
     *
     * @param data
     *            需要进行验签的原始数据
     * @param sign
     *            签名 （BASE64编码字符串）
     * @return
     * @throws Exception
     */
    public boolean verifyByPublicKey(String data, String sign) {
        return verifyByPublicKey(data, sign, pukStr);
    }

    /**
     * 使用公钥进行数字签名的校验
     *
     * @param data
     *            需要进行验签的原始数据
     * @param pukStr
     *            公钥（BASE64编码字符串）
     * @param sign
     *            签名 （BASE64编码字符串）
     * @return
     * @throws Exception
     */
    public boolean verifyByPublicKey(String data, String sign, String pukStr) {
        try {
            if (pukStr != null) {
                pukStr = pukStr.replaceAll(" ", "");
            }
            // 公钥字符串解码
            byte[] publicKeyByte = java.util.Base64.getDecoder().decode(pukStr);
            // x509格式结构的公钥
            X509EncodedKeySpec encodedKeySpec = new X509EncodedKeySpec(publicKeyByte);
            // RSA算法之密钥生成器
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 从结构性数据，转换为公钥对象
            RSAPublicKey publicKey = (RSAPublicKey)keyFactory.generatePublic(encodedKeySpec);
            // 签名算法
            Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
            // 初始化验证用的公钥
            signature.initVerify(publicKey);
            // 更新需要要整的内容
            signature.update(data.getBytes(CHARSET_NAME));
            // 判断是否是对应的签名
            return signature.verify(Base64.getDecoder().decode(sign));
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }

    /**
     * 公钥加密
     *
     * @param plainText
     * @return
     * @throws Exception
     */
    public String encryptByPublicKey(String plainText) {
        return encryptByPublicKey(plainText, pukStr);
    }

    /**
     * 公钥加密
     *
     * @param plainText
     * @param pukStr
     * @return
     * @throws Exception
     */
    @SneakyThrows
    public String encryptByPublicKey(String plainText, String pukStr) {
        if (null == plainText) {
            return null;
        }

        if (pukStr != null) {
            pukStr = pukStr.replaceAll(" ", "");
        }

        byte[] dataB = plainText.getBytes(CHARSET_NAME);
        byte[] publicKeyByte = Base64.getDecoder().decode(pukStr);
        X509EncodedKeySpec encodedKeySpec = new X509EncodedKeySpec(publicKeyByte);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        RSAPublicKey publicKey = (RSAPublicKey)keyFactory.generatePublic(encodedKeySpec);
        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        int inputLen = dataB.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段加密
        while (inputLen - offSet > 0) {
            if (inputLen - offSet > DEFAULT_MAX_ENCRYPT_BLOCK) {
                cache = cipher.doFinal(dataB, offSet, DEFAULT_MAX_ENCRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(dataB, offSet, inputLen - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * DEFAULT_MAX_ENCRYPT_BLOCK;
        }
        byte[] encryptedData = out.toByteArray();
        out.close();
        return Base64.getEncoder().encodeToString(encryptedData);
    }

    /**
     * 私钥解密
     *
     * @param encryptedStr
     * @return
     * @throws Exception
     */
    public String decryptByPrivateKey(String encryptedStr) {
        return decryptByPrivateKey(encryptedStr, prkStr);
    }

    /**
     * 私钥解密
     *
     * @param encryptedStr
     * @param prkStr
     * @return
     * @throws Exception
     */
    @SneakyThrows
    public String decryptByPrivateKey(String encryptedStr, String prkStr) {
        if (null == encryptedStr) {
            return null;
        }

        if (prkStr != null) {
            prkStr = prkStr.replaceAll(" ", "");
        }

        byte[] dataB = Base64.getDecoder().decode(encryptedStr);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        KeySpec keySpec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(prkStr));
        RSAPrivateKey privateKey = (RSAPrivateKey)keyFactory.generatePrivate(keySpec);
        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        // 模长
        int key_len = privateKey.getModulus().bitLength() / 8;
        byte[] decryptedData = null;
        ByteArrayOutputStream bout = new ByteArrayOutputStream();
        try {
            int dataLength = dataB.length;
            for (int i = 0; i < dataLength; i += key_len) {
                int decryptLength = dataLength - i < key_len ? dataLength - i : key_len;
                byte[] doFinal = cipher.doFinal(dataB, i, decryptLength);
                bout.write(doFinal);
            }
            decryptedData = bout.toByteArray();
        } finally {
            if (bout != null) {
                bout.close();
            }
        }
        return new String(decryptedData);
    }
}