/*
 * 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.factory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.security.KeyStore.PasswordProtection;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Arrays;

import org.junit.Before;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.derivation.cryptoapi.CryptoAPIDerivedKey;
import com.scytl.cryptolib.primitives.derivation.factory.CryptoPBKDFDeriver;
import com.scytl.cryptolib.scytl.keystore.configuration.ConfigScytlKeyStoreTypeAndProvider;
import com.scytl.cryptolib.scytl.keystore.tests.BaseScytlKeyStoreTests;
import com.scytl.cryptolib.symmetric.cipher.factory.CryptoSymmetricCipher;
import com.scytl.cryptolib.symmetric.cipher.factory.SymmetricCipherFactory;
import com.scytl.cryptolib.symmetric.key.configuration.ConfigSecretKeyAlgorithmAndSpec;

import mockit.Deencapsulation;
import mockit.Expectations;
import mockit.Injectable;

public class CryptoScytlKeyStoreWithPBKDFTest
        extends BaseScytlKeyStoreTests {

    private CryptoScytlKeyStoreWithPBKDF _target;

    @Injectable
    private KeyStore _keystore;

    @Injectable
    private SymmetricCipherFactory _symmetricCipherFactory;

    @Injectable
    private CryptoPBKDFDeriver _cryptoPBKDFDeriver;

    @Injectable
    private CryptoAPIDerivedKey _okPbeKey;

    private PrivateKey _okPrivateKey;

    private byte[] _okDerived;

    private final byte[] _okSalt = new byte[] {1 };

    @Injectable
    private ConfigSecretKeyAlgorithmAndSpec configSecretKeyAlgorithmAndSpec;

    @Injectable
    private CryptoSymmetricCipher _cryptoSymmetricCipher;

    @Before
    public void setup() throws GeneralSecurityException, IOException {

        _okPrivateKey = loadPrivateKey("privateKey.pem");
        _okCert = loadX509Certificate("cert.pem");
        _okChain = new Certificate[1];
        _okChain[0] = _okCert;

        _okDerived = new byte[] {1, 2, 3, 4, 5 };

        new Expectations() {
            {
                _cryptoPBKDFDeriver.generateRandomSalt();
                returns(_okSalt);
                _symmetricCipherFactory.create();
                returns(_cryptoSymmetricCipher);
            }
        };
        _target = Deencapsulation.newInstance(
            CryptoScytlKeyStoreWithPBKDF.class,
            ConfigScytlKeyStoreTypeAndProvider.PKCS12_SUN_JSSE.getType(),
            ConfigScytlKeyStoreTypeAndProvider.PKCS12_SUN_JSSE
                .getProvider(),
            _cryptoPBKDFDeriver, _symmetricCipherFactory,
            configSecretKeyAlgorithmAndSpec);
        Deencapsulation.setField(_target, _keystore);
    }

    private final PasswordProtection _okPassword =
        new PasswordProtection("01234567890abcdefghijk".toCharArray());

    private Certificate[] _okChain;

    private final String _okAlias = "myalias";

    private final String _okAliasWithUnderscore = "m_y_a_l_i_a_s-";

    private X509Certificate _okCert;

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

        new Expectations() {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    _okSalt);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.setKeyEntry(_okAlias, _okPrivateKey,
                    (char[]) any, _okChain);
            }
        };

        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setNullPrivateKeyTest()
            throws KeyStoreException, GeneralCryptoLibException,
            NoSuchAlgorithmException, CertificateException, IOException {
        _target.setSecretKeyEntry(_okAlias, null,
            _okPassword.getPassword());
    }

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

        new Expectations() {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    _okSalt);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.setKeyEntry(_okAliasWithUnderscore,
                    _okPrivateKey, (char[]) any, _okChain);
            }
        };

        _target.setPrivateKeyEntry(_okAliasWithUnderscore, _okPrivateKey,
            _okPassword.getPassword(), _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryShortAliasTest()
            throws GeneralCryptoLibException {

        final String alias = "";
        _target.setPrivateKeyEntry(alias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);
    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryNullAliasTest()
            throws GeneralCryptoLibException {

        final String alias = null;
        _target.setPrivateKeyEntry(alias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);
    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryLongAliasTest()
            throws GeneralCryptoLibException {

        final String alias =
            "sfkjasdhfkjashfjasfhksjfhsjkafhkawhfkweufhskufhkuefhksufhukefhksauhfkashefukhaskefhaskuefhsk";
        _target.setPrivateKeyEntry(alias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);
    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryInvalidAlphabetAliasTest()
            throws GeneralCryptoLibException {

        final String alias = "sfkjasdhfkj!ahfkashefukhaskefhaskuefhsk";
        _target.setPrivateKeyEntry(alias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);
    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryShortPasswordTest()
            throws GeneralSecurityException, IOException,
            GeneralCryptoLibException {

        new Expectations() {
            {

            }
        };
        final char[] password = "a".toCharArray();

        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey, password,
            _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryLongPasswordTest()
            throws GeneralCryptoLibException {
        char[] password = new char[1024];
        Arrays.fill(password, 'a');
        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey, password,
            _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryEmptyPasswordTest()
            throws GeneralCryptoLibException {
        final char[] password = "".toCharArray();

        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey, password,
            _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryNullPasswordTest()
            throws GeneralCryptoLibException {
        final char[] password = null;

        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey, password,
            _okChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryEmptyChainTest()
            throws GeneralCryptoLibException {

        Certificate[] emptyChain = new Certificate[0];
        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey,
            _okPassword.getPassword(), emptyChain);

    }

    @Test(expected = GeneralCryptoLibException.class)
    public void setKeyEntryNullChainTest()
            throws GeneralCryptoLibException {

        Certificate[] nullChain = null;
        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey,
            _okPassword.getPassword(), nullChain);

    }

    @Test(expected = CryptoLibException.class)
    public void getChainWithKeyStoreExceptionTest()
            throws KeyStoreException, GeneralCryptoLibException {

        new Expectations() {
            {
                _keystore.getCertificateChain(_okAlias);
                result = new KeyStoreException("test");
            }
        };
        _target.getCertificateChain(_okAlias);

    }

    @Test(expected = CryptoLibException.class)
    public void getKeyWithKeyStoreExceptionTest()
            throws GeneralSecurityException, GeneralCryptoLibException {

        new Expectations() {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    (byte[]) any);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.getKey(_okAlias, (char[]) any);

                result = new GeneralSecurityException("test");
            }
        };

        _target.getPrivateKeyEntry(_okAlias, _okPassword.getPassword());
    }

    @Test(expected = CryptoLibException.class)
    public void setKeyEntryWithExceptionTest()
            throws KeyStoreException, GeneralCryptoLibException {

        new Expectations() {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    (byte[]) any);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.setKeyEntry(_okAlias, _okPrivateKey,
                    (char[]) any, _okChain);
                result = new KeyStoreException("test");

            }
        };

        _target.setPrivateKeyEntry(_okAlias, _okPrivateKey,
            _okPassword.getPassword(), _okChain);
    }

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

        OutputStream mockedOut = new ByteArrayOutputStream();

        new Expectations(mockedOut) {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    _okSalt);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.store((OutputStream) any, (char[]) any);
            }
        };

        _target.store(mockedOut, _okPassword.getPassword());
    }

    @Test(expected = CryptoLibException.class)
    public void storeWithExceptionTest()
            throws IOException, CertificateException,
            GeneralSecurityException, GeneralCryptoLibException {

        OutputStream mockedOut = new ByteArrayOutputStream();

        new Expectations(mockedOut) {
            {
                _cryptoPBKDFDeriver.deriveKey(_okPassword.getPassword(),
                    _okSalt);
                returns(_okPbeKey);
                _okPbeKey.getEncoded();
                returns(_okDerived);
                _keystore.store((OutputStream) any, (char[]) any);
                result = new GeneralSecurityException("test");
            }
        };

        _target.store(mockedOut, _okPassword.getPassword());
    }
}
