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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.KeyException;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

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

import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncryptionParameters;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.service.ElGamalService;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.ov.config.shares.keys.elgamal.ElGamalKeyPairGenerator;
import com.scytl.products.ov.config.shares.keys.elgamal.ElGamalPrivateKeyAdapter;
import com.scytl.products.ov.config.shares.keys.elgamal.ElGamalPublicKeyAdapter;
import com.scytl.shares.ShamirAlgorithm;
import com.scytl.shares.Share;
import com.scytl.shares.ShareException;

public class MultipleSharesContainerTest {

    private static ShamirAlgorithm sa;

    private static List<MultipleSharesContainer> secrets2_shares2_threshold2;

    private static List<MultipleSharesContainer> secrets2_shares3_threshold3;

    @Rule
    public final ExpectedException exception = ExpectedException.none();

    @BeforeClass
    public static void setUp() throws ShareException {

        sa = new ShamirAlgorithm();

        secrets2_shares2_threshold2 = createListOfContainers_secrets2_shares2_threshold2();

        secrets2_shares3_threshold3 = createListOfContainers_secrets2_shares3_threshold3();
    }

    @Test
    public void whenSerializeAndReconstructThenOk() throws ShareException {

        MultipleSharesContainer container1 = secrets2_shares2_threshold2.get(0);
        MultipleSharesContainer container2 = secrets2_shares2_threshold2.get(1);

        byte[] container1Serailized = container1.serialize();
        byte[] container2Serailized = container2.serialize();

        MultipleSharesContainer reconstructed1 = new MultipleSharesContainer(container1Serailized);
        MultipleSharesContainer reconstructed2 = new MultipleSharesContainer(container2Serailized);

        assertTrue(checkIfContainersAreEqual(container1, reconstructed1));
        assertTrue(checkIfContainersAreEqual(container2, reconstructed2));
    }

    @Test
    public void whenCheckIfTwoCompatibleContainersAreCompatibleThenTrue() throws ShareException {

        MultipleSharesContainer container1 = secrets2_shares2_threshold2.get(0);
        MultipleSharesContainer container2 = secrets2_shares2_threshold2.get(1);

        assertTrue(container1.isCompatible(container2));
    }

    @Test
    public void whenCheckIfTwoCompatibleContainersAreCompatibleThenTrue2() throws ShareException {

        MultipleSharesContainer container1 = secrets2_shares3_threshold3.get(0);
        MultipleSharesContainer container2 = secrets2_shares3_threshold3.get(1);

        assertTrue(container1.isCompatible(container2));
    }

    @Test
    public void whenCheckIfTwoNotCompatibleContainersAreCompatibleThenFalse() throws ShareException {

        MultipleSharesContainer container1 = secrets2_shares2_threshold2.get(0);
        MultipleSharesContainer container2 = secrets2_shares3_threshold3.get(1);

        assertFalse(container1.isCompatible(container2));
    }

    @Test
    public void whenCheckThenOk() throws ShareException {

        secrets2_shares2_threshold2.get(0).check();
        secrets2_shares2_threshold2.get(1).check();
        secrets2_shares3_threshold3.get(0).check();
        secrets2_shares3_threshold3.get(1).check();
    }

    @Test
    public void whenGetSharesThenExpectedSize() throws ShareException {

        assertTrue(secrets2_shares2_threshold2.get(0).getShares().size() == 2);
        assertTrue(secrets2_shares2_threshold2.get(1).getShares().size() == 2);
        assertTrue(secrets2_shares3_threshold3.get(0).getShares().size() == 2);
        assertTrue(secrets2_shares3_threshold3.get(1).getShares().size() == 2);
    }

    @Test
    public void whenGetNumberOfPartsThenExpectedSize() throws ShareException {

        assertTrue(secrets2_shares2_threshold2.get(0).getNumberOfParts() == 2);
        assertTrue(secrets2_shares2_threshold2.get(1).getNumberOfParts() == 2);
        assertTrue(secrets2_shares3_threshold3.get(0).getNumberOfParts() == 3);
        assertTrue(secrets2_shares3_threshold3.get(1).getNumberOfParts() == 3);
    }

    @Test
    public void whenGetModulusThenExpectedSize() throws ShareException {

        int bitLenth = secrets2_shares2_threshold2.get(0).getModulus().bitLength();
        assertTrue((511 <= bitLenth) && (bitLenth <= 512));

        bitLenth = secrets2_shares2_threshold2.get(1).getModulus().bitLength();
        assertTrue((511 <= bitLenth) && (bitLenth <= 512));

        bitLenth = secrets2_shares3_threshold3.get(0).getModulus().bitLength();
        assertTrue((511 <= bitLenth) && (bitLenth <= 512));

        bitLenth = secrets2_shares3_threshold3.get(1).getModulus().bitLength();
        assertTrue((511 <= bitLenth) && (bitLenth <= 512));
    }

    @Test
    public void whenGetPointThenException() throws ShareException {

        exception.expect(UnsupportedOperationException.class);

        secrets2_shares2_threshold2.get(0).getPoint();
    }

    @Test
    public void whenSplitAndRecoverStringsThenOk() throws ShareException {

        int no = 2;
        int threshold = 2;

        String originalData1 = getSmallString1();
        String originalData2 = getSmallString2();

        ////////////////////////////////////////////////////
        //
        // split each piece of data into shares
        //
        ////////////////////////////////////////////////////

        Set<Share> shares1 = sa.split(originalData1.getBytes(StandardCharsets.UTF_8), no, threshold);
        Set<Share> shares2 = sa.split(originalData2.getBytes(StandardCharsets.UTF_8), no, threshold);

        ////////////////////////////////////////////////////
        //
        // arrange data into the structures to stored in the containers
        //
        ////////////////////////////////////////////////////

        List<Share> shares1AsList = new ArrayList<Share>(shares1);
        List<Share> shares2AsList = new ArrayList<Share>(shares2);

        List<byte[]> sharesForCard1 = new ArrayList<byte[]>();
        sharesForCard1.add(shares1AsList.get(0).serialize());
        sharesForCard1.add(shares2AsList.get(0).serialize());

        List<byte[]> sharesForCard2 = new ArrayList<byte[]>();
        sharesForCard2.add(shares1AsList.get(1).serialize());
        sharesForCard2.add(shares2AsList.get(1).serialize());

        MultipleSharesContainer container1 = new MultipleSharesContainer(no, threshold, sharesForCard1);
        MultipleSharesContainer container2 = new MultipleSharesContainer(no, threshold, sharesForCard2);

        ////////////////////////////////////////////////////
        //
        // extract data from the containers
        //
        ////////////////////////////////////////////////////

        List<byte[]> fromcard1 = container1.getShares();
        List<byte[]> fromcard2 = container2.getShares();

        Set<Share> data1Shares = new HashSet<Share>();
        data1Shares.add(new Share(fromcard1.get(0)));
        data1Shares.add(new Share(fromcard2.get(0)));

        Set<Share> data2Shares = new HashSet<Share>();
        data2Shares.add(new Share(fromcard1.get(1)));
        data2Shares.add(new Share(fromcard2.get(1)));

        byte[] recoveredData1 = sa.recover(data1Shares);
        byte[] recoveredData2 = sa.recover(data2Shares);

        String string1 = new String(recoveredData1, StandardCharsets.UTF_8);
        String string2 = new String(recoveredData2, StandardCharsets.UTF_8);

        assertEquals(originalData1, string1);
        assertEquals(originalData2, string2);
    }

    @Test
    public void whenSplitAndRecoverElGamalPrivateKeyThenOk()
            throws ShareException, GeneralCryptoLibException, KeyException {

        int numberShares = 2;
        int threshold = 2;

        // ENCRYPTION WHERE P AND Q HAVE 2048 BITS (EXPONENT IS 256 BYTES)
        BigInteger p = new BigInteger(
            "16370518994319586760319791526293535327576438646782139419846004180837103527129035954742043590609421369665944746587885814920851694546456891767644945459124422553763416586515339978014154452159687109161090635367600349264934924141746082060353483306855352192358732451955232000593777554431798981574529854314651092086488426390776811367125009551346089319315111509277347117467107914073639456805159094562593954195960531136052208019343392906816001017488051366518122404819967204601427304267380238263913892658950281593755894747339126531018026798982785331079065126375455293409065540731646939808640273393855256230820509217411510058759");
        BigInteger q = new BigInteger(
            "8185259497159793380159895763146767663788219323391069709923002090418551763564517977371021795304710684832972373293942907460425847273228445883822472729562211276881708293257669989007077226079843554580545317683800174632467462070873041030176741653427676096179366225977616000296888777215899490787264927157325546043244213195388405683562504775673044659657555754638673558733553957036819728402579547281296977097980265568026104009671696453408000508744025683259061202409983602300713652133690119131956946329475140796877947373669563265509013399491392665539532563187727646704532770365823469904320136696927628115410254608705755029379");
        BigInteger g = new BigInteger("2");
        ZpSubgroup group = new ZpSubgroup(g, p, q);

        ElGamalEncryptionParameters elGamalEncryptionParameters = new ElGamalEncryptionParameters(p, q, g);

        ElGamalServiceAPI elgamalServiceAPI = new ElGamalService();

        int numRequiredSubkeys = 3;
        ElGamalKeyPairGenerator elGamalKeyPairGenerator =
            new ElGamalKeyPairGenerator(elGamalEncryptionParameters, numRequiredSubkeys, elgamalServiceAPI);

        KeyPair keyPair = elGamalKeyPairGenerator.generate();

        PrivateKey privateKey = keyPair.getPrivate();
        PublicKey publickey = keyPair.getPublic();

        ElGamalPrivateKeyAdapter elGamalPrivateKeyAdapter = null;
        if (privateKey instanceof ElGamalPrivateKeyAdapter) {
            elGamalPrivateKeyAdapter = (ElGamalPrivateKeyAdapter) privateKey;
        }
        ElGamalPrivateKey elGamalPrivateKey = elGamalPrivateKeyAdapter.getPrivateKey();

        ElGamalPublicKeyAdapter elGamalPublicKeyAdapter = null;
        if (publickey instanceof ElGamalPublicKeyAdapter) {
            elGamalPublicKeyAdapter = (ElGamalPublicKeyAdapter) publickey;
        }
        ElGamalPublicKey elGamalPublicKey = elGamalPublicKeyAdapter.getPublicKey();

        List<Exponent> subkeys = elGamalPrivateKey.getKeys();
        int numSubkeys = subkeys.size();

        ////////////////////////////////////////////////////
        //
        // split each piece of data into shares
        //
        ////////////////////////////////////////////////////

        List<List<Share>> allShares = new ArrayList<List<Share>>();
        for (int i = 0; i < subkeys.size(); i++) {

            BigInteger value = subkeys.get(i).getValue();
            byte[] secret = value.toByteArray();
            Set<Share> shareForSubkey = sa.split(secret, numberShares, threshold);
            allShares.add(new ArrayList<Share>(shareForSubkey));
        }

        ////////////////////////////////////////////////////
        //
        // arrange data into the structures to be stored in the containers
        //
        ////////////////////////////////////////////////////

        List<MultipleSharesContainer> containers = new ArrayList<MultipleSharesContainer>();

        for (int i = 0; i < numberShares; i++) {

            List<byte[]> sharesForCard = new ArrayList<byte[]>();

            for (int j = 0; j < numSubkeys; j++) {
                sharesForCard.add(allShares.get(j).get(i).serialize());
            }
            containers.add(new MultipleSharesContainer(numberShares, threshold, sharesForCard));
        }

        ////////////////////////////////////////////////////
        //
        // extract data from the containers
        //
        ////////////////////////////////////////////////////

        List<Exponent> exps = new ArrayList<Exponent>();

        for (int i = 0; i < numSubkeys; i++) {

            Set<Share> sharesForSingleSubkey = new HashSet<Share>();
            for (int j = 0; j < numberShares; j++) {
                List<byte[]> fromcard1 = containers.get(j).getShares();
                sharesForSingleSubkey.add(new Share(fromcard1.get(i)));
            }

            byte[] recoveredData1 = sa.recover(sharesForSingleSubkey);
            Exponent recovered1 = new Exponent(q, new BigInteger(recoveredData1));
            exps.add(recovered1);
        }

        ElGamalPrivateKey recoveredKey = new ElGamalPrivateKey(exps, group);

        assertEquals(elGamalPrivateKey, recoveredKey);
    }

    private boolean checkIfContainersAreEqual(MultipleSharesContainer c1, MultipleSharesContainer c2) {

        if (c1 == null || c2 == null) {
            return false;
        }

        List<byte[]> sharesC1 = c1.getShares();
        List<byte[]> sharesC2 = c2.getShares();

        if (sharesC1.size() != sharesC2.size()) {
            return false;
        }

        for (int i = 0; i < sharesC1.size(); i++) {
            if (!Arrays.equals(sharesC1.get(i), sharesC2.get(i))) {
                return false;
            }
        }

        return true;
    }

    private static List<MultipleSharesContainer> createListOfContainers_secrets2_shares2_threshold2()
            throws ShareException {

        int no = 2;
        int threshold = 2;

        String originalData1 = getSmallString1();
        String originalData2 = getSmallString2();

        ////////////////////////////////////////////////////
        //
        // split each piece of data into shares
        //
        ////////////////////////////////////////////////////

        Set<Share> shares1 = sa.split(originalData1.getBytes(StandardCharsets.UTF_8), no, threshold);
        Set<Share> shares2 = sa.split(originalData2.getBytes(StandardCharsets.UTF_8), no, threshold);

        ////////////////////////////////////////////////////
        //
        // arrange data into the structures to stored in the containers
        //
        ////////////////////////////////////////////////////

        List<Share> shares1AsList = new ArrayList<Share>(shares1);
        List<Share> shares2AsList = new ArrayList<Share>(shares2);

        List<byte[]> sharesForCard1 = new ArrayList<byte[]>();
        sharesForCard1.add(shares1AsList.get(0).serialize());
        sharesForCard1.add(shares2AsList.get(0).serialize());

        List<byte[]> sharesForCard2 = new ArrayList<byte[]>();
        sharesForCard2.add(shares1AsList.get(1).serialize());
        sharesForCard2.add(shares2AsList.get(1).serialize());

        MultipleSharesContainer container1 = new MultipleSharesContainer(no, threshold, sharesForCard1);
        MultipleSharesContainer container2 = new MultipleSharesContainer(no, threshold, sharesForCard2);

        List<MultipleSharesContainer> listMultipleSharesContainers = new ArrayList<MultipleSharesContainer>();
        listMultipleSharesContainers.add(container1);
        listMultipleSharesContainers.add(container2);

        return listMultipleSharesContainers;
    }

    private static List<MultipleSharesContainer> createListOfContainers_secrets2_shares3_threshold3()
            throws ShareException {

        int no = 3;
        int threshold = 3;

        String originalData1 = getSmallString1();
        String originalData2 = getSmallString2();

        ////////////////////////////////////////////////////
        //
        // split each piece of data into shares
        //
        ////////////////////////////////////////////////////

        Set<Share> shares1 = sa.split(originalData1.getBytes(StandardCharsets.UTF_8), no, threshold);
        Set<Share> shares2 = sa.split(originalData2.getBytes(StandardCharsets.UTF_8), no, threshold);

        ////////////////////////////////////////////////////
        //
        // arrange data into the structures to stored in the containers
        //
        ////////////////////////////////////////////////////

        List<Share> shares1AsList = new ArrayList<Share>(shares1);
        List<Share> shares2AsList = new ArrayList<Share>(shares2);

        List<byte[]> sharesForCard1 = new ArrayList<byte[]>();
        sharesForCard1.add(shares1AsList.get(0).serialize());
        sharesForCard1.add(shares2AsList.get(0).serialize());

        List<byte[]> sharesForCard2 = new ArrayList<byte[]>();
        sharesForCard2.add(shares1AsList.get(1).serialize());
        sharesForCard2.add(shares2AsList.get(1).serialize());

        List<byte[]> sharesForCard3 = new ArrayList<byte[]>();
        sharesForCard3.add(shares1AsList.get(2).serialize());
        sharesForCard3.add(shares2AsList.get(2).serialize());

        MultipleSharesContainer container1 = new MultipleSharesContainer(no, threshold, sharesForCard1);
        MultipleSharesContainer container2 = new MultipleSharesContainer(no, threshold, sharesForCard2);
        MultipleSharesContainer container3 = new MultipleSharesContainer(no, threshold, sharesForCard3);

        List<MultipleSharesContainer> listMultipleSharesContainers = new ArrayList<MultipleSharesContainer>();
        listMultipleSharesContainers.add(container1);
        listMultipleSharesContainers.add(container2);
        listMultipleSharesContainers.add(container3);

        return listMultipleSharesContainers;
    }

    private static String getSmallString1() {
        return "aaaaaaaaaaaaaaaaaaaaaaaaaaaa111111111111111111111111111111111111";
    }

    private static String getSmallString2() {
        return "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb222222222222222222222222222222222222222222";
    }

    private static String getSmallString3() {
        return "DDDDDDDDDDDDDDDDDDDDDDDDDDAAAAAAAAAAAAAAAAAAa";
    }
}
