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

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.Reader;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyStore.PasswordProtection;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.List;

import javax.crypto.SecretKey;

import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.service.PrimitivesService;
import com.scytl.cryptolib.scytl.keystore.cryptoapi.CryptoAPIScytlKeyStore;
import com.scytl.cryptolib.scytl.keystore.tests.BaseScytlKeyStoreTests;
import com.scytl.cryptolib.test.tools.configuration.ZpSubgroupLoader;

import mockit.Tested;

public class ScytlKeyStoreServiceIT extends BaseScytlKeyStoreTests {

    private static final String ALIAS_1 = "alias1";

    private static final String ALIAS_2 = "alias2";

    @Tested
    private ScytlKeyStoreService _keyStoreService;

    @Tested
    ElGamalService _elGamalService;

    @Tested
    PrimitivesService _primitivesService;

    private static ElGamalPrivateKey _elGamalPrivateKey1;

    private static ElGamalPrivateKey _elGamalPrivateKey2;

    @BeforeClass
    public static void init() throws GeneralCryptoLibException {

        ZpSubgroupLoader zpSubgroupLoader = new ZpSubgroupLoader();
        BigInteger p = zpSubgroupLoader.getP();
        BigInteger q = zpSubgroupLoader.getQ();
        BigInteger g = zpSubgroupLoader.getG();
        ZpSubgroup zpSubgroup = new ZpSubgroup(g, p, q);

        _elGamalPrivateKey1 = getElGamalPrivateKeyEntry(zpSubgroup,
            Arrays.asList(new Exponent(q, g),
                new Exponent(q, g.add(BigInteger.TEN))));

        _elGamalPrivateKey2 = getElGamalPrivateKeyEntry(zpSubgroup,
            Arrays.asList(new Exponent(q, g.subtract(BigInteger.ONE)),
                new Exponent(q, g.subtract(BigInteger.TEN))));

    }

    @Test
    public void createP12WithPBKDF()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        X509Certificate cert = loadX509Certificate("cert.pem");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");

        Certificate[] chain = new Certificate[1];
        chain[0] = cert;

        final String alias = "myalias_-";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setPrivateKeyEntry(alias, privateKey,
            password.getPassword(), chain);

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

        assertTrue(privateKeyAliases.contains(alias));

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

        assertTrue(secretKeyAliases.isEmpty());

        try (FileOutputStream outKeyStore =
            new FileOutputStream("target/keystore.sks")) {
            cryptoKeyStore.store(outKeyStore, password.getPassword());
        }

        Assert.assertNotNull(cryptoKeyStore);
    }

    @Test
    public void openSksWithSoreBiggerThan10k()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        X509Certificate cert = loadX509Certificate("cert.pem");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");

        Certificate[] chain = new Certificate[1];
        chain[0] = cert;

        final String alias = "myalias_";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        int privateKeyCount = 50;
        for (int i = 0; i < privateKeyCount; i++) {
            cryptoKeyStore.setPrivateKeyEntry(alias + i, privateKey,
                password.getPassword(), chain);
        }

        OutputStream outputStream =
            new FileOutputStream("target/bigSks.sks");
        cryptoKeyStore.store(outputStream, password.getPassword());

        InputStream inputStream = new FileInputStream("target/bigSks.sks");

        CryptoAPIScytlKeyStore cryptoKeyStoreOpened =
            _keyStoreService.loadKeyStore(inputStream, password);

        Assert.assertNotNull(cryptoKeyStoreOpened);

        List<String> aliases = cryptoKeyStoreOpened.getPrivateKeyAliases();

        assertTrue(aliases.size() == privateKeyCount);
    }

    @Test
    public void retrieveKeyWithoutStoreCall()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        X509Certificate cert = loadX509Certificate("cert.pem");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");

        Certificate[] chain = new Certificate[1];
        chain[0] = cert;

        final String alias = "myalias_-";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setPrivateKeyEntry(alias, privateKey,
            password.getPassword(), chain);

        PrivateKey key = cryptoKeyStore.getPrivateKeyEntry(alias,
            password.getPassword());

        try (OutputStream outKeyStore =
            new FileOutputStream("target/keystore.sks")) {
            cryptoKeyStore.store(outKeyStore, password.getPassword());
        }

        Assert.assertNotNull(cryptoKeyStore);

        Assert.assertEquals(privateKey, key);
    }

    @Test
    public void loadP12WithPBKDF()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        X509Certificate cert = loadX509Certificate("cert.pem");
        Certificate[] chain = new Certificate[1];
        chain[0] = cert;

        PrivateKey privateKey = loadPrivateKey("privateKey.pem");

        final String alias = "myalias_-";

        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());
        Key key;
        CryptoAPIScytlKeyStore keystore;
        try (InputStream in =
            getClass().getResourceAsStream("/keystore.sks")) {
            
            keystore = _keyStoreService.loadKeyStore(in, password);

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

            assertTrue(privateKeyAliases.contains(alias));

            key =
                keystore.getPrivateKeyEntry(alias, password.getPassword());
        }

        Assert.assertEquals(privateKey, key);
        Assert.assertArrayEquals(chain,
            keystore.getCertificateChain(alias));
    }

    @Test
    public void secretKeyTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        SecretKey secretKey = loadSecretKey("secretKey.bin");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");
        X509Certificate cert = loadX509Certificate("cert.pem");
        Certificate[] chain = new Certificate[1];
        chain[0] = cert;
        List<String> privateKeyAliases;
        List<String> secretKeyAliases;

        final String alias = "myalias";
        final String symmetricAlias_1 = "myaliassymmetric_";
        final String symmetricAlias_2 = "myaliassymmetric-2_";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setPrivateKeyEntry(alias, privateKey,
            password.getPassword(), chain);

        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_1, secretKey,
            password.getPassword());
        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_2, secretKey,
            password.getPassword());

        privateKeyAliases = cryptoKeyStore.getPrivateKeyAliases();

        assertTrue(privateKeyAliases.contains(alias));
        Assert.assertFalse(privateKeyAliases.contains(symmetricAlias_1));

        secretKeyAliases = cryptoKeyStore.getSecretKeyAliases();

        assertTrue(secretKeyAliases.contains(symmetricAlias_1));
        assertTrue(secretKeyAliases.contains(symmetricAlias_2));
        Assert.assertFalse(secretKeyAliases.contains(alias));

        try (OutputStream outKeyStore =
            new FileOutputStream("target/keystoreSymmetric.sks")) {
            cryptoKeyStore.store(outKeyStore, password.getPassword());
        }

        Assert.assertNotNull(cryptoKeyStore);

        
        try (InputStream in =
                new FileInputStream("target/keystoreSymmetric.sks")) {
            cryptoKeyStore = _keyStoreService.loadKeyStore(in, password);
        }
        
        privateKeyAliases = cryptoKeyStore.getPrivateKeyAliases();
        
        assertTrue(privateKeyAliases.contains(alias));
        Assert
        .assertFalse(privateKeyAliases.contains(symmetricAlias_1));
        
        secretKeyAliases = cryptoKeyStore.getSecretKeyAliases();
        
        assertTrue(secretKeyAliases.contains(symmetricAlias_1));
        assertTrue(secretKeyAliases.contains(symmetricAlias_2));
        Assert.assertFalse(secretKeyAliases.contains(alias));
        
        Key key = cryptoKeyStore.getSecretKeyEntry(symmetricAlias_1,
            password.getPassword());

        Assert.assertEquals(secretKey, key);
    }

    @Test
    public void elGamalPrivateKeyTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        final PasswordProtection storePass =
            new PasswordProtection("01234567890abvgdeyozh".toCharArray());

        final PasswordProtection keyPass =
            new PasswordProtection("abvgdeyozh01234567890".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setElGamalPrivateKeyEntry(ALIAS_1,
            _elGamalPrivateKey1, keyPass.getPassword());

        List<String> aliases =
            cryptoKeyStore.getElGamalPrivateKeyAliases();

        assertTrue(aliases.contains(ALIAS_1));
        assertEquals(1, aliases.size());

        cryptoKeyStore.setElGamalPrivateKeyEntry(ALIAS_2,
            _elGamalPrivateKey2, keyPass.getPassword());

        aliases = cryptoKeyStore.getElGamalPrivateKeyAliases();

        assertTrue(aliases.contains(ALIAS_1));
        assertTrue(aliases.contains(ALIAS_2));
        assertEquals(2, aliases.size());

        
        try (OutputStream outKeyStore =
                new FileOutputStream("target/scytl_keystore.sks");) {
            cryptoKeyStore.store(outKeyStore, storePass.getPassword());
        }

        Assert.assertNotNull(cryptoKeyStore);

        try (InputStream in =
            new FileInputStream("target/scytl_keystore.sks")) {
            cryptoKeyStore = _keyStoreService.loadKeyStore(in, storePass);
        }

        aliases = cryptoKeyStore.getElGamalPrivateKeyAliases();
        
        assertTrue(aliases.contains(ALIAS_1));
        assertTrue(aliases.contains(ALIAS_2));
        assertEquals(2, aliases.size());
        ElGamalPrivateKey key1 = cryptoKeyStore.getElGamalPrivateKeyEntry(ALIAS_1,
            keyPass.getPassword());
        ElGamalPrivateKey key2 = cryptoKeyStore.getElGamalPrivateKeyEntry(ALIAS_2,
            keyPass.getPassword());
        Assert.assertEquals(_elGamalPrivateKey1, key1);
        Assert.assertEquals(_elGamalPrivateKey2, key2);
    }

    @Test
    public void longElGamalPrivateKeyTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        ElGamalPrivateKey longElGamalPrivateKey = ElGamalPrivateKey
            .fromJson(readFile("longElGamalPrivateKey.txt"));

        final PasswordProtection storePass =
            new PasswordProtection("01234567890abvgdeyozh".toCharArray());

        final PasswordProtection keyPass =
            new PasswordProtection("abvgdeyozh01234567890".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setElGamalPrivateKeyEntry(ALIAS_1,
            longElGamalPrivateKey, keyPass.getPassword());

        try (FileOutputStream outKeyStore =
            new FileOutputStream("target/scytl_keystore.sks")) {
            cryptoKeyStore.store(outKeyStore, storePass.getPassword());
        }

        ElGamalPrivateKey key;
        try (FileInputStream in =
            new FileInputStream("target/scytl_keystore.sks")) {
            cryptoKeyStore = _keyStoreService.loadKeyStore(in, storePass);
            key = cryptoKeyStore.getElGamalPrivateKeyEntry(ALIAS_1,
                keyPass.getPassword());
        }

        Assert.assertEquals(longElGamalPrivateKey, key);
    }

    private static String readFile(final String path) throws IOException {
        StringBuilder content = new StringBuilder();
        try (InputStream stream = ScytlKeyStoreServiceIT.class
            .getClassLoader().getResourceAsStream(path);
                Reader reader = new InputStreamReader(stream,
                    StandardCharsets.UTF_8)) {
            int c;
            while ((c = reader.read()) != -1) {
                content.append((char) c);
            }
        }
        return content.toString();
    }

    private static ElGamalPrivateKey getElGamalPrivateKeyEntry(
            final ZpSubgroup zpSubgroup,
            final List<Exponent> keyExponentList)
            throws GeneralCryptoLibException {
        return new ElGamalPrivateKey(keyExponentList, zpSubgroup);
    }

    @Test
    public void jsonTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        SecretKey secretKey = loadSecretKey("secretKey.bin");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");
        X509Certificate cert = loadX509Certificate("cert.pem");
        Certificate[] chain = new Certificate[1];
        chain[0] = cert;
        List<String> privateKeyAliases;

        final String alias = "myalias";
        final String symmetricAlias_1 = "myaliassymmetric_";
        final String symmetricAlias_2 = "myaliassymmetric-2_";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setPrivateKeyEntry(alias, privateKey,
            password.getPassword(), chain);

        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_1, secretKey,
            password.getPassword());
        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_2, secretKey,
            password.getPassword());

        privateKeyAliases = cryptoKeyStore.getPrivateKeyAliases();

        assertTrue(privateKeyAliases.contains(alias));
        Assert.assertFalse(privateKeyAliases.contains(symmetricAlias_1));

        Assert.assertNotNull(cryptoKeyStore);

        checkPrivateKey(cryptoKeyStore.toJSON(password.getPassword()),
            password, alias, privateKey);

        try (FileOutputStream outKeyStore =
            new FileOutputStream("target/keystoreSymmetric.sks")) {
            cryptoKeyStore.store(outKeyStore, password.getPassword());
        }
    }

    @Test
    public void formatKeystoreToJsonTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        SecretKey secretKey = loadSecretKey("secretKey.bin");
        PrivateKey privateKey = loadPrivateKey("privateKey.pem");
        X509Certificate cert = loadX509Certificate("cert.pem");
        Certificate[] chain = new Certificate[1];
        chain[0] = cert;

        final String alias = "myalias";
        final String symmetricAlias_1 = "myaliassymmetric_";
        final String symmetricAlias_2 = "myaliassymmetric-2_";
        final PasswordProtection password =
            new PasswordProtection("01234567890abcdefghijk".toCharArray());

        CryptoAPIScytlKeyStore cryptoKeyStore =
            _keyStoreService.createKeyStore();
        cryptoKeyStore.setPrivateKeyEntry(alias, privateKey,
            password.getPassword(), chain);

        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_1, secretKey,
            password.getPassword());
        cryptoKeyStore.setSecretKeyEntry(symmetricAlias_2, secretKey,
            password.getPassword());
        cryptoKeyStore.setElGamalPrivateKeyEntry(ALIAS_1,
            _elGamalPrivateKey1, password.getPassword());

        try (FileOutputStream outKeyStore =
            new FileOutputStream("target/keystoreSymmetric.sks")) {
            cryptoKeyStore.store(outKeyStore, password.getPassword());
        }
        String formatKeyStoreToJSON;
        try (FileInputStream fileInputStream =
            new FileInputStream("target/keystoreSymmetric.sks")) {
            formatKeyStoreToJSON =
                _keyStoreService.formatKeyStoreToJSON(fileInputStream);
        }
        CryptoAPIScytlKeyStore loadedKeyStore;
        try (InputStream stream = new ByteArrayInputStream(
            formatKeyStoreToJSON.getBytes(StandardCharsets.UTF_8))) {
            loadedKeyStore =
                _keyStoreService.loadKeyStoreFromJSON(stream, password);
        }
        assertEquals(privateKey, loadedKeyStore.getPrivateKeyEntry(alias,
            password.getPassword()));
        assertEquals(secretKey, loadedKeyStore
            .getSecretKeyEntry(symmetricAlias_1, password.getPassword()));
        assertEquals(secretKey, loadedKeyStore
            .getSecretKeyEntry(symmetricAlias_2, password.getPassword()));
        assertEquals(_elGamalPrivateKey1, loadedKeyStore
            .getElGamalPrivateKeyEntry(ALIAS_1, password.getPassword()));
    }

    private void checkPrivateKey(final String json,
            final PasswordProtection password, final String alias,
            final PrivateKey privateKey) throws GeneralCryptoLibException {
        byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
        CryptoAPIScytlKeyStore store;
        try (InputStream stream = new ByteArrayInputStream(bytes)) {
            store =
                _keyStoreService.loadKeyStoreFromJSON(stream, password);
        } catch (IOException e) {
            throw new AssertionError(e);
        }
        Key key = store.getPrivateKeyEntry(alias, password.getPassword());
        Assert.assertEquals(privateKey, key);
    }

    @Test
    public void testElGamalPrivateKeyToAndFromJSON()
            throws GeneralCryptoLibException, IOException {
        CryptoAPIScytlKeyStore keyStore =
            _keyStoreService.createKeyStore();

        final PasswordProtection storePass =
            new PasswordProtection("01234567890abvgdeyozh".toCharArray());

        final PasswordProtection keyPass =
            new PasswordProtection("abvgdeyozh01234567890".toCharArray());

        keyStore.setElGamalPrivateKeyEntry(ALIAS_1, _elGamalPrivateKey1,
            keyPass.getPassword());
        keyStore.setElGamalPrivateKeyEntry(ALIAS_2, _elGamalPrivateKey2,
            keyPass.getPassword());

        ByteArrayInputStream inStream = new ByteArrayInputStream(
            keyStore.toJSON(storePass.getPassword()).getBytes());

        CryptoAPIScytlKeyStore actualKeyStore =
            _keyStoreService.loadKeyStoreFromJSON(inStream, storePass);

        assertEquals(_elGamalPrivateKey1, actualKeyStore
            .getElGamalPrivateKeyEntry(ALIAS_1, keyPass.getPassword()));
        assertEquals(_elGamalPrivateKey2, actualKeyStore
            .getElGamalPrivateKeyEntry(ALIAS_2, keyPass.getPassword()));
    }
}
