/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package utils;

import hsm.HSMManager;
import hsm.HSMManagerImp;
import hsm.HSMUtils;
import hsmexampleconsole.HSMExampleConsole;
import iaik.pkcs.pkcs11.objects.ECDSAPrivateKey;
import iaik.pkcs.pkcs11.objects.ECDSAPublicKey;
import iaik.pkcs.pkcs11.objects.KeyPair;
import iaik.pkcs.pkcs11.wrapper.PKCS11Constants;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Security;
import java.security.interfaces.ECPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;
import java.util.List;
import javax.xml.bind.DatatypeConverter;
import objects.Keys;
import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.DEROctetString;
import org.bouncycastle.asn1.DERSequenceGenerator;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x9.X962Parameters;
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.ejbca.cvc.AlgorithmUtil;
import org.ejbca.cvc.CAReferenceField;
import org.ejbca.cvc.CVCPublicKey;
import org.ejbca.cvc.CVCertificate;
import org.ejbca.cvc.CVCertificateBody;
import org.ejbca.cvc.HolderReferenceField;
import org.ejbca.cvc.OIDField;
import org.ejbca.cvc.PublicKeyEC;
import org.ejbca.cvc.util.BCECUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * @author TuoiCM
 */
public class ExampleHSM {

    //<editor-fold defaultstate="collapsed" desc="VARIABLES">
    private static final Logger logger = LoggerFactory.getLogger(HSMExampleConsole.class);
    private static HSMManager hsmManager;
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="INIT">
    public static void init() {
        if (hsmManager == null) {
            try {
                hsmManager = HSMManagerImp.getInstance(
                        Configuration.getInstance().getHsmModule(),
                        HSMUtils.loadP11Wrapper(),
                        Configuration.getInstance().getHsmSlot(),
                        new String(DatatypeConverter.parseBase64Binary(Configuration.getInstance().getHsmPin())));
                logger.info("SUCCESSFULLY CONNECT TO HSM");
            } catch (Exception e) {
                logger.error("ERROR WHILE CONNECTING TO HSM. DETAILS: " + Utils.printStackTrace(e));
            }
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="GET HSM MANAGER">
    public static HSMManager getHsmManager() {
        return hsmManager;
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="DESTROY">
    public static void destroy() {
        if (hsmManager != null) {
            hsmManager.disconnectHSM();
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="GET LIST KEY">
    public static Keys[] getKeys() {
        try {
            if (null == hsmManager) {
                logger.error("HSM MANAGER IS NULL");
                return null;
            }
            List<ECDSAPrivateKey> ecdsaPrivateKey = hsmManager.listECKeys();
            Keys[] keys = new Keys[ecdsaPrivateKey.size()];
            for (int i = 0; i < ecdsaPrivateKey.size(); i++) {
                byte[] keyIdBytes = ecdsaPrivateKey.get(i).getId().getByteArrayValue();
                char[] keyLabelChars = ecdsaPrivateKey.get(i).getLabel().getCharArrayValue();

                String keyId = null;
                String keyLabel = null;

                if (keyIdBytes != null) {
                    keyId = new String(keyIdBytes);
                }

                if (keyLabelChars != null) {
                    keyLabel = new String(keyLabelChars);
                }

                Keys k = new Keys(keyId, keyLabel);
                keys[i] = k;
            }
            return keys;

        } catch (Exception e) {
            logger.error("GET LIST KEY HSM FAIL " + Utils.printStackTrace(e));
            return null;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="GEN KEY">
    public static Keys[] generateKeyPari(String curveName, String keyId, String keyLable) {
        try {
            if (null == hsmManager) {
                logger.error("HSM MANAGER IS NULL");
                return null;
            }
            String curveOID = null;
            if (Utils.isNullOrEmpty(curveName)) {
                curveOID = CurveName.DEFAULT_CURVE_OID;
            } else {
                curveOID = CurveName.getOID(curveName);
                if (Utils.isNullOrEmpty(curveOID)) {
                    logger.warn("INVALID CURVE NAME");
                    return null;
                }
            }
            String generatedKeyIDAndLabel = Utils.generateUUID();
            if (Utils.isNullOrEmpty(keyId)) {
                keyId = generatedKeyIDAndLabel;
            } else {
                if (hsmManager.idExists(keyId)) {
                    logger.error("KEY ID EXISTED");
                    return null;
                }
            }

            if (Utils.isNullOrEmpty(keyLable)) {
                keyLable = generatedKeyIDAndLabel;
            } else {
                if (hsmManager.labelExists(keyLable)) {
                    logger.error("KEY LABEL EXISTED");
                    return null;
                }
            }

            ASN1ObjectIdentifier curveId = new ASN1ObjectIdentifier(curveOID);
            KeyPair keypair = hsmManager.genECDSAKeyPair(keyId, keyLable, curveId);
            if (null == keypair) {
                return null;
            }
            Keys[] keys = new Keys[1];
            Keys key = new Keys(keyId, keyLable);
            keys[0] = key;
            return keys;

        } catch (Exception ex) {
            logger.error("GENERATE KEY PAIR FAIL " + Utils.printStackTrace(ex));
            return null;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="CVC REQUEST">
    public static String cvcRequest(String country, String mnemonic,
            String sequence, String signatureAlg,
            String keyId, String keyLabel) {
        try {
            if (null == hsmManager) {
                logger.error("HSM MANAGER IS NULl");
                return null;
            }
            Security.addProvider(new BouncyCastleProvider());
            if (Utils.isNullOrEmpty(keyId) && Utils.isNullOrEmpty(keyLabel)) {
                logger.error("KEY ID && KEY LABEL IS NULL");
                return null;
            }
            if (Utils.isNullOrEmpty(country)) {
                logger.error("COUNTRY IS NULL");
                return null;
            }
            if (Utils.isNullOrEmpty(sequence)) {
                logger.error("SEQUENCE IS NULL");
                return null;
            }
            if (Utils.isNullOrEmpty(signatureAlg)) {
                signatureAlg = SignatureAlgorithm.DEFAULT_SIG_ALGO;
            } else {
                signatureAlg = SignatureAlgorithm.getSignatureAlgorithm(signatureAlg);
                if (Utils.isNullOrEmpty(signatureAlg)) {
                    signatureAlg = SignatureAlgorithm.DEFAULT_SIG_ALGO;
                }
            }

            boolean findByLabel = false;
            List<ECDSAPrivateKey> ecdsaPrivateKey = null;
            List<ECDSAPublicKey> ecdsaPublicKey = null;

            if (!Utils.isNullOrEmpty(keyId)) {
                ecdsaPrivateKey = hsmManager.getECKeyByID(keyId);
            }

            if (Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                if (!Utils.isNullOrEmpty(keyLabel)) {
                    ecdsaPrivateKey = hsmManager.getECKeyByLabel(keyLabel);
                    findByLabel = true;
                }
            }

            if (Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                logger.error("NO KEY ID FOUND");
                return null;
            }

            if (findByLabel) {
                ecdsaPublicKey = hsmManager.getPublicECKeyByLabel(keyLabel);
            } else {
                ecdsaPublicKey = hsmManager.getPublicECKeyByID(keyId);
            }

            if (Utils.isNullOrEmpty(ecdsaPublicKey)) {
                logger.error("NO PUBLIC KEY ID FOUND");
                return null;
            }

            CAReferenceField caRef = null;
            HolderReferenceField holderRef = new HolderReferenceField(country, mnemonic, sequence);

            OIDField oid = AlgorithmUtil.getOIDField(signatureAlg);

            ECDSAPrivateKey signingKey = ecdsaPrivateKey.get(0);
            ECDSAPublicKey publicKey = ecdsaPublicKey.get(0);

            byte[] encodedAlgorithmIdParameters = publicKey.getEcdsaParams().getByteArrayValue();
            byte[] encodedPoint = DEROctetString.getInstance(
                    publicKey.getEcPoint().getByteArrayValue()).getOctets();
            ECPublicKey ecPublicKey = createECPublicKey(encodedAlgorithmIdParameters, encodedPoint);
            CVCPublicKey cvcPublicKey = new PublicKeyEC(oid, ecPublicKey, null);
            CVCertificateBody reqBody = new CVCertificateBody(caRef, cvcPublicKey, holderRef);
            CVCertificate cvc = new CVCertificate(reqBody);
            byte[] dataToBeSigned = cvc.getTBS();

            //hash the message
            MessageDigest md = MessageDigest.getInstance(SignatureAlgorithm.getHashAlgorithm(signatureAlg));
            md.update(dataToBeSigned);

            byte[] hashedMessage = md.digest();
            byte[] signature = hsmManager.sign(PKCS11Constants.CKM_ECDSA, hashedMessage, signingKey);

            if (null == signature) {
                logger.error("SIGNATURE IS NULL");
                return null;
            }
            byte[] r = new byte[signature.length / 2];
            byte[] s = new byte[signature.length / 2];

            System.arraycopy(signature, 0, r, 0, signature.length / 2);
            System.arraycopy(signature, signature.length / 2, s, 0, signature.length / 2);

            BigInteger[] bigSignature = new BigInteger[2];
            bigSignature[0] = new BigInteger(1, r);
            bigSignature[1] = new BigInteger(1, s);

            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            DERSequenceGenerator seq = new DERSequenceGenerator(byteArrayOutputStream);
            seq.addObject(new ASN1Integer(bigSignature[0]));
            seq.addObject(new ASN1Integer(bigSignature[1]));
            seq.close();
            byte[] finalSignature = byteArrayOutputStream.toByteArray();
            // Now convert the X9.62 signature to a CVC signature
            byte[] sig = BCECUtil.convertX962SigToCVC(signatureAlg, finalSignature);
            // Create and return the CVCRequest (which is an instance of CVCertificate)
            cvc.setSignature(sig);
            String cvcRequest = DatatypeConverter.printBase64Binary(cvc.getDEREncoded());
            //LOG.info("signature: " + DatatypeConverter.printHexBinary(signature));
            //LOG.info("r: " + DatatypeConverter.printHexBinary(r));
            //LOG.info("s: " + DatatypeConverter.printHexBinary(s));
            //LOG.info("finalSignature: " + DatatypeConverter.printHexBinary(finalSignature));
            //LOG.info("finalSignature: " + DatatypeConverter.printBase64Binary(finalSignature));
            //LOG.info("cvcRequest: " + cvcRequest);
            return cvcRequest;
        } catch (Exception ex) {
            logger.error("CVC REQUEST FAIL " + Utils.printStackTrace(ex));
            return null;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="DELETE KEYS">
    private boolean deleteKeys(String keyId, String keyLabel) {
        try {
            if (null == hsmManager) {
                logger.error("HSM MANAGER IS NULL");
                return false;
            }
            if (Utils.isNullOrEmpty(keyId) && Utils.isNullOrEmpty(keyLabel)) {
                logger.error("KEY ID && KEY LABEL IS NULL");
                return false;
            }
            boolean findByLabel = false;
            List<ECDSAPrivateKey> ecdsaPrivateKey = null;
            List<ECDSAPublicKey> ecdsaPublicKey = null;

            if (!Utils.isNullOrEmpty(keyId)) {
                ecdsaPrivateKey = hsmManager.getECKeyByID(keyId);
            }

            if (Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                if (!Utils.isNullOrEmpty(keyLabel)) {
                    ecdsaPrivateKey = hsmManager.getECKeyByLabel(keyLabel);
                    findByLabel = true;
                }
            }

            if (findByLabel) {
                ecdsaPublicKey = hsmManager.getPublicECKeyByLabel(keyLabel);
            } else {
                ecdsaPublicKey = hsmManager.getPublicECKeyByID(keyId);
            }
            if (Utils.isNullOrEmpty(ecdsaPrivateKey) && Utils.isNullOrEmpty(ecdsaPublicKey)) {
                logger.error("NO KEY ID FOUND");
                return false;
            } else {
                if (!Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                    for (ECDSAPrivateKey k : ecdsaPrivateKey) {
                        hsmManager.deleteKey(k);
                    }
                }

                if (!Utils.isNullOrEmpty(ecdsaPublicKey)) {
                    for (ECDSAPublicKey k : ecdsaPublicKey) {
                        hsmManager.deleteKey(k);
                    }
                }
                return true;
            }

        } catch (Exception ex) {
            logger.error("DELETE KEYS FAIL " + Utils.printStackTrace(ex));
            return false;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="GET TA SIGNATURE">
    public static String getTASginature(String keyId, String keyLabel,
            String hash, String encoding) {
        try {
            if (null == hsmManager) {
                logger.error("HSM MANAGER IS NULL");
                return null;
            }
            if(Utils.isNullOrEmpty(keyId) && Utils.isNullOrEmpty(keyLabel)) {
                logger.error("KEY ID && KEY LABEL IS NULL");
                return null;
            }
            if(Utils.isNullOrEmpty(hash)) {
                logger.error("HASH IS NULL");
                return null;
            }
            
            byte[] dataToBeSigned = null;
            try {
                if (Utils.isNullOrEmpty(encoding)) {
                    dataToBeSigned = DatatypeConverter.parseBase64Binary(hash);
                } else {
                    if (encoding.equalsIgnoreCase(Constant.ENCODING_BASE64)) {
                        dataToBeSigned = DatatypeConverter.parseBase64Binary(hash);
                    } else {
                        dataToBeSigned = DatatypeConverter.parseHexBinary(hash);
                    }
                }
            } catch (Exception e) {
                throw e;
            }
            
            if(null == dataToBeSigned) {
                logger.error("DATA TO BE SIGNED IS NULL");
                return null;
            }
            
            List<ECDSAPrivateKey> ecdsaPrivateKey = null;
            if (!Utils.isNullOrEmpty(keyId)) {
                ecdsaPrivateKey = hsmManager.getECKeyByID(keyId);
            }

            if (Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                if (!Utils.isNullOrEmpty(keyLabel)) {
                    ecdsaPrivateKey = hsmManager.getECKeyByLabel(keyLabel);
                }
            }
            
            if(Utils.isNullOrEmpty(ecdsaPrivateKey)) {
                logger.error("NO KEY ID FOUND");
                return null;
            }
            ECDSAPrivateKey signingKey = ecdsaPrivateKey.get(0);
            byte[] signature = hsmManager.sign(PKCS11Constants.CKM_ECDSA, dataToBeSigned, signingKey);
            if(null == signature) {
                logger.error("SIGNATURE IS NULL");
                return null;
            }
            return DatatypeConverter.printBase64Binary(signature);
            
        } catch (Exception ex) {
            logger.error("GET TA SIGNATURE FAIL " + ex);
            return null;
        }
    }
    //</editor-fold>

    //<editor-fold defaultstate="collapsed" desc="CREATE EC PUBLIC KEY">
    private static ECPublicKey createECPublicKey(byte[] encodedAlgorithmIdParameters,
            byte[] encodedPoint)
            throws InvalidKeySpecException, IOException {
        ASN1Encodable algParams;
        if (encodedAlgorithmIdParameters[0] == 6) {
            algParams = ASN1ObjectIdentifier.getInstance(encodedAlgorithmIdParameters);
        } else {
            algParams = X962Parameters.getInstance(encodedAlgorithmIdParameters);
        }
        AlgorithmIdentifier algId = new AlgorithmIdentifier(
                X9ObjectIdentifiers.id_ecPublicKey, algParams);

        SubjectPublicKeyInfo spki = new SubjectPublicKeyInfo(algId, encodedPoint);
        X509EncodedKeySpec keySpec;
        try {
            keySpec = new X509EncodedKeySpec(spki.getEncoded());
        } catch (IOException ex) {
            throw new InvalidKeySpecException(ex.getMessage(), ex);
        }

        KeyFactory kf;
        try {
            kf = KeyFactory.getInstance("EC", "BC");
        } catch (NoSuchAlgorithmException | NoSuchProviderException ex) {
            throw new InvalidKeySpecException(ex.getMessage(), ex);
        }
        return (ECPublicKey) kf.generatePublic(keySpec);
    }
    //</editor-fold>
}