package minsecurity.certificate.cert;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import minsecurity.Common;
import minsecurity.crypto.*;
import org.slf4j.LoggerFactory;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.Base64;

/*
 * @Author: hongyu guo
 * @Description:
 * @Version: 1.0.0
 * @Date: 17:09 2021/03/05
 * @Copyright: MIN-Group；国家重大科技基础设施——未来网络北大实验室；深圳市信息论与未来网络重点实验室
 */
public class CertUtils {
    private static final org.slf4j.Logger logger = LoggerFactory.getLogger(CertUtils.class);
    public static final int VERSION1 = 1;

    /**
     * Certificate -> InnerCertificate
     * @param certificate
     * @return InnerCertificate
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static InnerCertificate parseCertToInnerCert(Certificate certificate){
        TbsCertificate tbsCertificate = new TbsCertificate();
        tbsCertificate.setVersion(certificate.getVersion());
        tbsCertificate.setSerialNumber(certificate.getSerialNumber());
        tbsCertificate.setPublicKey(certificate.getPublicKey().getBytes());
        tbsCertificate.setSignatureAlgorithm(certificate.getSignatureAlgorithm());
        tbsCertificate.setPublicKeyAlgorithm(certificate.getPublicKeyAlgorithm());
        tbsCertificate.setIssueTo(certificate.getIssueTo());
        tbsCertificate.setIssuer(certificate.getIssuer());
        tbsCertificate.setNotAfter(certificate.getNotAfter());
        tbsCertificate.setNotBefore(certificate.getNotBefore());
        tbsCertificate.setKeyUsage(certificate.getKeyUsage());
        tbsCertificate.setCA(certificate.isCA());
        tbsCertificate.setTimestamp(certificate.getTimestamp());

        InnerCertificate innerCertificate = new InnerCertificate();
        innerCertificate.setSignatureAlgorithm(certificate.getSignatureAlgorithm());
        innerCertificate.setSignatureValue(certificate.getSignature());
        innerCertificate.setTbsCertificate(tbsCertificate);
        return innerCertificate;
    }


    /**
     * InnerCertificate -> Certificate
     * @param i
     * @return minsecurity.certificate.cert.Certificate
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static Certificate parseInnerCertToCert(InnerCertificate i){
        Certificate certificate = new Certificate();

        certificate.setVersion(i.getTbsCertificate().getVersion());
        certificate.setSerialNumber(i.getTbsCertificate().getSerialNumber());
        PublicKeyInterface publicKey = KeyUtils.unMarshalPublicKey(i.getTbsCertificate().getPublicKey(),i.getSignatureAlgorithm());
        certificate.setPublicKey(publicKey);
        certificate.setSignature(i.getSignatureValue());
        certificate.setSignatureAlgorithm(i.getTbsCertificate().getSignatureAlgorithm());
        certificate.setPublicKeyAlgorithm(i.getTbsCertificate().getPublicKeyAlgorithm());
        certificate.setIssueTo(i.getTbsCertificate().getIssueTo());
        certificate.setIssuer(i.getTbsCertificate().getIssuer());
        certificate.setNotBefore(i.getTbsCertificate().getNotBefore());
        certificate.setNotAfter(i.getTbsCertificate().getNotAfter());
        certificate.setKeyUsage(i.getTbsCertificate().getKeyUsage());
        certificate.setCA(i.getTbsCertificate().isCA());
        certificate.setTimestamp(i.getTbsCertificate().getTimestamp());

        return certificate;
    }

    /**
     * 检查证书是否在有效期内
     * @param innerCertificate
     * @return boolean
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static boolean checkDuration(InnerCertificate innerCertificate){
        long curTime = System.currentTimeMillis();
        // logger.debug("current time: {}", curTime);
        if(innerCertificate.getTbsCertificate().getNotAfter() < curTime)
            return false;
        if(innerCertificate.getTbsCertificate().getNotBefore() > curTime)
            return false;
        return true;
    }


    /**
     * 使用publicKey验证签名
     * @param cert
	 * @param publicKey
     * @return boolean
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static boolean checkSign(InnerCertificate cert, PublicKeyInterface publicKey) throws Exception {
        switch (cert.getTbsCertificate().getSignatureAlgorithm()){
            case Common.SM3withSM2:
                byte[] bytesOfCert = serialization(cert.getTbsCertificate());
                return publicKey.verify(bytesOfCert, cert.getSignatureValue());
            default:
                throw new Exception("未定义的签名方法");
        }
    }

    /**
     * 将TbsCertificate类型实例转化为字节流
     * @param tbsCertificate
     * @return byte[]
     * @throws CertException
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static byte[] serialization(TbsCertificate tbsCertificate) throws CertException {
        ObjectMapper mapper = new ObjectMapper();
        byte[] bytesOfCertJson;
        try {
            bytesOfCertJson = mapper.writeValueAsBytes(tbsCertificate);
            return bytesOfCertJson;
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        throw new CertException("tbsCertificate serialization failed");
    }

    /**
     * 将InnerCertificate类型实例转化为字节流
     * @param innerCertificate
     * @return byte[]
     * @throws CertException
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static byte[] serialization(InnerCertificate innerCertificate) throws CertException {
        ObjectMapper mapper = new ObjectMapper();
        byte[] bytesOfCertJson;
        try {
            bytesOfCertJson = mapper.writeValueAsBytes(innerCertificate);
            return bytesOfCertJson;
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        throw new CertException("innerCertificate serialization failed");
    }

    /**
     * 根据typeOfCert可将字节流反序列化为TbsCertificate或InnerCertificate,
     * 返回Object类型, 调用者调用时需要手动强制类型转化
     * @param bytesOfCert 字节流
	 * @param typeOfCert Common.TbsCertificate或Common.InnerCertificate
     * @return java.lang.Object
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static Object deserialization(byte[] bytesOfCert, int typeOfCert){
        ObjectMapper mapper = new ObjectMapper();
        try {
            switch (typeOfCert) {
                case Common.TbsCertificate:
                    return mapper.readValue(bytesOfCert, TbsCertificate.class);
                case Common.InnerCertificate:
                    return mapper.readValue(bytesOfCert, InnerCertificate.class);
                default:
                    throw new CertException("未知证书类型");
            }
        } catch (IOException | CertException e) {
            e.printStackTrace();
        }
        return null;
    }


    /**
     * 使用priv私钥给Certificate类型实例签名
     * @param certificate 未签名证书
	 * @param priv 私钥
     * @return void
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static void signCert(Certificate certificate, PrivateKeyInterface priv) throws Exception {
        byte[] bytesOfCert = serialization(CertUtils.parseCertToInnerCert(certificate).getTbsCertificate());
        // logger.debug("签名序列化bytes: {}", ByteUtils.toHexString(bytesOfCert));
        switch (certificate.getSignatureAlgorithm()){
            case Common.SM3withSM2:
                certificate.setSignature(priv.sign(bytesOfCert));
                break;
            default:
                logger.debug("signAlgo: {}", certificate.getSignatureAlgorithm());
                throw new Exception("未定义的签名方法");
        }
    }

    /**
     * 使用ca的公钥对sub进行验签, 如果时根证书验签, 调用时ca与sub指向同一个实例
     * @param ca
	 * @param sub
     * @return boolean
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static boolean verifyCert(Certificate ca, Certificate sub) throws Exception {
        InnerCertificate innerCertificate = CertUtils.parseCertToInnerCert(sub);
        switch (innerCertificate.getTbsCertificate().getVersion()){
            case CertUtils.VERSION1:
                boolean isValid = CertUtils.checkDuration(innerCertificate);
                // logger.debug("有效期内：{}", isValid);
                if(!isValid)
                    return false;
                if(sub.isCA()){ // 根证书, 自验证
                    return checkSign(innerCertificate, sub.getPublicKey());
                } else { // 使用ca验证sub
                    return checkSign(innerCertificate, ca.getPublicKey());
                }
            default:
                return false;
        }
    }

    /**
     * 证书序列化为String类型, 密码为null或为空时不对证书进行加密保存
     * @param certificate
	 * @param passwd
	 * @param symAlgoMode Common.SM4CBC / Common.SM4ECB
     * @return java.lang.String
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static String toPem(Certificate certificate, byte[] passwd, int symAlgoMode) throws CertException, NoSuchPaddingException, InvalidAlgorithmParameterException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, NoSuchProviderException, InvalidKeyException {
        if(certificate.getSignature() == null){
            throw new CertException("该证书未签名");
        }
        byte[] bytesOfCert = serialization(CertUtils.parseCertToInnerCert(certificate));
        byte[] ret;
        if(passwd != null && passwd.length > 0){
            // byte[] sm4Key;
            byte[] sm4Key = KeyUtils.get16bytePasswd(passwd);

            switch (symAlgoMode){
                case Common.SM4CBC:{
                    // TODO: PADDING???  IV?????
                    ret = SM4.encrypt_CBC_Padding(sm4Key, new byte[16],bytesOfCert);
                    break;
                }
                case Common.SM4ECB:{
                    ret = SM4.encrypt_ECB_Padding(sm4Key, bytesOfCert);
                    break;
                }
                default:
                    throw new CertException("未定义的对称加密算法");
            }
        } else {
            ret = bytesOfCert;
        }
        return Base64.getEncoder().encodeToString(ret);
    }

    /**
     * String类型证书转化为Certificate对象
     * @param pemString
	 * @param passwd
	 * @param symAlgoMode Common.SM4CBC / Common.SM4ECB
     * @return minsecurity.certificate.cert.Certificate
     * @author hongyu guo
     * @date 2021/3/7
    **/
    public static Certificate fromPem(String pemString, byte[] passwd, int symAlgoMode)
            throws CertException, BadPaddingException, NoSuchPaddingException,
                   InvalidAlgorithmParameterException, NoSuchAlgorithmException,
                   IllegalBlockSizeException, NoSuchProviderException, InvalidKeyException {
        if(pemString == null || pemString.equals("")){
            throw new CertException("pem串为空");
        }
        byte[] bytesOfPem = Base64.getDecoder().decode(pemString);
        byte[] bytesOfCert;
        if(passwd != null && passwd.length > 0){
            byte[] sm4Key = KeyUtils.get16bytePasswd(passwd);
            switch (symAlgoMode){
                case Common.SM4CBC:
                    bytesOfCert =  SM4.decrypt_CBC_Padding(sm4Key,new byte[16], bytesOfPem);
                    break;
                case Common.SM4ECB:
                    bytesOfCert =  SM4.decrypt_ECB_Padding(sm4Key, bytesOfPem);
                    break;
                default:
                    throw new CertException("未知的对称加密算法");
            }
        } else {
            bytesOfCert = bytesOfPem;
        }
        InnerCertificate innerCertificate = (InnerCertificate) deserialization(bytesOfCert, Common.InnerCertificate);
        return CertUtils.parseInnerCertToCert(innerCertificate);
    }
}