package cd.casic.framework.commons.util.encrypt;

import lombok.SneakyThrows;

import javax.crypto.*;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;


public class DHUtil {

    private static final String KEY_ALGORITHM = "DH";
    private static final String KEY_PROVIDER = "BC";
    private static final String SECRET_ALGORITHM = "DES";

    private static BigInteger p = new BigInteger("16560215747140417249215968347342080587", 16);
    private static BigInteger g = new BigInteger("1234567890", 16);

    public static DHKeyPair initKey(){
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM, KEY_PROVIDER);
            DHParameterSpec serverParam = new DHParameterSpec(p, g, 128);
            keyPairGenerator.initialize(serverParam, new SecureRandom());
            //keyPairGenerator.initialize(KEY_SIZE);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            return new DHKeyPair(keyPair.getPublic().getEncoded(), keyPair.getPrivate().getEncoded());
        } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException | NoSuchProviderException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static DHKeyPair initKey(byte[] partyAPublicKey) {
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(partyAPublicKey);
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            PrivateKey publicKey = keyFactory.generatePrivate(x509KeySpec);

            DHParameterSpec dhParameterSpec = ((DHPublicKey) publicKey).getParams();
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM, KEY_PROVIDER);
            keyPairGenerator.initialize(dhParameterSpec);
//            keyPairGenerator.initialize(KEY_SIZE);
            KeyPair keyPair = keyPairGenerator.genKeyPair();
            return new DHKeyPair(keyPair.getPublic().getEncoded(), keyPair.getPrivate().getEncoded());
        } catch (NoSuchAlgorithmException | InvalidKeySpecException | NoSuchProviderException | InvalidAlgorithmParameterException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] encrypt(byte[] data, byte[] partAPublicKey, byte[] partBPrivateKey) {
        byte[] key = getSecretKey(partAPublicKey, partBPrivateKey);
        SecretKeySpec secretKey = new SecretKeySpec(key, SECRET_ALGORITHM);
        try {
            Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, secretKey);
            return cipher.doFinal(data);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException | IllegalBlockSizeException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] decrypt(byte[] data, byte[] partBPublicKey, byte[] partAPrivateKey){
        byte[] key = getSecretKey(partBPublicKey, partAPrivateKey);
        SecretKeySpec secretKey = new SecretKeySpec(key, SECRET_ALGORITHM);
        Cipher cipher = null;
        try {
            cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.DECRYPT_MODE, secretKey);
            return cipher.doFinal(data);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }


    @SneakyThrows
    private static byte[] getSecretKey(byte[] publicKey, byte[] privateKey) {
        // 实例化密钥工厂
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        // 初始化公钥
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKey);
        // 产生公钥
        PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);
        // 初始化私钥
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(privateKey);
        // 产生私钥
        PrivateKey priKey = keyFactory.generatePrivate(pkcs8KeySpec);
        // 实例化
        KeyAgreement keyAgree = KeyAgreement.getInstance(KEY_ALGORITHM, KEY_PROVIDER);
        // 初始化
        keyAgree.init(priKey);
        keyAgree.doPhase(pubKey, true);
        // 生成本地密钥
        SecretKey secretKey = keyAgree.generateSecret(SECRET_ALGORITHM);
        return secretKey.getEncoded();
    }
}
