/*
 * 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.commons.sign;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;

import java.io.IOException;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import javax.naming.InvalidNameException;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.BeforeClass;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.scytl.cryptolib.api.asymmetric.AsymmetricServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.asymmetric.service.AsymmetricServiceFactoryHelper;
import com.scytl.products.ov.commons.beans.Ballot;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignature;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadSignatureException;
import com.scytl.products.ov.commons.beans.domain.model.messaging.PayloadVerificationException;
import com.scytl.products.ov.commons.beans.domain.model.vote.CiphertextEncryptedVote;
import com.scytl.products.ov.commons.beans.domain.model.vote.EncryptedVote;
import com.scytl.products.ov.commons.beans.exceptions.LambdaException;
import com.scytl.products.ov.commons.dto.mixing.MixingPayload;
import com.scytl.products.ov.commons.dto.mixing.MixingPayloadMockBuilder;
import com.scytl.products.ov.commons.verify.PayloadSigningCertificateValidator;
import com.scytl.products.ov.commons.verify.CryptolibPayloadSigningCertificateValidator;
import com.scytl.products.ov.commons.verify.CryptolibPayloadVerifier;

/**
 * Tests the Cryptolib-based payload signing and verification services.
 */
public class CryptolibPayloadSignerTest {

    private static final Logger logger = LoggerFactory.getLogger(CryptolibPayloadSignerTest.class);

    private static final String BALLOT_JSON = "ballot.json";

    private static final int NUM_VOTES = 10;

    private static final int ROUNDS = 25;

    private static final ExecutorService signingExecutor = Executors.newFixedThreadPool(ROUNDS);

    private static final ExecutorService verificationExecutor = Executors.newFixedThreadPool(ROUNDS);

    private static TestCertificateGenerator testCertificatesGenerator;

    private static AsymmetricServiceAPI asymmetricService;

    private static PayloadSigningCertificateValidator certificateChainValidator;

    private static SignedTestVotingDataService defaultVotingDataService;

    private static List<EncryptedVote> defaultVotes;

    private static CryptolibPayloadSigner signer;

    private static CryptolibPayloadVerifier verifier;

    @BeforeClass
    public static void init() throws GeneralCryptoLibException, IOException, InvalidNameException {
        Security.addProvider(new BouncyCastleProvider());

        // Set up the signing, verification and certification services.
        asymmetricService = AsymmetricServiceFactoryHelper.getFactoryOfThreadSafeServices().create();
        certificateChainValidator = new CryptolibPayloadSigningCertificateValidator();

        testCertificatesGenerator = TestCertificateGenerator.createDefault();

        // Set up votes.
        defaultVotingDataService =
            new SignedTestVotingDataService(getBallotDefinition(), asymmetricService, testCertificatesGenerator);
        defaultVotes = defaultVotingDataService.generateVotes(NUM_VOTES).stream().map(CiphertextEncryptedVote::new)
            .collect(Collectors.toList());

        // Set up service under test.
        signer = new CryptolibPayloadSigner(asymmetricService);
        verifier = new CryptolibPayloadVerifier(asymmetricService, certificateChainValidator);
    }

    /**
     * Get a ballot definition from a JSON file.
     *
     * @return the ballot definition
     */
    private static Ballot getBallotDefinition() throws IOException {
        return new ObjectMapper().readValue(
            CryptolibPayloadSignerTest.class.getClassLoader().getResourceAsStream(BALLOT_JSON), Ballot.class);
    }

    @Test
    public void signAndValidateADeserialisedPayload()
            throws PayloadSignatureException, PayloadVerificationException, IOException, GeneralCryptoLibException {
        MixingPayload payload = new MixingPayloadMockBuilder().withSignature(mock(PayloadSignature.class)).build();

        // Sign and store the signature in the payload.
        payload.setSignature(signer.sign(payload, defaultVotingDataService.getSigningKey(),
            defaultVotingDataService.getCertificateChain()));

        // Verify that the signature still holds.
        assertTrue(verifier.isValid(payload, defaultVotingDataService.getCertificateChain()[1]));
    }

    @Test
    public void notValidateATamperedWithPayload()
            throws PayloadSignatureException, PayloadVerificationException, IOException, GeneralCryptoLibException {
        // Build a payload, sign it and verify it.
        MixingPayload payload = new MixingPayloadMockBuilder().withVotes(defaultVotes).build();
        payload.setSignature(signer.sign(payload, defaultVotingDataService.getSigningKey(),
            defaultVotingDataService.getCertificateChain()));

        // The signature should be valid.
        assertTrue(verifier.isValid(payload, defaultVotingDataService.getCertificateChain()[1]));

        // Now, alter data in the payload.
        payload.getVotes().remove(payload.getVotes().iterator().next());

        // The signature should no longer be valid.
        assertFalse(verifier.isValid(payload, defaultVotingDataService.getCertificateChain()[1]));
    }

    @Test
    public void signConcurrently()
            throws PayloadSignatureException, IOException, GeneralCryptoLibException, InterruptedException,
            ExecutionException, PayloadVerificationException {

        List<MixingPayload> payloads = new ArrayList<>();
        List<Future<MixingPayload>> signedPayloadFutures = new ArrayList<>();
        // Generate and sign the test payloads concurrently.
        for (int i = 0; i < ROUNDS; i++) {
            signedPayloadFutures.add(signingExecutor.submit(() -> {
                SignedTestVotingDataService votingDataService = new SignedTestVotingDataService(getBallotDefinition(),
                    asymmetricService, testCertificatesGenerator);
                PrivateKey signingKey = votingDataService.getSigningKey();
                X509Certificate[] certificateChain = votingDataService.getCertificateChain();
                MixingPayload payload = new MixingPayloadMockBuilder().build();
                payload.setSignature(signer.sign(payload, signingKey, certificateChain));

                payloads.add(payload);
                return payload;
            }));
        }

        // The signing executor is no longer accepting any tasks.
        signingExecutor.shutdown();

        for (Future<MixingPayload> mixingPayloadFuture : signedPayloadFutures) {
            payloads.add(mixingPayloadFuture.get());
        }

        // Establish the root certificate for all validations.
        final X509Certificate rootCertificate = testCertificatesGenerator.getRootCertificate();

        // Validate all payloads sequentially to ensure no concurrency issues
        // interfere.
        for (MixingPayload payload : payloads) {
            assertTrue("Sequential validation of payloads has failed", verifier.isValid(payload, rootCertificate));
        }

        // Now run all validations concurrently.
        List<Future<Boolean>> verifications = verificationExecutor.invokeAll(payloads.stream()
            .map(payload -> new PayloadVerifier(payload, rootCertificate)).collect(Collectors.toSet()));

        // The verification executor is no longer accepting any tasks.
        verificationExecutor.shutdown();

        // Make sure that all payloads' signatures are valid.
        verifications.forEach(verification -> {
            try {
                assertTrue("All payloads should be valid", verification.get());
            } catch (InterruptedException | ExecutionException e) {
                throw new LambdaException(e);
            }
        });

        // Make sure the executors have been shut down.
        assertTrue("Payload signing did not finish in time", signingExecutor.awaitTermination(5, TimeUnit.SECONDS));
        assertTrue("Payload verification did not finish in time",
            verificationExecutor.awaitTermination(5, TimeUnit.SECONDS));
    }

    private class PayloadVerifier implements Callable<Boolean> {
        private final MixingPayload payload;

        private final X509Certificate rootCertificate;

        public PayloadVerifier(MixingPayload payload, X509Certificate rootCertificate) {
            this.payload = payload;
            this.rootCertificate = rootCertificate;
        }

        @Override
        public Boolean call() throws Exception {
            logger.info("Checking {}... \n", Thread.currentThread().getId());
            boolean isValid = verifier.isValid(payload, rootCertificate);

            logger.info("Payload in thread {} is {}\n", Thread.currentThread().getId(),
                isValid ? "valid" : "not valid");

            return isValid;
        }
    }
    // });
}
