/*
 * 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 org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.KeyManagementException;
import java.security.KeyPair;
import java.security.KeyStore.PasswordProtection;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.SignatureException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Calendar;
import java.util.Date;
import java.util.GregorianCalendar;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

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.asymmetric.service.AsymmetricService;
import com.scytl.cryptolib.certificates.bean.RootCertificateData;
import com.scytl.cryptolib.certificates.bean.ValidityDates;
import com.scytl.cryptolib.certificates.bean.X509DistinguishedName;
import com.scytl.cryptolib.certificates.service.CertificatesService;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.cryptoapi.EncryptionParameters;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.service.PrimitivesService;
import com.scytl.cryptolib.stores.service.StoresService;

/**
 * Tests of {@link GeneratorImpl}.
 */
public class GeneratorImplTest {
    private static final String ENCRYPTION_PARAMETERS_JSON =
        "{\"encryptionParams\":{\"p\":\"AMoJqsCm6UG4B+FGO1Vh6qJcCJcMAVyZFIbZnTIZyD3XAQSTLdxVZxrXm5ZNQMajKAmjUfb3/BV84OgDZaNMXQziV2Uw6hQ70jUcswZ6ZfHAQBy+kwSozDM1LA6ndqp0JIGmPD7LqFIhv9ly7427ksgX1LjNR1tUInVaCS26dfP5wkZyCxwa9/3TapFpFugek+WSJqozaDSd89WAuqXrbBu0FMJWvBVvJapsjfJLEOXTHzdeWef9X5M4FY4Uk+/0WaDVZXy5JhcdoxuNU6Oy0/9buDrPfzlg5a26gpV5ei7XHfykYxdBIPUQoAmraSyUp/byhSDcQu6GJxkfknfAORU=\",\"q\":\"AJT/piLUjl6zN9D/Pc4GYsAG0KAS57H7xJ/mqb+ycbIl\",\"g\":\"ALNkOUZfTDbTMsCYUAKBAAzPdl7422rYGhca6GUAwEPQMlexuIZfIoft+ww54NDqGrUb4TFyw8Yy+y+9Ryb7AXNockro4XYfUjN6eiwOrPxfWG3r9BxqReWZRp74uQ2Ec1toVBrpHaXQNETOJpPBThMKKqrgIuPeAdMRM9U9viku8X0IDmHcqoNaI2o2fvFLFaxlzt5h4kzeJ0A5Twn8wwp5Uln0KcTBJf3H904Icq1u9WTvCTa1jQkoSzkNDQ1VABNFWBa/hW/RGAG7YqTurYHQB+9sMRpOiN10X+6ihdF5VgGpCgXfquKCSqa7I44vq2Vr9+Bm34Gjy+uDOJWrilg=\"}}";

    private static final String ELECTORAL_AUTHORITY_ID = "electoralAuthorityId";

    private static final String ELECTION_EVENT_ID = "electionEventId";

    private static final String VERIFICATION_CARD_SET_ID = "verificationCardSetId";

    private static final String NODE_ID = "nodeId";

    private static final Date VALID_FROM;
    static {
        GregorianCalendar calendar = new GregorianCalendar();
        calendar.set(Calendar.MILLISECOND, 0);
        VALID_FROM = calendar.getTime();
    }

    private static final Date VALID_TO = new Date(VALID_FROM.getTime() + 1000);

    private static ElGamalServiceAPI elGamalService;

    private static EncryptionParameters parameters;

    private static AsymmetricServiceAPI asymmetricService;

    private static CertificatesServiceAPI certificatesService;

    private static PrivateKey nodeCAPrivateKey;

    private static X509Certificate[] nodeCACertificateChain;

    private PrimitivesServiceAPI primitivesService;

    private Codec codec;

    private GeneratorImpl generator;

    @BeforeClass
    public static void beforeClass() throws GeneralCryptoLibException {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }

        elGamalService = new ElGamalService();

        parameters = ElGamalEncryptionParameters.fromJson(ENCRYPTION_PARAMETERS_JSON);

        asymmetricService = new AsymmetricService();
        certificatesService = new CertificatesService();

        KeyPair pair = asymmetricService.getKeyPairForSigning();
        nodeCAPrivateKey = pair.getPrivate();

        RootCertificateData data = new RootCertificateData();
        data.setSubjectDn(new X509DistinguishedName.Builder("CA", "ES").addLocality("Barcelona")
            .addOrganization("Scytl").addOrganizationalUnit("Product Department").build());
        data.setSubjectPublicKey(pair.getPublic());
        data.setValidityDates(new ValidityDates(VALID_FROM, VALID_TO));
        X509Certificate nodeCACertificate =
            certificatesService.createRootAuthorityX509Certificate(data, nodeCAPrivateKey).getCertificate();
        nodeCACertificateChain = new X509Certificate[] {nodeCACertificate };

    }

    @Before
    public void setUp() throws Exception {
        primitivesService = new PrimitivesService();
        codec = new CodecImpl(new StoresService(), asymmetricService);
        generator = new GeneratorImpl(asymmetricService, certificatesService, elGamalService, primitivesService, codec,
            NODE_ID);
    }

    @Test
    public void testGenerateChoiceCodeKeys() throws KeyManagementException, GeneralCryptoLibException {
        NodeKeys nodeKeys = generator.generateNodeKeys(nodeCAPrivateKey, nodeCACertificateChain);
        ElectionSigningKeys electionSigningKeys =
            generator.generateElectionSigningKeys(ELECTION_EVENT_ID, VALID_FROM, VALID_TO, nodeKeys);
        ChoiceCodeKeysSpec spec = new ChoiceCodeKeysSpec.Builder().setElectionEventId(ELECTION_EVENT_ID)
            .setVerificationCardSetId(VERIFICATION_CARD_SET_ID).setParameters(parameters).setGenerationLength(1)
            .setDecryptionLength(2).build();

        ChoiceCodeKeys keys = generator.generateChoiceCodeKeys(spec, electionSigningKeys);

        ZpSubgroup group = keys.generationPrivateKey().getGroup();
        assertEquals(parameters.getGroup(), group);

        assertEquals(1, keys.generationPrivateKey().getKeys().size());
        assertEquals(1, keys.generationPublicKey().getKeys().size());
        assertTrue(asymmetricService.verifySignature(keys.generationPublicKeySignature(),
            electionSigningKeys.publicKey(), codec.encodeElGamalPublicKey(keys.generationPublicKey()),
            ELECTION_EVENT_ID.getBytes(StandardCharsets.UTF_8),
            VERIFICATION_CARD_SET_ID.getBytes(StandardCharsets.UTF_8)));

        assertEquals(2, keys.decryptionPrivateKey().getKeys().size());
        assertEquals(2, keys.decryptionPublicKey().getKeys().size());
        assertTrue(asymmetricService.verifySignature(keys.decryptionPublicKeySignature(),
            electionSigningKeys.publicKey(), codec.encodeElGamalPublicKey(keys.decryptionPublicKey()),
            ELECTION_EVENT_ID.getBytes(StandardCharsets.UTF_8),
            VERIFICATION_CARD_SET_ID.getBytes(StandardCharsets.UTF_8)));
    }

    @Test
    public void testGenerateElectionSigningKeys()
            throws KeyManagementException, GeneralCryptoLibException, InvalidKeyException, CertificateException,
            NoSuchAlgorithmException, NoSuchProviderException, SignatureException {
        NodeKeys nodeKeys = generator.generateNodeKeys(nodeCAPrivateKey, nodeCACertificateChain);

        ElectionSigningKeys keys =
            generator.generateElectionSigningKeys(ELECTION_EVENT_ID, VALID_FROM, VALID_TO, nodeKeys);

        PrivateKey privateKey = keys.privateKey();
        PublicKey publicKey = keys.publicKey();
        byte[] bytes = publicKey.getEncoded();
        byte[] signature = asymmetricService.sign(privateKey, bytes);
        assertTrue(asymmetricService.verifySignature(signature, publicKey, bytes));

        X509Certificate[] certificateChain = keys.certificateChain();
        assertEquals(2, certificateChain.length);
        assertEquals(nodeKeys.caCertificate(), certificateChain[1]);
        X509Certificate certificate = certificateChain[0];
        assertEquals(VALID_FROM, certificate.getNotBefore());
        assertEquals(VALID_TO, certificate.getNotAfter());
        certificate.verify(nodeKeys.caPublicKey());
    }

    @Test
    public void testGenerateMixingKeys()
            throws KeyManagementException, GeneralCryptoLibException, FingerprintGeneratorException {
        NodeKeys nodeKeys = generator.generateNodeKeys(nodeCAPrivateKey, nodeCACertificateChain);
        ElectionSigningKeys electionSigningKeys =
            generator.generateElectionSigningKeys(ELECTION_EVENT_ID, VALID_FROM, VALID_TO, nodeKeys);
        MixingKeysSpec spec = new MixingKeysSpec.Builder().setElectionEventId(ELECTION_EVENT_ID)
            .setElectoralAuthorityId(ELECTORAL_AUTHORITY_ID).setParameters(parameters).setLength(1).build();

        MixingKeys keys = generator.generateMixingKeys(spec, electionSigningKeys);

        ZpSubgroup group = keys.privateKey().getGroup();
        assertEquals(parameters.getGroup(), group);

        assertEquals(1, keys.privateKey().getKeys().size());
        assertEquals(1, keys.publicKey().getKeys().size());
        assertTrue(asymmetricService.verifySignature(keys.publicKeySignature(), electionSigningKeys.publicKey(),
            codec.encodeElGamalPublicKey(keys.publicKey()), ELECTION_EVENT_ID.getBytes(StandardCharsets.UTF_8),
            ELECTORAL_AUTHORITY_ID.getBytes(StandardCharsets.UTF_8)));
    }

    @Test
    public void testGenerateNodeKeys()
            throws KeyManagementException, GeneralCryptoLibException, InvalidKeyException, CertificateException,
            NoSuchAlgorithmException, NoSuchProviderException, SignatureException {
        NodeKeys keys = generator.generateNodeKeys(nodeCAPrivateKey, nodeCACertificateChain);

        assertEquals(nodeCAPrivateKey, keys.caPrivateKey());
        assertArrayEquals(nodeCACertificateChain, keys.caCertificateChain());

        X509Certificate caCertificate = keys.caCertificate();

        PrivateKey encryptionPrivateKey = keys.encryptionPrivateKey();
        PublicKey encryptionPublicKey = keys.encryptionPublicKey();
        byte[] bytes = encryptionPublicKey.getEncoded();
        byte[] encrypted = asymmetricService.encrypt(encryptionPublicKey, bytes);
        assertArrayEquals(bytes, asymmetricService.decrypt(encryptionPrivateKey, encrypted));

        X509Certificate[] encryptionCertificateChain = keys.encryptionCertificateChain();
        assertEquals(2, encryptionCertificateChain.length);
        assertEquals(caCertificate, encryptionCertificateChain[1]);
        X509Certificate encryptionCertificate = encryptionCertificateChain[0];
        assertEquals(VALID_FROM, caCertificate.getNotBefore());
        assertEquals(VALID_TO, caCertificate.getNotAfter());
        encryptionCertificate.verify(caCertificate.getPublicKey());

        PrivateKey logSigningPrivateKey = keys.logSigningPrivateKey();
        PublicKey logSigningPublicKey = keys.logSigningPublicKey();
        bytes = logSigningPublicKey.getEncoded();
        encrypted = asymmetricService.encrypt(logSigningPublicKey, bytes);
        assertArrayEquals(bytes, asymmetricService.decrypt(logSigningPrivateKey, encrypted));

        X509Certificate[] logSigningCertificateChain1 = keys.logSigningCertificateChain();
        assertEquals(2, logSigningCertificateChain1.length);
        assertEquals(caCertificate, logSigningCertificateChain1[1]);
        X509Certificate logSigningCertificate = logSigningCertificateChain1[0];
        assertEquals(VALID_FROM, caCertificate.getNotBefore());
        assertEquals(VALID_TO, caCertificate.getNotAfter());
        logSigningCertificate.verify(caCertificate.getPublicKey());

        PrivateKey logEncryptionPrivateKey = keys.logEncryptionPrivateKey();
        PublicKey logEncryptionPublicKey = keys.logEncryptionPublicKey();
        bytes = logEncryptionPublicKey.getEncoded();
        encrypted = asymmetricService.encrypt(logEncryptionPublicKey, bytes);
        assertArrayEquals(bytes, asymmetricService.decrypt(logEncryptionPrivateKey, encrypted));

        X509Certificate[] logEncryptionCertificateChain = keys.logEncryptionCertificateChain();
        assertEquals(2, logEncryptionCertificateChain.length);
        assertEquals(caCertificate, logEncryptionCertificateChain[1]);
        X509Certificate logEncryptionCertificate = logEncryptionCertificateChain[0];
        assertEquals(VALID_FROM, caCertificate.getNotBefore());
        assertEquals(VALID_TO, caCertificate.getNotAfter());
        logEncryptionCertificate.verify(caCertificate.getPublicKey());
    }

    @Test
    public void testGeneratePassword() throws KeyManagementException {
        PasswordProtection protection = generator.generatePassword();
        assertTrue(protection.getPassword().length >= 16);
    }
}
