/*
 * 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.cryptolib.scytl.keystore.service;

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.KeyStore.PasswordProtection;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.util.List;

import javax.crypto.SecretKey;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.asymmetric.utils.AsymmetricTestDataGenerator;
import com.scytl.cryptolib.certificates.utils.X509CertificateTestDataGenerator;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.utils.ElGamalTestDataGenerator;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.mathematical.groups.utils.MathematicalTestDataGenerator;
import com.scytl.cryptolib.scytl.keystore.cryptoapi.CryptoAPIScytlKeyStore;
import com.scytl.cryptolib.scytl.keystore.utils.SksTestDataGenerator;
import com.scytl.cryptolib.symmetric.utils.SymmetricTestDataGenerator;
import com.scytl.cryptolib.test.tools.utils.CommonTestDataGenerator;

/**
 * Tests of the Scytl key store service API.
 */
public class ScytlKeyStoreServiceTest {

    private static final String SCYTL_KEY_STORE_PATH =
        "target/keystore.sks";

    private static final int NUM_ELGAMAL_KEY_ELEMENTS = 5;

    private static ScytlKeyStoreService _scytlKeyStoreService;

    private static CryptoAPIScytlKeyStore _scytlKeyStore;

    private static String _rootPrivateKeyAlias;

    private static String _leafPrivateKeyAlias;

    private static String _secretKey1Alias;

    private static String _secretKey2Alias;

    private static String _elGamalPrivateKey1Alias;

    private static String _elGamalPrivateKey2Alias;

    private static char[] _privateKeyPassword;

    private static char[] _secretKeyPassword;

    private static char[] _elGamalPrivateKeyPassword;

    private static PrivateKey _rootPrivateKey;

    private static Certificate _rootCertificate;

    private static Certificate[] _rootCertificateChain;

    private static PrivateKey _leafPrivateKey;

    private static Certificate _leafCertificate;

    private static Certificate[] _leafCertificateChain;

    private static SecretKey _secretKey1;

    private static SecretKey _secretKey2;

    private static ElGamalPrivateKey _elGamalPrivateKey1;

    private static ElGamalPrivateKey _elGamalPrivateKey2;

    private static char[] _scytlKeyStorePassword;

    private static CryptoAPIScytlKeyStore _loadedScytlKeyStore;

    private static String _scytlKeyStoreAsJson;

    private static String _aliasWithUnderscoreChar;

    private static String _aliasWithHyphenChar;

    @BeforeClass
    public static void setUp()
            throws GeneralCryptoLibException, IOException,
            GeneralSecurityException {

        Security.addProvider(new BouncyCastleProvider());

        _scytlKeyStoreService = new ScytlKeyStoreService();

        _scytlKeyStore = _scytlKeyStoreService.createKeyStore();

        _rootPrivateKeyAlias = SksTestDataGenerator.getAlias();
        do {
            _leafPrivateKeyAlias = SksTestDataGenerator.getAlias();
        } while (_leafPrivateKeyAlias.equals(_rootPrivateKeyAlias));

        _secretKey1Alias = SksTestDataGenerator.getAlias();
        do {
            _secretKey2Alias = SksTestDataGenerator.getAlias();
        } while (_secretKey2Alias.equals(_secretKey1Alias));

        _elGamalPrivateKey1Alias = SksTestDataGenerator.getAlias();
        do {
            _elGamalPrivateKey2Alias = SksTestDataGenerator.getAlias();
        } while (_elGamalPrivateKey2Alias
            .equals(_elGamalPrivateKey1Alias));

        _privateKeyPassword = SksTestDataGenerator.getPassword();
        _secretKeyPassword = SksTestDataGenerator.getPassword();
        _elGamalPrivateKeyPassword = SksTestDataGenerator.getPassword();

        KeyPair rootKeyPair =
            AsymmetricTestDataGenerator.getKeyPairForSigning();
        _rootPrivateKey = rootKeyPair.getPrivate();
        _rootCertificate = X509CertificateTestDataGenerator
            .getRootAuthorityX509Certificate(rootKeyPair).getCertificate();
        _rootCertificateChain = new Certificate[1];
        _rootCertificateChain[0] = _rootCertificate;

        KeyPair leafKeyPair =
            AsymmetricTestDataGenerator.getKeyPairForSigning();
        _leafCertificate = X509CertificateTestDataGenerator
            .getSignX509Certificate(leafKeyPair, rootKeyPair)
            .getCertificate();
        _leafPrivateKey = leafKeyPair.getPrivate();
        _leafCertificateChain = new Certificate[2];
        _leafCertificateChain[0] = _leafCertificate;
        _leafCertificateChain[1] = _rootCertificate;

        _secretKey1 =
            SymmetricTestDataGenerator.getSecretKeyForEncryption();
        _secretKey2 =
            SymmetricTestDataGenerator.getSecretKeyForEncryption();

        ZpSubgroup group = MathematicalTestDataGenerator.getZpSubgroup();
        _elGamalPrivateKey1 = ElGamalTestDataGenerator
            .getKeyPair(group, NUM_ELGAMAL_KEY_ELEMENTS).getPrivateKeys();
        _elGamalPrivateKey2 = ElGamalTestDataGenerator
            .getKeyPair(group, NUM_ELGAMAL_KEY_ELEMENTS).getPrivateKeys();

        _scytlKeyStore.setPrivateKeyEntry(_rootPrivateKeyAlias,
            _rootPrivateKey, _privateKeyPassword, _rootCertificateChain);
        _scytlKeyStore.setPrivateKeyEntry(_leafPrivateKeyAlias,
            _leafPrivateKey, _privateKeyPassword, _leafCertificateChain);

        _scytlKeyStore.setSecretKeyEntry(_secretKey1Alias, _secretKey1,
            _secretKeyPassword);
        _scytlKeyStore.setSecretKeyEntry(_secretKey2Alias, _secretKey2,
            _secretKeyPassword);

        _scytlKeyStore.setElGamalPrivateKeyEntry(_elGamalPrivateKey1Alias,
            _elGamalPrivateKey1, _elGamalPrivateKeyPassword);
        _scytlKeyStore.setElGamalPrivateKeyEntry(_elGamalPrivateKey2Alias,
            _elGamalPrivateKey2, _elGamalPrivateKeyPassword);

        _scytlKeyStorePassword = SksTestDataGenerator.getPassword();

        try (OutputStream outStream =
            new FileOutputStream(SCYTL_KEY_STORE_PATH)) {
            _scytlKeyStore.store(outStream, _scytlKeyStorePassword);
        } catch (IOException e) {
            throw new GeneralCryptoLibException(
                "Could not store Scytl key store to path "
                    + SCYTL_KEY_STORE_PATH,
                e);
        }

        try (InputStream inStream =
            new FileInputStream(SCYTL_KEY_STORE_PATH)) {
            _loadedScytlKeyStore = _scytlKeyStoreService.loadKeyStore(
                inStream, new PasswordProtection(_scytlKeyStorePassword));
        } catch (IOException e) {
            throw new GeneralCryptoLibException(
                "Could not load Scytl key store from path "
                    + SCYTL_KEY_STORE_PATH,
                e);
        }

        _scytlKeyStoreAsJson =
            _scytlKeyStore.toJSON(_scytlKeyStorePassword);

        int index = CommonTestDataGenerator.getInt(0,
            (_leafPrivateKeyAlias.length() - 1));
        StringBuilder builder = new StringBuilder(_leafPrivateKeyAlias);
        builder.setCharAt(index, '_');
        _aliasWithUnderscoreChar = builder.toString();
        builder = new StringBuilder(_leafPrivateKeyAlias);
        builder.setCharAt(index, '-');
        _aliasWithHyphenChar = builder.toString();
    }

    @AfterClass
    public static void tearDown() {

        Security.removeProvider(BouncyCastleProvider.PROVIDER_NAME);
    }

    @Test
    public void whenSetAndGetRootPrivateKeyFromKeyStoreThenOk()
            throws GeneralCryptoLibException {

        PrivateKey privateKey = _scytlKeyStore
            .getPrivateKeyEntry(_rootPrivateKeyAlias, _privateKeyPassword);

        Assert.assertArrayEquals(privateKey.getEncoded(),
            _rootPrivateKey.getEncoded());
    }

    @Test
    public void whenSetAndGetRootCertificateChainFromKeyStoreThenOk()
            throws KeyStoreException, GeneralCryptoLibException,
            NoSuchAlgorithmException, CertificateException, IOException {

        Certificate[] certificateChain =
            _scytlKeyStore.getCertificateChain(_rootPrivateKeyAlias);

        Certificate certificate = certificateChain[0];

        Assert.assertArrayEquals(certificate.getEncoded(),
            _rootCertificate.getEncoded());
    }

    @Test
    public void whenSetAndGetLeafPrivateKeyFromKeyStoreThenOk()
            throws GeneralCryptoLibException {

        PrivateKey privateKey = _scytlKeyStore
            .getPrivateKeyEntry(_leafPrivateKeyAlias, _privateKeyPassword);

        Assert.assertArrayEquals(privateKey.getEncoded(),
            _leafPrivateKey.getEncoded());
    }

    @Test
    public void whenSetAndGetLeafCertificateChainFromKeyStoreThenOk()
            throws KeyStoreException, GeneralCryptoLibException,
            NoSuchAlgorithmException, CertificateException, IOException {

        Certificate[] certificateChain =
            _scytlKeyStore.getCertificateChain(_leafPrivateKeyAlias);

        Assert.assertArrayEquals(certificateChain[0].getEncoded(),
            _leafCertificate.getEncoded());

        Assert.assertArrayEquals(certificateChain[1].getEncoded(),
            _rootCertificate.getEncoded());
    }

    @Test
    public void whenSetAndGetSecretKeysForEncryptionFromKeyStoreThenOk()
            throws GeneralCryptoLibException {

        SecretKey secretKey = _scytlKeyStore
            .getSecretKeyEntry(_secretKey1Alias, _secretKeyPassword);

        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey1.getEncoded());

        secretKey = _scytlKeyStore.getSecretKeyEntry(_secretKey2Alias,
            _secretKeyPassword);

        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey2.getEncoded());
    }

    @Test
    public void whenSetAndGetElGamalPrivateKeysFromKeyStoreThenOk()
            throws GeneralCryptoLibException {

        ElGamalPrivateKey elGamalPrivateKey =
            _scytlKeyStore.getElGamalPrivateKeyEntry(
                _elGamalPrivateKey1Alias, _elGamalPrivateKeyPassword);

        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey1);

        elGamalPrivateKey = _scytlKeyStore.getElGamalPrivateKeyEntry(
            _elGamalPrivateKey2Alias, _elGamalPrivateKeyPassword);

        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey2);
    }

    @Test
    public void whenPrivateKeyGetAliasesFromKeyStoreThenOk() {

        List<String> privateKeyAliases =
            _scytlKeyStore.getPrivateKeyAliases();

        Assert.assertEquals(privateKeyAliases.get(0),
            _rootPrivateKeyAlias);

        Assert.assertEquals(privateKeyAliases.get(1),
            _leafPrivateKeyAlias);
    }

    @Test
    public void whenGetSecretAliasesFromKeyStoreThenOk() {

        List<String> secretKeyAliases =
            _scytlKeyStore.getSecretKeyAliases();

        Assert.assertEquals(secretKeyAliases.get(0), _secretKey1Alias);

        Assert.assertEquals(secretKeyAliases.get(1), _secretKey2Alias);
    }

    @Test
    public void whenGetElGamalPrivateKeyAliasesFromKeyStoreThenOk() {

        List<String> elGamalPrivateKeyAliases =
            _scytlKeyStore.getElGamalPrivateKeyAliases();

        Assert.assertEquals(elGamalPrivateKeyAliases.get(0),
            _elGamalPrivateKey1Alias);

        Assert.assertEquals(elGamalPrivateKeyAliases.get(1),
            _elGamalPrivateKey2Alias);
    }

    @Test
    public void whenSetKeyEntryUsingAliasWithUnderscoreCharThenOk()
            throws GeneralCryptoLibException {

        _scytlKeyStore.setPrivateKeyEntry(_aliasWithUnderscoreChar,
            _rootPrivateKey, _privateKeyPassword, _rootCertificateChain);
        _scytlKeyStore.setPrivateKeyEntry(_aliasWithUnderscoreChar,
            _leafPrivateKey, _privateKeyPassword, _leafCertificateChain);

        _scytlKeyStore.setSecretKeyEntry(_aliasWithUnderscoreChar,
            _secretKey1, _secretKeyPassword);
        _scytlKeyStore.setSecretKeyEntry(_aliasWithUnderscoreChar,
            _secretKey2, _secretKeyPassword);

        _scytlKeyStore.setElGamalPrivateKeyEntry(_aliasWithUnderscoreChar,
            _elGamalPrivateKey1, _elGamalPrivateKeyPassword);
        _scytlKeyStore.setElGamalPrivateKeyEntry(_aliasWithUnderscoreChar,
            _elGamalPrivateKey2, _elGamalPrivateKeyPassword);
    }

    @Test
    public void whenSetKeyEntryUsingAliasWithHyphenCharThenOk()
            throws GeneralCryptoLibException {

        _scytlKeyStore.setPrivateKeyEntry(_aliasWithHyphenChar,
            _leafPrivateKey, _privateKeyPassword, _rootCertificateChain);

        _scytlKeyStore.setSecretKeyEntry(_aliasWithHyphenChar, _secretKey1,
            _secretKeyPassword);

        _scytlKeyStore.setElGamalPrivateKeyEntry(_aliasWithHyphenChar,
            _elGamalPrivateKey1, _elGamalPrivateKeyPassword);
    }

    @Test
    public void whenStoreAndLoadKeyStoreThenOk()
            throws FileNotFoundException, GeneralCryptoLibException {

        PrivateKey privateKey = _loadedScytlKeyStore
            .getPrivateKeyEntry(_rootPrivateKeyAlias, _privateKeyPassword);
        Assert.assertArrayEquals(privateKey.getEncoded(),
            _rootPrivateKey.getEncoded());
        privateKey = _loadedScytlKeyStore
            .getPrivateKeyEntry(_leafPrivateKeyAlias, _privateKeyPassword);
        Assert.assertArrayEquals(privateKey.getEncoded(),
            _leafPrivateKey.getEncoded());

        SecretKey secretKey = _loadedScytlKeyStore
            .getSecretKeyEntry(_secretKey1Alias, _secretKeyPassword);
        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey1.getEncoded());
        secretKey = _loadedScytlKeyStore
            .getSecretKeyEntry(_secretKey2Alias, _secretKeyPassword);
        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey2.getEncoded());

        ElGamalPrivateKey elGamalPrivateKey =
            _loadedScytlKeyStore.getElGamalPrivateKeyEntry(
                _elGamalPrivateKey1Alias, _elGamalPrivateKeyPassword);
        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey1);
        elGamalPrivateKey = _loadedScytlKeyStore.getElGamalPrivateKeyEntry(
            _elGamalPrivateKey2Alias, _elGamalPrivateKeyPassword);
        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey2);
    }

    @Test
    public void whenSerializeAndDeserializeKeyStoreThenOk()
            throws GeneralCryptoLibException {

        InputStream inStream = new ByteArrayInputStream(
            _scytlKeyStoreAsJson.getBytes(StandardCharsets.UTF_8));

        CryptoAPIScytlKeyStore deserializedScytlKeyStore =
            _scytlKeyStoreService.loadKeyStoreFromJSON(inStream,
                new PasswordProtection(_scytlKeyStorePassword));

        PrivateKey privateKey = deserializedScytlKeyStore
            .getPrivateKeyEntry(_rootPrivateKeyAlias, _privateKeyPassword);
        Assert.assertArrayEquals(privateKey.getEncoded(),
            _rootPrivateKey.getEncoded());
        privateKey = deserializedScytlKeyStore
            .getPrivateKeyEntry(_leafPrivateKeyAlias, _privateKeyPassword);
        Assert.assertArrayEquals(privateKey.getEncoded(),
            _leafPrivateKey.getEncoded());

        SecretKey secretKey = deserializedScytlKeyStore
            .getSecretKeyEntry(_secretKey1Alias, _secretKeyPassword);
        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey1.getEncoded());
        secretKey = deserializedScytlKeyStore
            .getSecretKeyEntry(_secretKey2Alias, _secretKeyPassword);
        Assert.assertArrayEquals(secretKey.getEncoded(),
            _secretKey2.getEncoded());

        ElGamalPrivateKey elGamalPrivateKey =
            deserializedScytlKeyStore.getElGamalPrivateKeyEntry(
                _elGamalPrivateKey1Alias, _elGamalPrivateKeyPassword);
        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey1);
        elGamalPrivateKey =
            deserializedScytlKeyStore.getElGamalPrivateKeyEntry(
                _elGamalPrivateKey2Alias, _elGamalPrivateKeyPassword);
        Assert.assertEquals(elGamalPrivateKey, _elGamalPrivateKey2);
    }
}
