/*
 * 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.channel.cc.commons.test;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import org.apache.commons.lang.StringUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.annotation.DirtiesContext.ClassMode;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.TestPropertySource;
import org.springframework.test.context.junit4.SpringRunner;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.api.primitives.PrimitivesServiceAPI;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.cryptolib.primitives.service.PrimitivesService;
import com.scytl.products.ov.channel.cc.choicecodes.generation.service.ChoiceCodesGenerationService;
import com.scytl.products.ov.channel.cc.choicecodes.verification.service.ChoiceCodesComputationService;
import com.scytl.products.ov.channel.cc.choicecodes.verification.service.ComputedVerificationCardExistsException;
import com.scytl.products.ov.channel.cc.choicecodes.verification.service.ComputedVerificationCardsRepository;
import com.scytl.products.ov.commons.beans.domain.model.confirmation.ConfirmationMessage;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignature;
import com.scytl.products.ov.commons.beans.domain.model.vote.Vote;
import com.scytl.products.ov.commons.dto.BallotCastingKeyVerificationInput;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationDTO;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationInput;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationReqPayload;
import com.scytl.products.ov.commons.dto.ChoiceCodeGenerationResPayload;
import com.scytl.products.ov.commons.dto.ChoiceCodeVerificationDTO;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationInput;
import com.scytl.products.ov.commons.dto.ChoiceCodesVerificationResPayload;
import com.scytl.products.ov.commons.dto.ElGamalComputationsValuesCodec;
import com.scytl.products.ov.commons.dto.PartialCodesInput;
import com.scytl.products.ov.commons.messaging.MessageListener;
import com.scytl.products.ov.commons.messaging.MessagingService;
import com.scytl.products.ov.commons.ui.ws.rs.ObjectMappers;

@RunWith(SpringRunner.class)
@ContextConfiguration(classes = {TestConfiguration.class })
@DirtiesContext(classMode = ClassMode.AFTER_EACH_TEST_METHOD)
@TestPropertySource(properties = {
        "queues.cg-comp.res=" + VoteCastConfirmationTest.TEST_GENERATION_COMPUTATION_RESPONSE_QUEUE,
        "queues.cv-comp.res=" + VoteCastConfirmationTest.TEST_VERIFICATION_COMPUTATION_RESPONSE_QUEUE })
public class VoteCastConfirmationTest {

    static final String TEST_GENERATION_COMPUTATION_RESPONSE_QUEUE = "cg-comp-res";

    static final String TEST_VERIFICATION_COMPUTATION_RESPONSE_QUEUE = "cv-comp-res";

    private static ZpSubgroup encryptionParameters = null;

    private final static String TENANT_ID = "100";

    private final static String ELECTION_EVENT_ID = "electionEventId";

    private final static String VERIFICATION_CARD_SET_ID = "verificationCardSetId";

    private final static String VERIFICATION_CARD_ID = "verificationCardId";

    private final static String GENERATION_REQUEST_ID = "generationRequestId";

    private final static int CHUNK_ID = 1;

    @Autowired
    MessagingService messagingService;

    @Autowired
    ChoiceCodesGenerationService choiceCodesGenerationService;

    @Autowired
    ChoiceCodesComputationService choiceCodesComputationService;

    @Autowired
    ElGamalComputationsValuesCodec codec;

    @Autowired
    ComputedVerificationCardsRepository computedVerificationCardsRepository;

    @BeforeClass
    public static void init() throws GeneralCryptoLibException {
        encryptionParameters = new ZpSubgroup(new BigInteger("2"),
            new BigInteger(
                "16370518994319586760319791526293535327576438646782139419846004180837103527129035954742043590609421369665944746587885814920851694546456891767644945459124422553763416586515339978014154452159687109161090635367600349264934924141746082060353483306855352192358732451955232000593777554431798981574529854314651092086488426390776811367125009551346089319315111509277347117467107914073639456805159094562593954195960531136052208019343392906816001017488051366518122404819967204601427304267380238263913892658950281593755894747339126531018026798982785331079065126375455293409065540731646939808640273393855256230820509217411510058759"),
            new BigInteger(
                "8185259497159793380159895763146767663788219323391069709923002090418551763564517977371021795304710684832972373293942907460425847273228445883822472729562211276881708293257669989007077226079843554580545317683800174632467462070873041030176741653427676096179366225977616000296888777215899490787264927157325546043244213195388405683562504775673044659657555754638673558733553957036819728402579547281296977097980265568026104009671696453408000508744025683259061202409983602300713652133690119131956946329475140796877947373669563265509013399491392665539532563187727646704532770365823469904320136696927628115410254608705755029379"));

        Security.addProvider(new BouncyCastleProvider());
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testVoteConfirmationSuccessful()
            throws GeneralCryptoLibException, ComputedVerificationCardExistsException, JsonProcessingException {

        // Prepare computed state of VCid
        computedVerificationCardsRepository.saveComputedVerificationCard(ELECTION_EVENT_ID, VERIFICATION_CARD_ID);

        // Vote Cast Code Generation
        List<ElGamalComputationsValues> representations = new ArrayList<>();
        representations.add(new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("941"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("3089"), encryptionParameters) })));
        representations.add(new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("577"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("631"), encryptionParameters) })));

        String encryptedRepresentations = codec.encodeList(representations);
        
        BigInteger prime = new BigInteger("9887");
        PrimitivesServiceAPI primitivesService = new PrimitivesService();
        BigInteger hashedPrime = new BigInteger(primitivesService.getHash(prime.toByteArray()));
        BigInteger hashedPrimeSquare = hashedPrime.multiply(hashedPrime);
        ZpGroupElement egHashedPrimeSquared = new ZpGroupElement(hashedPrimeSquare, encryptionParameters);

        ElGamalComputationsValues ballotCastingKey = new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {egHashedPrimeSquared, egHashedPrimeSquared }));

        String encryptedBallotCastingKey = codec.encodeSingle(ballotCastingKey);

        List<ChoiceCodeGenerationInput> choiceCodeGenerationInputList = new ArrayList<>();
        choiceCodeGenerationInputList.add(
            new ChoiceCodeGenerationInput(VERIFICATION_CARD_ID, encryptedBallotCastingKey, encryptedBallotCastingKey, encryptedRepresentations));

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationReqPayload> dtoGeneration = new ChoiceCodeGenerationDTO<>();
        dtoGeneration.setCorrelationId(UUID.randomUUID());
        dtoGeneration.setRequestId(GENERATION_REQUEST_ID);
        dtoGeneration.setPayload(new ChoiceCodeGenerationReqPayload());
        dtoGeneration.getPayload().setElectionEventId(ELECTION_EVENT_ID);
        dtoGeneration.getPayload().setTenantId(TENANT_ID);
        dtoGeneration.getPayload().setChunkId(CHUNK_ID);
        dtoGeneration.getPayload().setVerificationCardSetId(VERIFICATION_CARD_SET_ID);

        dtoGeneration.getPayload().setChoiceCodeGenerationInputList(choiceCodeGenerationInputList);

        dtoGeneration.getPayload().setSignature(Mockito.mock(PayloadSignature.class));

        ((MessageListener) choiceCodesGenerationService).onMessage(dtoGeneration);

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload> generationResult =
            (ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_GENERATION_COMPUTATION_RESPONSE_QUEUE);

        String encodedComputedBallotCastingKey =
            generationResult.getPayload().getChoiceCodeGenerationOutputList().get(0).getComputedBallotCastingKey();
        ElGamalComputationsValues computedBallotCastingKeyGeneration =
            codec.decodeSingle(encodedComputedBallotCastingKey);

        // Vote Cast Code Verification
        List<BigInteger> partialCodesElements = new ArrayList<>();
        partialCodesElements.add(prime);
        
        ChoiceCodeVerificationDTO<PartialCodesInput> dtoVerification = new ChoiceCodeVerificationDTO<>();
        dtoVerification.setCorrelationId(UUID.randomUUID());
        dtoVerification.setRequestId(GENERATION_REQUEST_ID);
        dtoVerification.setElectionEventId(ELECTION_EVENT_ID);
        dtoVerification.setVerificationCardSetId(VERIFICATION_CARD_SET_ID);
        dtoVerification.setVerificationCardId(VERIFICATION_CARD_ID);

        BallotCastingKeyVerificationInput bckVerificationInput = new BallotCastingKeyVerificationInput();
        ConfirmationMessage confirmationMessage = new ConfirmationMessage();
        String confirmationKey =
                Base64.getEncoder().encodeToString(partialCodesElements.get(0).toString().getBytes(StandardCharsets.UTF_8));
        confirmationMessage.setConfirmationKey(confirmationKey);
        bckVerificationInput.setConfirmationMessage(ObjectMappers.toJson(confirmationMessage));
        bckVerificationInput.setVotingCardId(VERIFICATION_CARD_ID);

        dtoVerification.setPayload(new PartialCodesInput());
        dtoVerification.getPayload().setBallotCastingKeyVerificationInput(bckVerificationInput);
        dtoVerification.getPayload().setPartialCodesElements(partialCodesElements);

        ((MessageListener) choiceCodesComputationService).onMessage(dtoVerification);

        ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload> verificationComputationResult =
            (ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_VERIFICATION_COMPUTATION_RESPONSE_QUEUE);

        Map<BigInteger, BigInteger> result = verificationComputationResult.getPayload().getPrimeToComputedPrime();
        List<BigInteger> computedBallotCastingKeyVerification = new ArrayList<>(result.values());

        // Test result
        Assert.assertEquals(computedBallotCastingKeyVerification.get(0),
            computedBallotCastingKeyGeneration.getValues().get(0).getValue());
    }

    // NOTE: This test omits the decryption step of the vote cast process as it
    // is not needed for the purpose of the test
    @SuppressWarnings("unchecked")
    @Test
    public void testVoteCastSuccessful() throws GeneralCryptoLibException, JsonProcessingException {

        // Vote Cast Code Generation
        List<ElGamalComputationsValues> representations = new ArrayList<>();
        representations.add(new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("941"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("3089"), encryptionParameters) })));
        representations.add(new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("577"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("631"), encryptionParameters) })));

        String encryptedRepresentations = codec.encodeList(representations);

        ElGamalComputationsValues ballotCastingKey = new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("9887"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("9887"), encryptionParameters) }));

        String encryptedBallotCastingKey = codec.encodeSingle(ballotCastingKey);

        List<ChoiceCodeGenerationInput> choiceCodeGenerationInputList = new ArrayList<>();
        choiceCodeGenerationInputList.add(
            new ChoiceCodeGenerationInput(VERIFICATION_CARD_ID, encryptedBallotCastingKey, encryptedBallotCastingKey, encryptedRepresentations));

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationReqPayload> dtoGeneration = new ChoiceCodeGenerationDTO<>();
        dtoGeneration.setCorrelationId(UUID.randomUUID());
        dtoGeneration.setRequestId(GENERATION_REQUEST_ID);
        dtoGeneration.setPayload(new ChoiceCodeGenerationReqPayload());
        dtoGeneration.getPayload().setElectionEventId(ELECTION_EVENT_ID);
        dtoGeneration.getPayload().setTenantId(TENANT_ID);
        dtoGeneration.getPayload().setChunkId(CHUNK_ID);
        dtoGeneration.getPayload().setVerificationCardSetId(VERIFICATION_CARD_SET_ID);

        dtoGeneration.getPayload().setChoiceCodeGenerationInputList(choiceCodeGenerationInputList);

        dtoGeneration.getPayload().setSignature(Mockito.mock(PayloadSignature.class));

        ((MessageListener) choiceCodesGenerationService).onMessage(dtoGeneration);

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload> generationResult =
            (ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_GENERATION_COMPUTATION_RESPONSE_QUEUE);

        Map<ElGamalComputationsValues, ElGamalComputationsValues> computedChoiceCodesGeneration = codec.decodeMap(
            generationResult.getPayload().getChoiceCodeGenerationOutputList().get(0).getComputedRepresentations());

        // Vote Choice Code Verification
        List<BigInteger> partialCodesElements = new ArrayList<>();

        for (ElGamalComputationsValues representation : representations) {
            for (ZpGroupElement zpGroupElement : representation.getValues()) {
                partialCodesElements.add(zpGroupElement.getValue());
            }
        }

        ChoiceCodeVerificationDTO<PartialCodesInput> dtoVerification = new ChoiceCodeVerificationDTO<>();
        dtoVerification.setCorrelationId(UUID.randomUUID());
        dtoVerification.setRequestId(GENERATION_REQUEST_ID);
        dtoVerification.setElectionEventId(ELECTION_EVENT_ID);
        dtoVerification.setVerificationCardSetId(VERIFICATION_CARD_SET_ID);
        dtoVerification.setVerificationCardId(VERIFICATION_CARD_ID);

        ChoiceCodesVerificationInput choiceCodesVerificationInput = new ChoiceCodesVerificationInput();
        choiceCodesVerificationInput.setCompressesedChoiceCodesEncryptionPK("UNUSED");
        Vote vote = new Vote();
        vote.setEncryptedPartialChoiceCodes(StringUtils.join(partialCodesElements, ";"));
        choiceCodesVerificationInput.setVote(ObjectMappers.toJson(vote));

        dtoVerification.setPayload(new PartialCodesInput());
        dtoVerification.getPayload().setChoiceCodesVerificationInput(choiceCodesVerificationInput);
        dtoVerification.getPayload().setPartialCodesElements(partialCodesElements);

        ((MessageListener) choiceCodesComputationService).onMessage(dtoVerification);

        ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload> verificationComputationResult =
            (ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_VERIFICATION_COMPUTATION_RESPONSE_QUEUE);

        Map<BigInteger, BigInteger> result = verificationComputationResult.getPayload().getPrimeToComputedPrime();
        List<BigInteger> computedChoiceCodesVerification = new ArrayList<>(result.values());

        // Test result
        Assert.assertEquals(computedChoiceCodesVerification.get(0),
            computedChoiceCodesGeneration.get(representations.get(0)).getValues().get(0).getValue());
        Assert.assertEquals(computedChoiceCodesVerification.get(1),
            computedChoiceCodesGeneration.get(representations.get(0)).getValues().get(1).getValue());
        Assert.assertEquals(computedChoiceCodesVerification.get(2),
            computedChoiceCodesGeneration.get(representations.get(1)).getValues().get(0).getValue());
        Assert.assertEquals(computedChoiceCodesVerification.get(3),
            computedChoiceCodesGeneration.get(representations.get(1)).getValues().get(1).getValue());
    }

    // NOTE: This test omits the decryption step of the vote cast process as it
    // is not needed for the purpose of the test
    @SuppressWarnings("unchecked")
    @Test
    public void testConfirmationCodesAttackFails()
            throws GeneralCryptoLibException, ComputedVerificationCardExistsException, JsonProcessingException {

        // Prepare computed state of VCid
        computedVerificationCardsRepository.saveComputedVerificationCard(ELECTION_EVENT_ID, VERIFICATION_CARD_ID);

        // Vote Cast Code Generation
        List<ElGamalComputationsValues> representations = new ArrayList<>();
        representations.add(new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("941"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("941"), encryptionParameters) })));

        String encryptedRepresentations = codec.encodeList(representations);

        ElGamalComputationsValues ballotCastingKey = new ElGamalComputationsValues(
            Arrays.asList(new ZpGroupElement[] {new ZpGroupElement(new BigInteger("9887"), encryptionParameters),
                    new ZpGroupElement(new BigInteger("9887"), encryptionParameters) }));

        String encryptedBallotCastingKey = codec.encodeSingle(ballotCastingKey);

        List<ChoiceCodeGenerationInput> choiceCodeGenerationInputList = new ArrayList<>();
        choiceCodeGenerationInputList.add(
            new ChoiceCodeGenerationInput(VERIFICATION_CARD_ID, encryptedBallotCastingKey, encryptedBallotCastingKey, encryptedRepresentations));

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationReqPayload> dtoGeneration = new ChoiceCodeGenerationDTO<>();
        dtoGeneration.setCorrelationId(UUID.randomUUID());
        dtoGeneration.setRequestId(GENERATION_REQUEST_ID);
        dtoGeneration.setPayload(new ChoiceCodeGenerationReqPayload());
        dtoGeneration.getPayload().setElectionEventId(ELECTION_EVENT_ID);
        dtoGeneration.getPayload().setTenantId(TENANT_ID);
        dtoGeneration.getPayload().setChunkId(CHUNK_ID);
        dtoGeneration.getPayload().setVerificationCardSetId(VERIFICATION_CARD_SET_ID);

        dtoGeneration.getPayload().setChoiceCodeGenerationInputList(choiceCodeGenerationInputList);

        dtoGeneration.getPayload().setSignature(Mockito.mock(PayloadSignature.class));

        ((MessageListener) choiceCodesGenerationService).onMessage(dtoGeneration);

        ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload> generationResult =
            (ChoiceCodeGenerationDTO<ChoiceCodeGenerationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_GENERATION_COMPUTATION_RESPONSE_QUEUE);

        Map<ElGamalComputationsValues, ElGamalComputationsValues> computedChoiceCodesGeneration = codec.decodeMap(
            generationResult.getPayload().getChoiceCodeGenerationOutputList().get(0).getComputedRepresentations());

        // Vote Choice Code Verification
        List<BigInteger> partialCodesElements = new ArrayList<>();

        // Just use the first of the representations to bypass a validation
        partialCodesElements.add(representations.get(0).getValues().get(0).getValue());

        ChoiceCodeVerificationDTO<PartialCodesInput> dtoVerification = new ChoiceCodeVerificationDTO<>();
        dtoVerification.setCorrelationId(UUID.randomUUID());
        dtoVerification.setRequestId(GENERATION_REQUEST_ID);
        dtoVerification.setElectionEventId(ELECTION_EVENT_ID);
        dtoVerification.setVerificationCardSetId(VERIFICATION_CARD_SET_ID);
        dtoVerification.setVerificationCardId(VERIFICATION_CARD_ID);

        // The attacker tries to disguise choice codes as if they were cast
        // codes into the computation
        BallotCastingKeyVerificationInput bckVerificationInput = new BallotCastingKeyVerificationInput();
        ConfirmationMessage confirmationMessage = new ConfirmationMessage();
        String confirmationKey =
                Base64.getEncoder().encodeToString(partialCodesElements.get(0).toString().getBytes(StandardCharsets.UTF_8));
        confirmationMessage.setConfirmationKey(confirmationKey);
        bckVerificationInput.setConfirmationMessage(ObjectMappers.toJson(confirmationMessage));
        bckVerificationInput.setVotingCardId(VERIFICATION_CARD_ID);

        dtoVerification.setPayload(new PartialCodesInput());
        dtoVerification.getPayload().setBallotCastingKeyVerificationInput(bckVerificationInput);
        dtoVerification.getPayload().setPartialCodesElements(partialCodesElements);

        ((MessageListener) choiceCodesComputationService).onMessage(dtoVerification);

        ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload> verificationComputationResult =
            (ChoiceCodeVerificationDTO<ChoiceCodesVerificationResPayload>) ((MessagingServiceMockImpl) messagingService)
                .getQueueMock().get(TEST_VERIFICATION_COMPUTATION_RESPONSE_QUEUE);

        Map<BigInteger, BigInteger> result = verificationComputationResult.getPayload().getPrimeToComputedPrime();
        List<BigInteger> computedChoiceCodesVerification = new ArrayList<>(result.values());

        // Test result
        Assert.assertNotEquals(computedChoiceCodesVerification.get(0),
            computedChoiceCodesGeneration.get(representations.get(0)).getValues().get(0).getValue());
    }

}
