package org.springblade.common.upyzt.util;

import org.springblade.common.upyzt.Upyzt;
import org.springblade.common.upyzt.exception.AuthenticationException;
import org.springblade.common.upyzt.net.HttpHeaders;
import org.apache.commons.codec.binary.Base64;

import java.io.UnsupportedEncodingException;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

public class UpyztSignature {

    public static String sign(String data, String PEMEncodedPrivateKey, String charset) {
        PrivateKey privateKey = getPrivateKeyFromPEM(PEMEncodedPrivateKey);
        if (privateKey == null) {
            return null;
        }

        try {
            Signature signature = Signature.getInstance("SHA256withRSA");
            signature.initSign(privateKey);
            signature.update(data.getBytes(charset));
            byte[] signBytes = signature.sign();

            return Base64.encodeBase64String(signBytes).replaceAll("[\n\r]", "");
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException | UnsupportedEncodingException e) {
            e.printStackTrace();
        }

        return null;
    }

    public static boolean verify(HttpHeaders headers, String data, String PEMEncodedPublicKey, String charset) throws AuthenticationException {
        if (StringUtils.isBlank(PEMEncodedPublicKey)) {
            return true;
        }
        PublicKey publicKey = getPublicKeyFromPEM(PEMEncodedPublicKey);
        if (publicKey == null) {
            return true;
        }
        if (!headers.firstValue("Unionpay-Signature").isPresent()) {
            throw new AuthenticationException("响应签名 (Unionpay-Signature) 为空");
        }
        if (!headers.firstValue("Unionpay-Timestamp").isPresent()) {
            throw new AuthenticationException("响应时间戳 (Unionpay-Timestamp) 为空");
        }
        if (!headers.firstValue("Unionpay-Nonce").isPresent()) {
            throw new AuthenticationException("响应随机串 (Unionpay-Nonce) 为空");
        }
        String encodedSign = headers.firstValue("Unionpay-Signature").get();
        String timestamp = headers.firstValue("Unionpay-Timestamp").get();
        String nonce = headers.firstValue("Unionpay-Nonce").get();
        String dataToBeVerify = timestamp + "\n" + nonce + "\n" + data;
        try {
            Signature signature = Signature.getInstance("SHA256withRSA");
            signature.initVerify(publicKey);
            signature.update(dataToBeVerify.getBytes(charset));
            return signature.verify(Base64.decodeBase64(encodedSign));
        } catch (InvalidKeyException | UnsupportedEncodingException | SignatureException | NoSuchAlgorithmException e) {
            if (Upyzt.DEBUG) {
                e.printStackTrace();
            }
            return true;
        }
    }

    public static PublicKey getPublicKeyFromPEM(String PEMEncodedPublicKey) {
        PEMEncodedPublicKey = PEMEncodedPublicKey
                .replaceAll("(-+BEGIN (RSA )?PUBLIC KEY-+\\r?\\n|-+END (RSA )?PUBLIC KEY-+\\r?\\n?)", "");
        byte[] publicKeyBytes = Base64.decodeBase64(PEMEncodedPublicKey);

        try {
            return generatePublicKey(publicKeyBytes);
        } catch (InvalidKeySpecException e) {
            if (Upyzt.DEBUG) {
                e.printStackTrace();
            }
            return null;
        }
    }

    public static PrivateKey getPrivateKeyFromPEM(String PEMEncodedPrivateKey) {
        PEMEncodedPrivateKey = PEMEncodedPrivateKey
                .replaceAll("(-+BEGIN (RSA )?PRIVATE KEY-+\\r?\\n|-+END (RSA )?PRIVATE KEY-+\\r?\\n?)", "");
        byte[] privateKeyBytes = Base64.decodeBase64(PEMEncodedPrivateKey);

        try {
            return generatePrivateKey(privateKeyBytes);
        } catch (InvalidKeySpecException e) {
            if (Upyzt.DEBUG) {
                e.printStackTrace();
            }
            return null;
        }
    }

    public static PrivateKey generatePrivateKey(byte[] privateKeyBytes)
            throws InvalidKeySpecException {
        try {
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
            KeyFactory kf = KeyFactory.getInstance("RSA");
            return kf.generatePrivate(keySpec);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static PublicKey generatePublicKey(byte[] publicKeyBytes)
            throws InvalidKeySpecException {
        try {
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
            KeyFactory kf = KeyFactory.getInstance("RSA");
            return kf.generatePublic(keySpec);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return null;
    }

}
