/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.channel.cc.commons.keymanagement;

import static java.lang.System.arraycopy;
import static java.text.MessageFormat.format;

import java.nio.charset.StandardCharsets;
import java.security.KeyManagementException;
import java.security.KeyPair;
import java.security.KeyStore.PasswordProtection;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.Date;

import com.scytl.cryptolib.api.asymmetric.AsymmetricServiceAPI;
import com.scytl.cryptolib.api.certificates.CertificatesServiceAPI;
import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.primitives.PrimitivesServiceAPI;
import com.scytl.cryptolib.certificates.bean.CertificateData;
import com.scytl.cryptolib.certificates.bean.ValidityDates;
import com.scytl.cryptolib.certificates.bean.X509DistinguishedName;
import com.scytl.cryptolib.certificates.factory.CryptoX509Certificate;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalKeyPairGenerator;
import com.scytl.cryptolib.elgamal.cryptoapi.EncryptionParameters;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.securerandom.cryptoapi.CryptoAPIRandomString;

/**
 * Implementation of {@link Generator}.
 */
class GeneratorImpl implements Generator {
    private static final int PASSWORD_LENGTH = 26;

    private final AsymmetricServiceAPI asymmetricService;

    private final CertificatesServiceAPI certificatesService;

    private final ElGamalServiceAPI elGamalService;

    private final PrimitivesServiceAPI primitivesService;

    private final Codec codec;

    private final String nodeId;

    /**
     * Constructor.
     * 
     * @param asymmetricService
     *            the asymmetric service
     * @param certificatesService
     *            the certificate service
     * @param elGamalService
     *            the ElGamal service
     * @param primitivesService
     *            the primitives service
     * @param codec
     *            the codec
     * @param nodeId
     *            the node identifier.
     */
    public GeneratorImpl(AsymmetricServiceAPI asymmetricService, CertificatesServiceAPI certificatesService,
            ElGamalServiceAPI elGamalService, PrimitivesServiceAPI primitivesService, Codec codec, String nodeId) {
        this.asymmetricService = asymmetricService;
        this.certificatesService = certificatesService;
        this.elGamalService = elGamalService;
        this.primitivesService = primitivesService;
        this.codec = codec;
        this.nodeId = nodeId;
    }

    private static X509DistinguishedName generateSubjectDn(String commonName, X509DistinguishedName issuerDn)
            throws GeneralCryptoLibException {
        return new X509DistinguishedName.Builder(commonName, issuerDn.getCountry()).addLocality(issuerDn.getLocality())
            .addOrganization(issuerDn.getOrganization()).addOrganizationalUnit(issuerDn.getOrganizationalUnit())
            .build();
    }

    private static X509Certificate[] newCertificateChain(X509Certificate certificate, X509Certificate[] issuerChain) {
        X509Certificate[] chain = new X509Certificate[issuerChain.length + 1];
        chain[0] = certificate;
        arraycopy(issuerChain, 0, chain, 1, issuerChain.length);
        return chain;
    }

    @Override
    public ChoiceCodeKeys generateChoiceCodeKeys(ChoiceCodeKeysSpec spec, ElectionSigningKeys electionSigningKeys)
            throws KeyManagementException {
        ChoiceCodeKeys.Builder builder = new ChoiceCodeKeys.Builder();

        ElGamalKeyPair generationPair = generateElGamalKeyPair(spec.parameters(), spec.generationLength());
        byte[] generationPublicKeySignature = signElGamalPublicKey(electionSigningKeys.privateKey(),
            generationPair.getPublicKeys(), spec.electionEventId(), spec.verificationCardSetId());
        builder.setGenerationKeys(generationPair.getPrivateKeys(), generationPair.getPublicKeys(),
            generationPublicKeySignature);

        ElGamalKeyPair decryptionPair = generateElGamalKeyPair(spec.parameters(), spec.decryptionLength());
        byte[] decryptionPublicKeySignature = signElGamalPublicKey(electionSigningKeys.privateKey(),
            decryptionPair.getPublicKeys(), spec.electionEventId(), spec.verificationCardSetId());
        builder.setDecryptionKeys(decryptionPair.getPrivateKeys(), decryptionPair.getPublicKeys(),
            decryptionPublicKeySignature);

        return builder.build();
    }

    @Override
    public ElectionSigningKeys generateElectionSigningKeys(String electionEventId, Date validFrom, Date validTo,
            NodeKeys nodeKeys) throws KeyManagementException {
        KeyPair pair = asymmetricService.getKeyPairForSigning();
        X509Certificate certificate = generateElectionSigningCertificate(electionEventId, validFrom, validTo,
            pair.getPublic(), nodeKeys.caPrivateKey(), nodeKeys.caCertificate());
        X509Certificate[] certificateChain = newCertificateChain(certificate, nodeKeys.caCertificateChain());
        return new ElectionSigningKeys(pair.getPrivate(), certificateChain);
    }

    @Override
    public MixingKeys generateMixingKeys(MixingKeysSpec spec, ElectionSigningKeys electionSigningKeys)
            throws KeyManagementException {
        ElGamalKeyPair pair = generateElGamalKeyPair(spec.parameters(), spec.length());

        byte[] publicKeySignature = signElGamalPublicKey(electionSigningKeys.privateKey(), pair.getPublicKeys(),
            spec.electionEventId(), spec.electoralAuthorityId());

        return new MixingKeys(pair.getPrivateKeys(), pair.getPublicKeys(), publicKeySignature);
    }

    @Override
    public NodeKeys generateNodeKeys(PrivateKey caPrivateKey, X509Certificate[] caCertificateChain)
            throws KeyManagementException {
        NodeKeys.Builder builder = new NodeKeys.Builder();

        builder.setCAKeys(caPrivateKey, caCertificateChain);

        KeyPair encryptionPair = asymmetricService.getKeyPairForEncryption();
        X509Certificate encryptionCertificate =
            generateNodeEncryptionCertificate(encryptionPair.getPublic(), caPrivateKey, caCertificateChain[0]);

        X509Certificate[] encryptionCertificateChain = newCertificateChain(encryptionCertificate, caCertificateChain);
        builder.setEncryptionKeys(encryptionPair.getPrivate(), encryptionCertificateChain);

        KeyPair logSignPair = asymmetricService.getKeyPairForSigning();
        X509Certificate logSignCertificate =
            generateNodeLogSignCertificate(logSignPair.getPublic(), caPrivateKey, caCertificateChain[0]);
        X509Certificate[] logSignCertificateChain = newCertificateChain(logSignCertificate, caCertificateChain);
        builder.setLogSigningKeys(logSignPair.getPrivate(), logSignCertificateChain);

        KeyPair logEncryptionPair = asymmetricService.getKeyPairForEncryption();
        X509Certificate logEncryptionCertificate =
            generateNodeLogEncryptionCertificate(logEncryptionPair.getPublic(), caPrivateKey, caCertificateChain[0]);
        X509Certificate[] logEncryptionCertificateChain =
            newCertificateChain(logEncryptionCertificate, caCertificateChain);
        builder.setLogEncryptionKeys(logEncryptionPair.getPrivate(), logEncryptionCertificateChain);

        return builder.build();
    }

    @Override
    public PasswordProtection generatePassword() throws KeyManagementException {
        CryptoAPIRandomString generator = primitivesService.get32CharAlphabetCryptoRandomString();
        String password;
        try {
            password = generator.nextRandom(PASSWORD_LENGTH);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to generate password.", e);
        }
        return new PasswordProtection(password.toCharArray());
    }

    private X509Certificate generateElectionSigningCertificate(String electionEventId, Date validFrom, Date validTo,
            PublicKey publicKey, PrivateKey caPrivateKey, X509Certificate caCertificate) throws KeyManagementException {
        try {
            CryptoX509Certificate issuerCertificate = new CryptoX509Certificate(caCertificate);
            X509DistinguishedName issuerDn = issuerCertificate.getSubjectDn();

            CertificateData data = new CertificateData();
            X509DistinguishedName subjectDn = generateSubjectDn(electionEventId, issuerDn);
            data.setSubjectDn(subjectDn);
            data.setSubjectPublicKey(publicKey);
            ValidityDates dates = new ValidityDates(validFrom, validTo);
            data.setValidityDates(dates);
            data.setIssuerDn(issuerDn);

            return certificatesService.createSignX509Certificate(data, caPrivateKey).getCertificate();
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to generate election signing certificate.", e);
        }
    }

    private ElGamalKeyPair generateElGamalKeyPair(EncryptionParameters parameters, int length)
            throws KeyManagementException {
        try {
            ElGamalEncryptionParameters elGamalEncryptionParameters;
            if (!(parameters instanceof ElGamalEncryptionParameters)) {
                if (!(parameters.getGroup() instanceof ZpSubgroup)) {
                    throw new KeyManagementException(format("Unsupported encryption parameters ''{0}''", parameters));
                }
                ZpSubgroup group = (ZpSubgroup) parameters.getGroup();
                elGamalEncryptionParameters = new ElGamalEncryptionParameters(group.getP(), group.getQ(), group.getG());
            } else {
                elGamalEncryptionParameters = (ElGamalEncryptionParameters) parameters;
            }
            CryptoAPIElGamalKeyPairGenerator generator = elGamalService.getElGamalKeyPairGenerator();
            return generator.generateKeys(elGamalEncryptionParameters, length);
        } catch (GeneralCryptoLibException e) {
            throw new KeyAlreadyExistsException("Failed to generate ElGamal key pair.", e);
        }
    }

    private X509Certificate generateNodeEncryptionCertificate(PublicKey publicKey, PrivateKey caPrivateKey,
            X509Certificate caCertificate) throws KeyManagementException {
        try {
            CryptoX509Certificate issuerCertificate = new CryptoX509Certificate(caCertificate);
            X509DistinguishedName issuerDn = issuerCertificate.getSubjectDn();

            CertificateData data = new CertificateData();
            String commonName = nodeId + " Encryption";
            data.setSubjectDn(generateSubjectDn(commonName, issuerDn));
            data.setSubjectPublicKey(publicKey);
            data.setValidityDates(new ValidityDates(issuerCertificate.getNotBefore(), issuerCertificate.getNotAfter()));
            data.setIssuerDn(issuerDn);

            return certificatesService.createEncryptionX509Certificate(data, caPrivateKey).getCertificate();
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to generate node encryption certificate.", e);
        }
    }

    private X509Certificate generateNodeLogEncryptionCertificate(PublicKey publicKey, PrivateKey caPrivateKey,
            X509Certificate caCertificate) throws KeyManagementException {
        try {
            CryptoX509Certificate issuerCertificate = new CryptoX509Certificate(caCertificate);
            X509DistinguishedName issuerDn = issuerCertificate.getSubjectDn();

            CertificateData data = new CertificateData();
            String commonName = nodeId + " Log Encryption";
            data.setSubjectDn(generateSubjectDn(commonName, issuerDn));
            data.setSubjectPublicKey(publicKey);
            data.setValidityDates(new ValidityDates(issuerCertificate.getNotBefore(), issuerCertificate.getNotAfter()));
            data.setIssuerDn(issuerDn);

            return certificatesService.createEncryptionX509Certificate(data, caPrivateKey).getCertificate();
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to generate node encryption certificate.", e);
        }
    }

    private X509Certificate generateNodeLogSignCertificate(PublicKey publicKey, PrivateKey caPrivateKey,
            X509Certificate caCertificate) throws KeyManagementException {
        try {
            CryptoX509Certificate issuerCertificate = new CryptoX509Certificate(caCertificate);
            X509DistinguishedName issuerDn = issuerCertificate.getSubjectDn();

            CertificateData data = new CertificateData();
            String commonName = nodeId + " Log Sign";
            data.setSubjectDn(generateSubjectDn(commonName, issuerDn));
            data.setSubjectPublicKey(publicKey);
            data.setValidityDates(new ValidityDates(issuerCertificate.getNotBefore(), issuerCertificate.getNotAfter()));
            data.setIssuerDn(issuerDn);

            return certificatesService.createSignX509Certificate(data, caPrivateKey).getCertificate();
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to generate node encryption certificate.", e);
        }
    }

    private byte[] signElGamalPublicKey(PrivateKey signingKey, ElGamalPublicKey key, String... attributes)
            throws KeyManagementException {
        byte[][] data = new byte[attributes.length + 1][];
        data[0] = codec.encodeElGamalPublicKey(key);
        for (int i = 0; i < attributes.length; i++) {
            data[i + 1] = attributes[i].getBytes(StandardCharsets.UTF_8);
        }
        try {
            return asymmetricService.sign(signingKey, data);
        } catch (GeneralCryptoLibException e) {
            throw new KeyManagementException("Failed to sign ElGamal public key.", e);
        }
    }
}
