/*
 * 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.products.ov.mixnet.commons.homomorphic.impl;

import static org.junit.Assert.assertTrue;

import java.math.BigInteger;
import java.util.Arrays;

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

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
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.products.ov.mixnet.commons.tools.GroupTools;

public class GjosteenElGamalTest {

    private static GjosteenElGamal gjosteenElGamal;

    private static GjosteenElGamal gjosteenElGamalMultiElementKeys;

    private static ZpSubgroup group;

    private static Exponent[] privateKeyExponentArray;

    private static Exponent[] otherPrivateKeyExponentArray;

    private static Exponent[] multiElementPrivateKeyExponentArray;

    @BeforeClass
    public static void setUp() {

        try {

            final BigInteger p = new BigInteger("15294034768093677312256663166625633354362303");
            final BigInteger q = new BigInteger("7647017384046838656128331583312816677181151");

            group = GroupTools.createGroupWithRandomGenerator(p, q);

            // create private key exponent array
            Exponent privateKeyExponent = new Exponent(group.getQ(), new BigInteger("222222222"));
            privateKeyExponentArray = new Exponent[1];
            privateKeyExponentArray[0] = privateKeyExponent;

            gjosteenElGamal = new GjosteenElGamal(group, privateKeyExponentArray);

            // create other private key exponent array
            Exponent privateKeyExponentOther = new Exponent(group.getQ(), new BigInteger("11111111"));
            otherPrivateKeyExponentArray = new Exponent[1];
            otherPrivateKeyExponentArray[0] = privateKeyExponentOther;

            // create multi-element private key exponent array
            multiElementPrivateKeyExponentArray = new Exponent[10];
            multiElementPrivateKeyExponentArray[0] = new Exponent(group.getQ(), new BigInteger("11111111"));
            multiElementPrivateKeyExponentArray[1] = new Exponent(group.getQ(), new BigInteger("22222222"));
            multiElementPrivateKeyExponentArray[2] = new Exponent(group.getQ(), new BigInteger("33333333"));
            multiElementPrivateKeyExponentArray[3] = new Exponent(group.getQ(), new BigInteger("44444444"));
            multiElementPrivateKeyExponentArray[4] = new Exponent(group.getQ(), new BigInteger("55555555"));
            multiElementPrivateKeyExponentArray[5] = new Exponent(group.getQ(), new BigInteger("66666666"));
            multiElementPrivateKeyExponentArray[6] = new Exponent(group.getQ(), new BigInteger("77777777"));
            multiElementPrivateKeyExponentArray[7] = new Exponent(group.getQ(), new BigInteger("88888888"));
            multiElementPrivateKeyExponentArray[8] = new Exponent(group.getQ(), new BigInteger("99999999"));
            multiElementPrivateKeyExponentArray[9] = new Exponent(group.getQ(), new BigInteger("12121212"));

            gjosteenElGamalMultiElementKeys = new GjosteenElGamal(group, multiElementPrivateKeyExponentArray);

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }

    }

    @Test
    public void whenEncryptAndDecryptUsingKeysThatAreSetWithCryptosystemThenOk() {

        int numElementsToEncrypt = 1;

        ZpGroupElement[] plaintextAsGroupElements = createListOfGroupElements(numElementsToEncrypt);

        final GjosteenElGamalPlaintext plaintext = new GjosteenElGamalPlaintext(plaintextAsGroupElements);

        final GjosteenElGamalRandomness randomness = createRandomness();

        final Ciphertext ciphertext = gjosteenElGamal.encrypt(plaintext, randomness);

        final ZpGroupElement[] decryptedPlaintext = gjosteenElGamal
            .decrypt(ciphertext.getElements().toArray(new ZpGroupElement[ciphertext.getElements().size()]));

        final String errorMsg = "Expected that the decrypted data would match the plaintext";
        assertTrue(errorMsg, Arrays.equals(plaintextAsGroupElements, decryptedPlaintext));
    }

    @Test
    public void whenEncryptAndDecryptByPassingCorrectPrivateKeyAsArgumentThenOk() {

        int numElementsToEncrypt = 1;

        ZpGroupElement[] plaintextAsGroupElements = createListOfGroupElements(numElementsToEncrypt);

        final GjosteenElGamalPlaintext plaintext = new GjosteenElGamalPlaintext(plaintextAsGroupElements);

        final GjosteenElGamalRandomness randomness = createRandomness();

        final Ciphertext ciphertext = gjosteenElGamal.encrypt(plaintext, randomness);

        final ZpGroupElement[] decryptedPlaintext = gjosteenElGamal.decrypt(
            ciphertext.getElements().toArray(new ZpGroupElement[ciphertext.getElements().size()]),
            privateKeyExponentArray);

        final String errorMsg = "Expected that the decrypted data would match the plaintext";
        assertTrue(errorMsg, Arrays.equals(plaintextAsGroupElements, decryptedPlaintext));
    }

    @Test
    public void whenEncryptAndDecryptByPassingDifferntPrivateKeyAsArgumentThenDifferentPlaintext() {

        int numElementsToEncrypt = 1;

        ZpGroupElement[] plaintextAsGroupElements = createListOfGroupElements(numElementsToEncrypt);

        final GjosteenElGamalPlaintext plaintext = new GjosteenElGamalPlaintext(plaintextAsGroupElements);

        final GjosteenElGamalRandomness randomness = createRandomness();

        final Ciphertext ciphertext = gjosteenElGamal.encrypt(plaintext, randomness);

        final ZpGroupElement[] decryptedPlaintext = gjosteenElGamal.decrypt(
            ciphertext.getElements().toArray(new ZpGroupElement[ciphertext.getElements().size()]),
            otherPrivateKeyExponentArray);

        final String errorMsg = "Expected that the decrypted data would match the plaintext";
        assertTrue(errorMsg, !Arrays.equals(plaintextAsGroupElements, decryptedPlaintext));
    }

    @Test
    public void whenEncryptAndDecryptMultiElementPlaintextThenOk() {

        int numElementsToEncrypt = 10;

        ZpGroupElement[] plaintextAsGroupElements = createListOfGroupElements(numElementsToEncrypt);

        final GjosteenElGamalPlaintext plaintext = new GjosteenElGamalPlaintext(plaintextAsGroupElements);

        final GjosteenElGamalRandomness randomness = createRandomness();

        final Ciphertext ciphertext = gjosteenElGamalMultiElementKeys.encrypt(plaintext, randomness);

        final ZpGroupElement[] decryptedPlaintext = gjosteenElGamalMultiElementKeys
            .decrypt(ciphertext.getElements().toArray(new ZpGroupElement[ciphertext.getElements().size()]));

        final String errorMsg = "Expected that the decrypted data would match the plaintext";
        assertTrue(errorMsg, Arrays.equals(plaintextAsGroupElements, decryptedPlaintext));
    }

    private GjosteenElGamalRandomness createRandomness() {

        try {
            final Exponent exponentUsedInEncryption = new Exponent(group.getQ(), BigInteger.TEN);
            return new GjosteenElGamalRandomness(exponentUsedInEncryption);
        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }

    private ZpGroupElement[] createListOfGroupElements(final int numElements) {

        try {

            ZpGroupElement[] plaintextElements = new ZpGroupElement[numElements];
            for (int i = 0; i < numElements; i++) {

                int exponentUsedToGeneratePlaintextAsInt = i + 5;
                final Exponent exponentUsedToGeneratePlaintext =
                    new Exponent(group.getQ(), BigInteger.valueOf(exponentUsedToGeneratePlaintextAsInt));
                ZpGroupElement plaintextElement = group.getGenerator().exponentiate(exponentUsedToGeneratePlaintext);
                plaintextElements[i] = plaintextElement;
            }

            return plaintextElements;

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }
}
