package meerkat.crypto.concrete;

import com.google.protobuf.ByteString;
import com.google.protobuf.Message;
import meerkat.crypto.DigitalSignatureGenerator;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Crypto.Signature;
import meerkat.util.Hex;
import org.spongycastle.asn1.ASN1EncodableVector;
import org.spongycastle.asn1.DERSequence;
import org.spongycastle.asn1.x500.X500Name;
import org.spongycastle.asn1.x509.*;
import org.spongycastle.asn1.x509.Extension;
import org.spongycastle.cert.X509v3CertificateBuilder;
import org.spongycastle.cert.jcajce.JcaX509CertificateConverter;
import org.spongycastle.cert.jcajce.JcaX509ExtensionUtils;
import org.spongycastle.jce.provider.BouncyCastleProvider;
import org.spongycastle.operator.ContentSigner;
import org.spongycastle.operator.jcajce.JcaContentSignerBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.*;
import java.security.spec.ECGenParameterSpec;
import java.util.*;


/**
 * Sign and verify digital signatures.
 *
 * This class is not thread-safe (each thread should have its own instance).
 */
public class ECDSASignature implements DigitalSignatureGenerator {

    private static GlobalCryptoSetup globalCryptoSetup = GlobalCryptoSetup.getInstance();

    final Logger logger = LoggerFactory.getLogger(getClass());

    final public static String KEYSTORE_TYPE = "PKCS12";
    final public static String DEFAULT_SIGNATURE_ALGORITHM = "SHA256withECDSA";
    final public static String DEFAULT_KEY_ALGORITHM = "ECDSA";
    final public static String DEFAULT_KEY_CURVE = "secp256k1";


    SHA256Digest certDigest = new SHA256Digest();

    /**
     * Buffer used to hold length in for hash update
     */
    ByteBuffer lenBuf = ByteBuffer.allocate(4);


    Map<ByteString, Certificate> loadedCertificates = new HashMap<ByteString, Certificate>();

    /**
     * Signature currently loaded (will be used in calls to {@link #verify()}).
     */
    ByteString loadedSignature = null;

    ByteString loadedSigningKeyId = null;

    Certificate loadedSigningCert = null;

    /**
     * The actual signing implementation. (used for both signing and verifying)
     */
    java.security.Signature signer;

    /**
     * The currently loaded signing key.
     */
    PrivateKey loadedSigningKey;


    /**
     * Compute a fingerprint of a cert as a SHA256 hash.
     *
     * @param cert
     * @return
     */
    public ByteString computeCertificateFingerprint(Certificate cert) {
        try {
            certDigest.reset();
            byte[] data = cert.getEncoded();
            certDigest.update(data);
            return ByteString.copyFrom(certDigest.digest());
        } catch (CertificateEncodingException e) {
            // Shouldn't happen
            logger.error("Certificate encoding error", e);
            throw new RuntimeException("Certificate encoding error", e);
        }

    }

    public ECDSASignature(java.security.Signature signer) {
        this.signer = signer;
    }

    public ECDSASignature() {
        try {
            this.signer = java.security.Signature.getInstance(DEFAULT_SIGNATURE_ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            // Should never happen
            logger.error("Couldn't find implementation for " + DEFAULT_SIGNATURE_ALGORITHM + " signatures", e);
        }
    }



    public void loadVerificationCertificate(Certificate cert) throws CertificateException {
        // Just checking
        if (!(cert instanceof X509Certificate)) {
            logger.error("Certificate must be in X509 format; got {} instead!", cert.getClass().getCanonicalName());
            return;
        }
        ByteString keyId = computeCertificateFingerprint(cert);
        loadedCertificates.put(keyId, cert);
    }


    @Override
    public void loadVerificationCertificate(Crypto.SignatureVerificationKey cert) throws CertificateException {
        ByteArrayInputStream certStream = new ByteArrayInputStream(cert.getData().toByteArray());
        loadVerificationCertificates(certStream);
    }


    public Collection<? extends Certificate> deserializeCertificates(InputStream certStream)
            throws CertificateException  {
        CertificateFactory certificateFactory = CertificateFactory.getInstance(CERTIFICATE_ENCODING_X509);

        Collection<? extends Certificate> certs = certificateFactory.generateCertificates(certStream);
        return certs;
    }

    public Collection<? extends Certificate> deserializeCertificates(Crypto.SignatureVerificationKey cert)
            throws CertificateException  {
        ByteArrayInputStream certStream = new ByteArrayInputStream(cert.getData().toByteArray());
        return deserializeCertificates(certStream);
    }

    @Override
    public ByteString computeCertificateFingerprint(Crypto.SignatureVerificationKey cert)
            throws CertificateException {
        Collection<? extends Certificate> certs = deserializeCertificates(cert);
        return computeCertificateFingerprint(certs.iterator().next());
    }


    @Override
    public void loadVerificationCertificates(InputStream certStream)
            throws CertificateException {
        Collection<? extends Certificate> certs = deserializeCertificates(certStream);
        for (Certificate cert : certs) {
            loadVerificationCertificate(cert);
        }
    }

    @Override
    public void clearVerificationCertificates() {
        loadedCertificates.clear();
    }


    /**
     * Add the list of messages to the stream that is being verified/signed.
     * Messages are prepended with their length in 32-bit big-endian format.
     *
     * @param msg
     * @throws SignatureException
     */
    @Override
    public void updateContent(Message msg) throws SignatureException {
        assert msg != null;

        lenBuf.clear();
        lenBuf.putInt(msg.getSerializedSize());
        lenBuf.flip();
        signer.update(lenBuf);
        signer.update(msg.toByteString().asReadOnlyByteBuffer());
    }

    @Override
    public void updateContent(byte[] data) throws SignatureException {
        signer.update(data);
    }

    public void updateContent(InputStream in) throws IOException, SignatureException {
        ByteString inStr = ByteString.readFrom(in);
        signer.update(inStr.asReadOnlyByteBuffer());
    }

    @Override
    public Signature sign() throws SignatureException {
        Signature.Builder sig = Signature.newBuilder();
        sig.setType(Crypto.SignatureType.ECDSA);
        sig.setData(ByteString.copyFrom(signer.sign()));
        sig.setSignerId(loadedSigningKeyId);
        return sig.build();
    }

    @Override
    public void initVerify(Signature sig)
            throws CertificateException, InvalidKeyException {
        Certificate cert = loadedCertificates.get(sig.getSignerId());
        if (cert == null) {
            logger.warn("No certificate loaded for ID {}!", sig.getSignerId());
            throw new CertificateException("No certificate loaded for " + sig.getSignerId());
        }
        signer.initVerify(cert.getPublicKey());
        loadedSignature = sig.getData();
        loadedSigningKeyId = null;
    }

    @Override
    public boolean verify() {
        try {
            return signer.verify(loadedSignature.toByteArray());
        } catch (SignatureException e) {
            // Happens only if signature is invalid!
            logger.error("Signature exception", e);
            return false;
        }
    }

    /**
     * Utility method to more easily deal with simple password-protected files.
     *
     * @param password
     * @return
     */
    public CallbackHandler getFixedPasswordHandler(final char[] password) {
        return new CallbackHandler() {
            @Override
            public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
                for (Callback callback : callbacks) {
                    if (callback instanceof PasswordCallback) {
                        PasswordCallback passwordCallback = (PasswordCallback) callback;
                        logger.debug("Requested password ({})", passwordCallback.getPrompt());
                        passwordCallback.setPassword(password);
                    }
                }

            }
        };
    }

    /**
     * Load a keystore from an input stream in PKCS12 format.
     *
     * @param keyStream
     * @param password
     * @return
     * @throws IOException
     * @throws CertificateException
     * @throws KeyStoreException
     * @throws NoSuchAlgorithmException
     */
    @Override
    public KeyStore.Builder getPKCS12KeyStoreBuilder(InputStream keyStream, char[] password)
            throws IOException, CertificateException, KeyStoreException, NoSuchAlgorithmException {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        keyStore.load(keyStream, password);
        return KeyStore.Builder.newInstance(keyStore, new KeyStore.CallbackHandlerProtection(getFixedPasswordHandler(password)));
    }



    public void loadSigningCertificate(Certificate cert, PrivateKey key) throws InvalidKeyException, CertificateException {
        loadedSigningKey = (PrivateKey) key;
        loadedSigningKeyId = computeCertificateFingerprint(cert);
        loadedSigningCert = cert;
        signer.initSign(loadedSigningKey);
        logger.debug("Loaded signing key with ID {}", Hex.encode(loadedSigningKeyId));
    }

    /**
     * For now we only support PKCS12.
     * TODO: Support for PKCS11 as well.
     *
     * @param keyStoreBuilder
     * @throws IOException
     * @throws CertificateException
     * @throws UnrecoverableKeyException
     */
    @Override
    public void loadSigningCertificate(KeyStore.Builder keyStoreBuilder)
            throws IOException, CertificateException, UnrecoverableKeyException {
        try {

            KeyStore keyStore = keyStoreBuilder.getKeyStore();

            // Iterate through all aliases until we find the first privatekey
            Enumeration<String> aliases = keyStore.aliases();
            while (aliases.hasMoreElements()) {
                String alias = aliases.nextElement();
                logger.trace("Testing keystore entry {}", alias);


                try {
                    Certificate cert = keyStore.getCertificate(alias);
                    logger.trace("keystore entry {}, has cert type {}", alias, cert.getClass());

                    Key key;

                    try {
                        key = keyStore.getKey(alias, null);
                    } catch (UnrecoverableKeyException e) {
                        // This might be a keystore that doesn't support callback handlers
                        // (e.g., Java 8 PKCS12)
                        // Manually extract password using callback handler
                        char[] password = null;
                        KeyStore.ProtectionParameter prot = keyStoreBuilder.getProtectionParameter(alias);

                        if (prot instanceof KeyStore.PasswordProtection) {
                            password = ((KeyStore.PasswordProtection) prot).getPassword();
                        } else if (prot instanceof KeyStore.CallbackHandlerProtection) {
                            PasswordCallback callback = new PasswordCallback("Password for " + alias + "?", false);
                            Callback[] callbacks = { callback };
                            try {
                                ((KeyStore.CallbackHandlerProtection) prot).getCallbackHandler().handle(callbacks);
                                password = callback.getPassword();
                            } catch (UnsupportedCallbackException e1) {
                                logger.error("PasswordCallback fallback not supported!", e1);
                                throw new UnrecoverableKeyException("Couldn't use password callback to get key");
                            }
                        } else {
                            logger.error("Unrecognized protection handler for keystore: {}", prot.getClass());
                            throw new UnrecoverableKeyException("Unrecognized protection handler for keystore");
                        }
                        key = keyStore.getKey(alias, password);
                    }
                    logger.trace("keystore entry {}, has key type {}", alias, key.getClass());
                    if (key instanceof PrivateKey) {

                        // Here is where we *actually* load the private key
                        loadSigningCertificate(cert, (PrivateKey) key);
                        return;
                    } else {
                        logger.info("Certificate {} in keystore does not have a private key", cert.toString());
                    }
                } catch(InvalidKeyException e) {
                    logger.info("Read invalid key", e);
                } catch(UnrecoverableEntryException e) {
                    logger.info("Read unrecoverable entry", e);
                }
            }

        } catch (KeyStoreException e) {
            logger.error("Keystore exception", e);
        } catch (NoSuchAlgorithmException e) {
            logger.error("NoSuchAlgorithmException exception", e);
            throw new CertificateException(e);
        }
        logger.error("Didn't find valid private key entry in keystore");
        throw new UnrecoverableKeyException("Didn't find valid private key entry in keystore!");
    }

    @Override
    public ByteString getSignerID() {
        return loadedSigningKeyId;
    }

    @Override
    public Crypto.SignatureVerificationKey getSignerPublicKey() {
        Crypto.SignatureType certType;

        if (loadedSigningCert instanceof X509Certificate) {
            X509Certificate xcert = (X509Certificate) loadedSigningCert;
            String algName = xcert.getSigAlgName();
            if (algName.endsWith("ECDSA")) {
                certType = Crypto.SignatureType.ECDSA;
            } else if (algName.endsWith("DSA")) {
                certType = Crypto.SignatureType.DSA;
            } else if (algName.endsWith("RSA")) {
                certType = Crypto.SignatureType.RSA;
            } else {
                logger.error("Unknown certificate signature algorithm: {}", algName);
                throw new RuntimeException("Unknown certificate signature algorithm: " + algName);
            }
        } else {
            logger.error("We don't know how to handle non-X509 certificates");
            throw new RuntimeException("Unknown Certificate type: " + loadedSigningCert.getClass());
        }
        try {
            return Crypto.SignatureVerificationKey.newBuilder()
                    .setType(certType)
                    .setData(ByteString.copyFrom(loadedSigningCert.getEncoded()))
                    .build();
        } catch (CertificateEncodingException e) {
            logger.error("CertificateEncodingException: ", e);
            throw new RuntimeException(e);
        }
    }

    public void clearSigningKey() {
        try {
            // TODO: Check if this really clears the key from memory
            if (loadedSigningKeyId != null)
                signer.initSign(null);
            loadedSigningKeyId = null;
            loadedSigningKey = null;
            loadedSigningCert = null;
            // Start garbage collection?
        } catch (InvalidKeyException e) {
            // Do nothing
        }
    }


    @Override
    public void generateSigningCertificate(BigInteger serial, Date notBefore, Date notAfter,
                                           String subjectName) {
        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(DEFAULT_KEY_ALGORITHM);
            ECGenParameterSpec kpgparams = new ECGenParameterSpec(DEFAULT_KEY_CURVE);
            keyGen.initialize(kpgparams);
            KeyPair pair = keyGen.generateKeyPair();

            SubjectPublicKeyInfo subPubKeyInfo = SubjectPublicKeyInfo.getInstance(pair.getPublic().getEncoded());

            // build a certificate generator
            X500Name subject = new X500Name("cn=" + subjectName);
            X509v3CertificateBuilder builder = new X509v3CertificateBuilder(subject, serial, notBefore, notAfter,
                    subject, subPubKeyInfo);

            // TODO: Might need to explicitly force SHA256 digest
            JcaX509ExtensionUtils extensionUtils = new JcaX509ExtensionUtils();

            builder.addExtension(Extension.subjectKeyIdentifier, false, extensionUtils.createSubjectKeyIdentifier(pair.getPublic()));
            builder.addExtension(Extension.basicConstraints, true, new BasicConstraints(true));
            KeyUsage usage = new KeyUsage(KeyUsage.keyCertSign | KeyUsage.digitalSignature | KeyUsage.cRLSign);
            builder.addExtension(Extension.keyUsage, false, usage);
            ASN1EncodableVector purposes = new ASN1EncodableVector();
            purposes.add(KeyPurposeId.id_kp_serverAuth);
            purposes.add(KeyPurposeId.id_kp_clientAuth);
            purposes.add(KeyPurposeId.anyExtendedKeyUsage);
            builder.addExtension(Extension.extendedKeyUsage, false, new DERSequence(purposes));

            String signatureAlgorithm = DEFAULT_SIGNATURE_ALGORITHM;

            ContentSigner contentSigner = new JcaContentSignerBuilder(signatureAlgorithm).build(pair.getPrivate());
            X509Certificate cert = new JcaX509CertificateConverter().getCertificate(builder.build(contentSigner));

            loadSigningCertificate(cert, pair.getPrivate());
        } catch (Exception e) {
            // Convert all checked exceptions into runtime exceptions.
            // TODO: Handle exceptions more elegantly.
            throw new RuntimeException(e);
        }
    }

    @Override
    public void storeSignatureKeypair(KeyStore store, String alias, char[] password) throws IOException, KeyStoreException {
        Certificate chain[] = { loadedSigningCert };
        store.setKeyEntry(alias, loadedSigningKey, password, chain);
    }
}
