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

import static org.junit.Assert.assertEquals;

import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Security;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;

import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.digests.SHA256Digest;
import org.bouncycastle.crypto.generators.KDF2BytesGenerator;
import org.bouncycastle.crypto.kems.RSAKeyEncapsulation;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.RSAKeyParameters;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.BeforeClass;
import org.junit.Test;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.asymmetric.keypair.configuration.ConfigEncryptionKeyPairAlgorithmAndSpec;
import com.scytl.cryptolib.asymmetric.keypair.configuration.ConfigSigningKeyPairAlgorithmAndSpec;
import com.scytl.cryptolib.asymmetric.keypair.configuration.KeyPairPolicy;
import com.scytl.cryptolib.commons.utils.system.OperatingSystem;
import com.scytl.cryptolib.primitives.securerandom.configuration.ConfigSecureRandomAlgorithmAndProvider;

public class KeyPairFactoryTest {

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

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

    @Test
    public void givenRsaCryptoPolicySunRsaSignWhenGenerateSigningKeyPairThenExpectedFieldsFromLinux() {
        Assume.assumeTrue(OperatingSystem.LINUX.isCurrent());
        createKeyPairThenAssertValues(getKeyPairPolicySunRSASignProviderLinuxSecureRandomAlg());
    }

    @Test
    public void givenRsaCryptoPolicyBCWhenGenerateSigningKeyPairThenExpectedFieldsFromLinux() {
        Assume.assumeTrue(OperatingSystem.LINUX.isCurrent());
        createKeyPairThenAssertValues(getKeyPairPolicyBCProviderLinuxSecureRandomAlg());
    }

    @Test
    public void givenRsaCryptoPolicySunRsaSignWhenGenerateSigningKeyPairThenExpectedFieldsFromWindows() {
        Assume.assumeTrue(OperatingSystem.WINDOWS.isCurrent());
        createKeyPairThenAssertValues(getKeyPairPolicySunRSASignProviderWindowsSecureRandomAlg());
    }

    @Test
    public void givenRsaCryptoPolicyBCWhenGenerateSigningKeyPairThenExpectedFieldsFromWindows() {
        Assume.assumeTrue(OperatingSystem.WINDOWS.isCurrent());
        createKeyPairThenAssertValues(getKeyPairPolicyBCProviderWindowsSecureRandomAlg());
    }

    @Test(expected = CryptoLibException.class)
    public void whenCreateKeyPairFactoryUsingSecureRandomAlgFromWrongOperatingSystem() {

        Assume.assumeTrue(OperatingSystem.LINUX.isCurrent());

        createKeyPairThenAssertValues(getKeyPairPolicyBCProviderWindowsSecureRandomAlg());

    }

    private void createKeyPairThenAssertValues(
            final KeyPairPolicy keyPairPolicy) {

        KeyPairGeneratorFactory keyPairFactory =
            new KeyPairGeneratorFactory(keyPairPolicy);

        CryptoKeyPairGenerator keyPairGenerator =
            keyPairFactory.createSigning();

        KeyPair keyPair = keyPairGenerator.genKeyPair();
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAKeyParameters bcPublicKey =
            new RSAKeyParameters(false, publicKey.getModulus(),
                publicKey.getPublicExponent());
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        RSAKeyParameters bcPrivateKey =
            new RSAKeyParameters(true, privateKey.getModulus(),
                privateKey.getPrivateExponent());

        assertPublicKeyHasExpectedAlgorithmAndFormat(keyPairPolicy
            .getSigningKeyPairAlgorithmAndSpec().getAlgorithm(),
            keyPair.getPublic());

        assertPrivateKeyHasExpectedAlgorithmAndFormat(keyPairPolicy
            .getSigningKeyPairAlgorithmAndSpec().getAlgorithm(),
            keyPair.getPrivate());

        // Generate RSA key pair.
        // RSAKeyPairGenerator rsaGen = new RSAKeyPairGenerator();
        // rsaGen.init(new RSAKeyGenerationParameters(BigInteger
        // .valueOf(65537), new SecureRandom(), 1024, 5));
        // AsymmetricCipherKeyPair keys = rsaGen.generateKeyPair();

        AsymmetricCipherKeyPair keys =
            new AsymmetricCipherKeyPair(bcPublicKey, bcPrivateKey);
        // Set RSA-KEM parameters
        RSAKeyEncapsulation kem;
        KDF2BytesGenerator kdf =
            new KDF2BytesGenerator(new SHA256Digest());
        SecureRandom rnd = new SecureRandom();
        int keyByteSize = publicKey.getModulus().bitLength() / Byte.SIZE;
        byte[] out = new byte[keyByteSize];
        KeyParameter key1, key2;
        // Test RSA-KEM
        kem = new RSAKeyEncapsulation(kdf, rnd);
        kem.init(keys.getPublic());
        key1 = (KeyParameter) kem.encrypt(out, 16);

        SecretKey secretKey1 = new SecretKeySpec(key1.getKey(), "AES");

        KDF2BytesGenerator kdf2 =
            new KDF2BytesGenerator(new SHA256Digest());
        SecureRandom rnd2 = new SecureRandom();
        RSAKeyEncapsulation kem2 =
            kem = new RSAKeyEncapsulation(kdf2, rnd2);
        kem2.init(keys.getPrivate());
        key2 = (KeyParameter) kem2.decrypt(out, 16);

        SecretKey secretKey2 = new SecretKeySpec(key2.getKey(), "AES");

        Assert.assertArrayEquals("ARRAYS NOT EQUAL", key1.getKey(),
            key2.getKey());
        Assert.assertArrayEquals("ARRAYS NOT EQUAL",
            secretKey1.getEncoded(), secretKey2.getEncoded());
    }

    private void assertPublicKeyHasExpectedAlgorithmAndFormat(
            final String expectedAlgorithm, final PublicKey publicKey) {

        assertEquals(expectedAlgorithm, publicKey.getAlgorithm());
        assertEquals("X.509", publicKey.getFormat());
    }

    private void assertPrivateKeyHasExpectedAlgorithmAndFormat(
            final String expectedAlgorithm, final PrivateKey privateKey) {

        assertEquals(expectedAlgorithm, privateKey.getAlgorithm());
        assertEquals("PKCS#8", privateKey.getFormat());
    }

    private KeyPairPolicy getKeyPairPolicyBCProviderLinuxSecureRandomAlg() {
        return new KeyPairPolicy() {

            @Override
            public ConfigSigningKeyPairAlgorithmAndSpec getSigningKeyPairAlgorithmAndSpec() {
                return ConfigSigningKeyPairAlgorithmAndSpec.RSA_2048_F4_BC;
            }

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

            @Override
            public ConfigEncryptionKeyPairAlgorithmAndSpec getEncryptingKeyPairAlgorithmAndSpec() {
                return null;
            }

        };
    }

    private KeyPairPolicy getKeyPairPolicySunRSASignProviderLinuxSecureRandomAlg() {
        return new KeyPairPolicy() {

            @Override
            public ConfigSigningKeyPairAlgorithmAndSpec getSigningKeyPairAlgorithmAndSpec() {
                return ConfigSigningKeyPairAlgorithmAndSpec.RSA_2048_F4_SUN_RSA_SIGN;
            }

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

            @Override
            public ConfigEncryptionKeyPairAlgorithmAndSpec getEncryptingKeyPairAlgorithmAndSpec() {
                return null;
            }

        };
    }

    private KeyPairPolicy getKeyPairPolicySunRSASignProviderWindowsSecureRandomAlg() {
        return new KeyPairPolicy() {

            @Override
            public ConfigSigningKeyPairAlgorithmAndSpec getSigningKeyPairAlgorithmAndSpec() {
                return ConfigSigningKeyPairAlgorithmAndSpec.RSA_2048_F4_SUN_RSA_SIGN;
            }

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

            @Override
            public ConfigEncryptionKeyPairAlgorithmAndSpec getEncryptingKeyPairAlgorithmAndSpec() {
                return null;
            }

        };
    }

    private KeyPairPolicy getKeyPairPolicyBCProviderWindowsSecureRandomAlg() {
        return new KeyPairPolicy() {

            @Override
            public ConfigSigningKeyPairAlgorithmAndSpec getSigningKeyPairAlgorithmAndSpec() {
                return ConfigSigningKeyPairAlgorithmAndSpec.RSA_2048_F4_BC;
            }

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

            @Override
            public ConfigEncryptionKeyPairAlgorithmAndSpec getEncryptingKeyPairAlgorithmAndSpec() {
                return null;
            }

        };
    }
}
