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

/**
 * @date 19/12/2013 18:55:41
 *
 * Copyright (C) 2013 Scytl Secure Electronic Voting SA
 *
 * All rights reserved.
 */

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

import javax.naming.ConfigurationException;

import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.commons.utils.system.OperatingSystem;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncrypterValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalKeyPair;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.configuration.ConfigGroupType;
import com.scytl.cryptolib.elgamal.configuration.ElGamalPolicy;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.mathematical.groups.utils.GroupUtils;
import com.scytl.cryptolib.primitives.securerandom.configuration.ConfigSecureRandomAlgorithmAndProvider;
import com.scytl.cryptolib.test.tools.configuration.QrSubgroupLoader;
import com.scytl.cryptolib.test.tools.configuration.ZpSubgroupLoader;

import mockit.Mock;
import mockit.MockUp;

public class CryptoElGamalEncrypterTest {

    private static CryptoElGamalEncrypter _encrypterForSmallZpGroup;

    private static ZpSubgroup _small_group;

    private static ElGamalPublicKey _publicKey;

    private static int _numKeys;

    private static final int SHORT_EXPONENT_LEN = 256;

    @Rule
    public ExpectedException expectedEx = ExpectedException.none();

    @BeforeClass
    public static void setUp() throws GeneralCryptoLibException {

        BigInteger p = new BigInteger("23");
        BigInteger q = new BigInteger("11");
        BigInteger g = new BigInteger("2");

        _small_group = new ZpSubgroup(g, p, q);

        _numKeys = 4;

        ElGamalFactory elGamalFactoryForZpGroup = new ElGamalFactory(getCryptoElGamalZpPolicy());

        ElGamalEncryptionParameters elGamalEncryptionParameters = new ElGamalEncryptionParameters(p, q, g);
        ElGamalKeyPair cryptoElGamalKeyPair = elGamalFactoryForZpGroup.createCryptoElGamalKeyPairGenerator()
            .generateKeys(elGamalEncryptionParameters, _numKeys);

        _publicKey = cryptoElGamalKeyPair.getPublicKeys();

        _encrypterForSmallZpGroup = elGamalFactoryForZpGroup.createEncrypter(_publicKey);
    }

    @Test
    public void given4Keys6CiphertextsWhenEncryptWith4PrecomputedValuesThenException()
            throws GeneralCryptoLibException {

        expectedEx.expect(GeneralCryptoLibException.class);
        expectedEx.expectMessage(
            "Number of mathematical group elements to ElGamal encrypt must be less than or equal to encrypter public key length: 4; Found 6");

        ZpGroupElement element = new ZpGroupElement(new BigInteger("9"), _small_group);

        int messageSize = 6;
        List<ZpGroupElement> ciphertext = new ArrayList<>(messageSize);
        for (int i = 0; i < messageSize; i++) {
            ciphertext.add(element);
        }

        List<ZpGroupElement> prephis = new ArrayList<>(_numKeys);
        for (int i = 0; i < _numKeys; i++) {
            prephis.add(element);
        }

        ElGamalComputationsValues preComputedValues = new ElGamalComputationsValues(element, prephis);

        Exponent randomExponent = new Exponent(_small_group.getQ(), new BigInteger("2"));

        ElGamalEncrypterValues values = new ElGamalEncrypterValues(randomExponent, preComputedValues);

        _encrypterForSmallZpGroup.encryptGroupElements(ciphertext, values);
    }

    @Test
    public void given4Keys4CiphertextsWhenEncryptWith6PrecomputerValuesThenException()
            throws GeneralCryptoLibException {

        expectedEx.expect(GeneralCryptoLibException.class);
        expectedEx.expectMessage(
            "Number of ElGamal pre-computed phi elements  must be less than or equal to encrypter public key length: 4; Found 6");

        ZpGroupElement element = new ZpGroupElement(new BigInteger("9"), _small_group);

        List<ZpGroupElement> ciphertext = new ArrayList<>(_numKeys);
        for (int i = 0; i < _numKeys; i++) {
            ciphertext.add(element);
        }

        int phiSize = 6;
        List<ZpGroupElement> prephis = new ArrayList<>(phiSize);
        for (int i = 0; i < 6; i++) {
            prephis.add(element);
        }

        ElGamalComputationsValues preComputedValues = new ElGamalComputationsValues(element, prephis);

        Exponent randomExponent = new Exponent(_small_group.getQ(), new BigInteger("2"));

        ElGamalEncrypterValues values = new ElGamalEncrypterValues(randomExponent, preComputedValues);

        _encrypterForSmallZpGroup.encryptGroupElements(ciphertext, values);
    }

    @Test
    public void whenEncryptThenGammaAndPhisAreGroupMembers() throws GeneralCryptoLibException {

        ZpGroupElement element = new ZpGroupElement(BigInteger.ONE, _small_group);

        List<ZpGroupElement> ciphertext = new ArrayList<>();
        for (int i = 0; i < _numKeys; i++) {
            ciphertext.add(element);
        }

        ElGamalComputationsValues encryptionValues =
            _encrypterForSmallZpGroup.encryptGroupElements(ciphertext).getComputationValues();

        // Assert that has gamma and phis
        assertHasGammaAndPhis(encryptionValues);

        // Check that the number of generated phis is the same as the number of
        // elements to encrypt
        Assert.assertEquals("The number of generated phis is the same that the number of elements to encrypt",
            ciphertext.size(), encryptionValues.getPhis().size());

        // Assert that encryption values are group members
        assertEncryptionValuesAreGroupMembers(encryptionValues);
    }

    @Test
    public void whenEncryptNullStringListThenException() throws GeneralCryptoLibException {

        expectedEx.expect(GeneralCryptoLibException.class);
        expectedEx.expectMessage("List of stringified mathematical group elements to ElGamal encrypt is null.");

        _encrypterForSmallZpGroup.encryptStrings(null).getComputationValues();
    }

    @Test
    public void whenEncryptEmptyStringListThenException() throws GeneralCryptoLibException {

        expectedEx.expect(GeneralCryptoLibException.class);
        expectedEx.expectMessage("List of stringified mathematical group elements to ElGamal encrypt is empty.");

        List<String> ciphertext = new ArrayList<>();

        _encrypterForSmallZpGroup.encryptStrings(ciphertext).getComputationValues();
    }

    @Test
    public void whenEncryptNonNumericStringListThenException() throws GeneralCryptoLibException {

        expectedEx.expect(CryptoLibException.class);
        expectedEx.expectMessage("I am not a number 'is not a valid representation of a BigInteger.");

        List<String> ciphertext = new ArrayList<>();
        ciphertext.add("I am not a number");

        _encrypterForSmallZpGroup.encryptStrings(ciphertext).getComputationValues();
    }

    @Test
    public void whenEncryptGroupElementListThenOk() throws GeneralCryptoLibException {

        List<String> ciphertext = new ArrayList<>();
        ciphertext.add("4");

        _encrypterForSmallZpGroup.encryptStrings(ciphertext).getComputationValues();
    }

    @Test
    public void whenEncryptListMultipleGroupElementListThenOk() throws GeneralCryptoLibException {

        List<String> ciphertext = new ArrayList<>();
        ciphertext.add("4");
        ciphertext.add("4");
        ciphertext.add("4");
        ciphertext.add("4");

        _encrypterForSmallZpGroup.encryptStrings(ciphertext).getComputationValues();
    }

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

        List<String> ciphertext = new ArrayList<>();
        ciphertext.add("4");
        ciphertext.add("4");
        ciphertext.add("4");
        ciphertext.add("4");
        ciphertext.add("4");

        _encrypterForSmallZpGroup.encryptStrings(ciphertext).getComputationValues();
    }

    @Test
    public void whenPreComputeThenGammaAndPrePhisAreGroupMembers() throws GeneralCryptoLibException {

        ElGamalComputationsValues preComputedValues = _encrypterForSmallZpGroup.preCompute().getComputationValues();

        // Assert that has gamma and prephis
        assertHasGammaAndPrePhis(preComputedValues);

        // Assert that precomputation values are group members
        assertPreComputeValuesAreGroupMembers(preComputedValues);
    }

    @Test
    public void testWhenEncryptExponentZeroThenPhisEqualsMessages() throws GeneralCryptoLibException {

        // result identity (in this case 1) to be returned when
        // exponentiate to exponent=0
        final ZpGroupElement identityElement = _small_group.getIdentity();

        new MockUp<ZpGroupElement>() {
            @Mock
            public ZpGroupElement exponentiate(final Exponent exponent) {
                return identityElement;
            }
        };

        // create ciphertext
        List<ZpGroupElement> ciphertext = createCiphertext(_small_group);

        ElGamalComputationsValues encryptionResult =
            _encrypterForSmallZpGroup.encryptGroupElements(ciphertext).getComputationValues();

        Assert.assertEquals("Gamma is not 1", encryptionResult.getGamma().getValue(), BigInteger.ONE);
        Assert.assertEquals("Phis are not equal with messages", encryptionResult.getPhis(), ciphertext);
    }

    @Test
    public void whenEncryptQrElementWithNotShortExponentThenExponentSizeCorrect()
            throws GeneralCryptoLibException {

        ZpSubgroup group = obtainQrGroupFromFile();
        boolean useShortExponent = false;

        Exponent exponent = encryptAndObtainRandomExponent(useShortExponent, group, getCryptoElGamalQrPolicy());

        Assert.assertTrue("Exponent was larger than expected",
            exponent.getValue().bitLength() <= group.getQ().bitLength());
    }

    @Test
    public void whenEncryptQrElementWithShortExponentThenExponentSizeCorrect()
            throws GeneralCryptoLibException {

        ZpSubgroup group = obtainQrGroupFromFile();
        boolean useShortExponent = true;

        Exponent exponent = encryptAndObtainRandomExponent(useShortExponent, group, getCryptoElGamalQrPolicy());

        Assert.assertTrue("Short Exponent was not used!", exponent.getValue().bitLength() <= SHORT_EXPONENT_LEN);
    }

    @Test
    public void whenEncryptZpElementWithShortExponentThenExceptionForNonQrGroup()
            throws GeneralCryptoLibException, ConfigurationException {

        expectedEx.expect(CryptoLibException.class);
        expectedEx.expectMessage(
            "Short Exponents can only be used with Quadratic Residue Groups. Please check policy settings.");

        ZpSubgroup group = obtainZpGroupFromFile();
        boolean useShortExponent = true;

        encryptAndObtainRandomExponent(useShortExponent, group, getCryptoElGamalZpPolicy());
    }

    @Test
    public void whenEncryptZpElementWithNotShortExponentThenExponentSizeCorrect() throws GeneralCryptoLibException, ConfigurationException {

        ZpSubgroup group = obtainZpGroupFromFile();
        boolean useShortExponent = false;

        Exponent exponent = encryptAndObtainRandomExponent(useShortExponent, group, getCryptoElGamalQrPolicy());

        Assert.assertTrue("Exponent was larger than expected",
            exponent.getValue().bitLength() <= group.getQ().bitLength());
    }

    private Exponent encryptAndObtainRandomExponent(final boolean isShort, ZpSubgroup group, final ElGamalPolicy policy)
            throws GeneralCryptoLibException {

        List<ZpGroupElement> ciphertext = new ArrayList<>();

        // for test purpose the q value is used
        BigInteger randomGroupElementValue = group.getQ();
        ciphertext.add(new ZpGroupElement(randomGroupElementValue, group));

        List<ZpGroupElement> pubKeys = new ArrayList<>();
        pubKeys.add(new ZpGroupElement(randomGroupElementValue, group));

        CryptoElGamalEncrypter encrypter =
            new ElGamalFactory(policy).createEncrypter(new ElGamalPublicKey(pubKeys, group));

        if (isShort) {
            return encrypter.encryptGroupElementsWithShortExponent(ciphertext).getExponent();
        } else {
            return encrypter.encryptGroupElements(ciphertext).getExponent();
        }
    }

    private ZpSubgroup obtainQrGroupFromFile() throws GeneralCryptoLibException {

        QrSubgroupLoader qrSubgroupLoader = new QrSubgroupLoader();
        BigInteger p = qrSubgroupLoader.getP();
        BigInteger g = qrSubgroupLoader.getG();
        return new GroupUtils().buildQuadraticResidueGroupFromPAndG(p, g);
    }

    private ZpSubgroup obtainZpGroupFromFile() throws GeneralCryptoLibException {

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

    private List<ZpGroupElement> createCiphertext(final ZpSubgroup group) throws GeneralCryptoLibException {

        ZpGroupElement ciphertextElement1 = new ZpGroupElement(new BigInteger("2"), group);
        ZpGroupElement ciphertextElement2 = new ZpGroupElement(new BigInteger("3"), group);
        ZpGroupElement message3 = new ZpGroupElement(new BigInteger("5"), group);

        List<ZpGroupElement> ciphertext = new ArrayList<>();
        ciphertext.add(ciphertextElement1);
        ciphertext.add(ciphertextElement2);
        ciphertext.add(message3);

        return ciphertext;
    }

    private void assertPreComputeValuesAreGroupMembers(final ElGamalComputationsValues preComputationValues) {

        String errorMessage = "All precomputed elements should be group members";

        Assert.assertEquals(errorMessage, true, _small_group.isGroupMember(preComputationValues.getGamma()));

        for (ZpGroupElement element : preComputationValues.getPhis()) {
            Assert.assertEquals(errorMessage, true, _small_group.isGroupMember(element));
        }
    }

    private void assertHasGammaAndPrePhis(final ElGamalComputationsValues preComputationValues) {

        Assert.assertNotNull("The gamma should not be null", preComputationValues.getGamma());
        Assert.assertNotNull("The phis set should not be null", preComputationValues.getPhis());
    }

    private void assertEncryptionValuesAreGroupMembers(final ElGamalComputationsValues encryptionValues) {

        String errorMessage = "All encrypted elements should be group members";

        Assert.assertEquals(errorMessage, true, _small_group.isGroupMember(encryptionValues.getGamma()));

        for (ZpGroupElement element : encryptionValues.getPhis()) {
            Assert.assertEquals(errorMessage, true, _small_group.isGroupMember(element));
        }
    }

    private void assertHasGammaAndPhis(final ElGamalComputationsValues encryptionValues) {

        Assert.assertNotNull("The gamma should not be null", encryptionValues.getGamma());
        Assert.assertNotNull("The phis set should not be null", encryptionValues.getPhis());
    }

    private static ElGamalPolicy getCryptoElGamalZpPolicy() {

        return new ElGamalPolicy() {

            @Override
            public ConfigGroupType getGroupType() {
                return ConfigGroupType.ZP_2048_256;
            }

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {

                switch (OperatingSystem.current()) {

                case WINDOWS:
                    return ConfigSecureRandomAlgorithmAndProvider.PRNG_SUN_MSCAPI;
                case LINUX:
                    return ConfigSecureRandomAlgorithmAndProvider.NATIVE_PRNG_SUN;
                default:
                    throw new CryptoLibException("OS not supported");
                }
            }
        };
    }

    private static ElGamalPolicy getCryptoElGamalQrPolicy() {

        return new ElGamalPolicy() {

            @Override
            public ConfigGroupType getGroupType() {
                return ConfigGroupType.QR_2048;
            }

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {

                switch (OperatingSystem.current()) {

                case WINDOWS:
                    return ConfigSecureRandomAlgorithmAndProvider.PRNG_SUN_MSCAPI;
                case LINUX:
                    return ConfigSecureRandomAlgorithmAndProvider.NATIVE_PRNG_SUN;
                default:
                    throw new CryptoLibException("OS not supported");
                }
            }
        };
    }
}
