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

import static org.junit.Assert.assertNotNull;
import static org.junit.Assume.assumeTrue;

import java.security.Security;
import java.util.Arrays;

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

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.commons.configuration.Provider;
import com.scytl.cryptolib.commons.utils.system.OperatingSystem;
import com.scytl.cryptolib.derivation.cryptoapi.CryptoAPIDerivedKey;
import com.scytl.cryptolib.derivation.cryptoapi.CryptoAPIPBKDFDeriver;
import com.scytl.cryptolib.primitives.derivation.configuration.ConfigKDFDerivationParameters;
import com.scytl.cryptolib.primitives.derivation.configuration.ConfigPBKDFDerivationParameters;
import com.scytl.cryptolib.primitives.derivation.configuration.DerivationPolicy;
import com.scytl.cryptolib.primitives.securerandom.configuration.ConfigSecureRandomAlgorithmAndProvider;

public class CryptoPBKDFDeriverTest {

    public static final double SUN_JCE_VERSION = 1.8d;

    private final CryptoKeyDeriverFactory _cryptoKeyDeriverFactory =
        new CryptoKeyDeriverFactory(getKeyDerivationPolicy());

    @BeforeClass
    public static void setUp() {
        Security.addProvider(new BouncyCastleProvider());
    }

    @AfterClass
    public static void removeBouncyCastleProvider() {
        Security.removeProvider(BouncyCastleProvider.PROVIDER_NAME);
    }

    @Test
    public void testThatDerivesKeyForPBKDF2SHA256BCProviderKL128()
            throws GeneralCryptoLibException {

        char[] password = "passwordTESTpassword".toCharArray();

        CryptoPBKDFDeriver cryptoPBKDFDeriver =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        CryptoAPIDerivedKey cryptoPBKDFDerivedKey =
            cryptoPBKDFDeriver.deriveKey(password);

        assertNotNull(cryptoPBKDFDerivedKey);
    }

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

        char[] password = "passwordTESTpassword".toCharArray();
        byte[] invalidSalt = "0".getBytes();

        CryptoPBKDFDeriver cryptoPBKDFDeriver =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        CryptoAPIDerivedKey cryptoPBKDFDerivedKey =
            cryptoPBKDFDeriver.deriveKey(password, invalidSalt);

        assertNotNull(cryptoPBKDFDerivedKey);
    }

    @Test
    public void testThatDerivesPasswordPBKDF2SHA256BCProviderKL128()
            throws GeneralCryptoLibException {

        char[] password = "passwordTESTpassword".toCharArray();

        CryptoAPIPBKDFDeriver deriver =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        CryptoAPIDerivedKey deriveKey = deriver.deriveKey(password);

        assertNotNull(deriveKey);
    }

    @Test
    public void testThatDerivesPasswordPBKDF2SHA256SunJCEProviderKL128()
            throws GeneralCryptoLibException {

        double version =
            Security.getProvider(Provider.SUN_JCE.getProviderName())
                .getVersion();
        assumeTrue(version >= SUN_JCE_VERSION);

        char[] password = "passwordTESTpassword".toCharArray();

        CryptoKeyDeriverFactory cryptoKeyDeriverFactory =
            new CryptoKeyDeriverFactory(getKeyDerivationSHA256Policy());

        CryptoAPIPBKDFDeriver deriver =
            cryptoKeyDeriverFactory.createPBKDFDeriver();

        CryptoAPIDerivedKey deriveKey = deriver.deriveKey(password);

        assertNotNull(deriveKey);
    }

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

        char[] password = "short".toCharArray();

        CryptoPBKDFDeriver cryptoPBKDFDeriver =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        cryptoPBKDFDeriver.deriveKey(password);
    }

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

        char[] password = "short".toCharArray();

        CryptoPBKDFDeriver cryptoDerivedFromPassword =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        cryptoDerivedFromPassword.deriveKey(password);
    }

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

        char[] password = new char[1024];
        Arrays.fill(password, 'o');

        CryptoPBKDFDeriver cryptoPBKDFDeriver =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        cryptoPBKDFDeriver.deriveKey(password);
    }

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

        char[] password = new char[1024];
        Arrays.fill(password, 'o');

        CryptoPBKDFDeriver cryptoDerivedFromPassword =
            _cryptoKeyDeriverFactory.createPBKDFDeriver();

        cryptoDerivedFromPassword.deriveKey(password);
    }

    private DerivationPolicy getKeyDerivationPolicyForLinux() {
        return new DerivationPolicy() {

            @Override
            public ConfigKDFDerivationParameters getKDFDerivationParameters() {
                return ConfigKDFDerivationParameters.MGF1_SHA256_BC;
            }

            @Override
            public ConfigPBKDFDerivationParameters getPBKDFDerivationParameters() {
                return ConfigPBKDFDerivationParameters.PBKDF2_1_SHA256_256_BC_KL128;
            }

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {
                return ConfigSecureRandomAlgorithmAndProvider.NATIVE_PRNG_SUN;
            }
        };
    }

    private DerivationPolicy getKeyDerivationPolicyForWindows() {
        return new DerivationPolicy() {

            @Override
            public ConfigKDFDerivationParameters getKDFDerivationParameters() {
                return ConfigKDFDerivationParameters.MGF1_SHA256_BC;
            }

            @Override
            public ConfigPBKDFDerivationParameters getPBKDFDerivationParameters() {
                return ConfigPBKDFDerivationParameters.PBKDF2_1_SHA256_256_BC_KL128;
            }

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {
                return ConfigSecureRandomAlgorithmAndProvider.PRNG_SUN_MSCAPI;
            }
        };
    }

    private DerivationPolicy getKeyDerivationSHA256Policy() {
        return new DerivationPolicy() {

            @Override
            public ConfigKDFDerivationParameters getKDFDerivationParameters() {
                return ConfigKDFDerivationParameters.MGF1_SHA256_BC;
            }

            @Override
            public ConfigPBKDFDerivationParameters getPBKDFDerivationParameters() {
                return ConfigPBKDFDerivationParameters.PBKDF2_32000_SHA256_256_BC_KL128;
            }

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {
                if (OperatingSystem.WINDOWS.isCurrent()) {
                    return ConfigSecureRandomAlgorithmAndProvider.PRNG_SUN_MSCAPI;
                }
                return ConfigSecureRandomAlgorithmAndProvider.NATIVE_PRNG_SUN;
            }
        };
    }

    private DerivationPolicy getKeyDerivationPolicy() {
        if (OperatingSystem.WINDOWS.isCurrent()) {
            return getKeyDerivationPolicyForWindows();
        }
        return getKeyDerivationPolicyForLinux();
    }
}
