package meerkat.crypto.concrete;

import com.google.protobuf.ByteString;
import com.google.protobuf.Message;
import meerkat.protobuf.Crypto;
import org.junit.Before;
import org.junit.Test;

import java.io.InputStream;
import java.math.BigInteger;
import java.security.KeyStore;
import java.util.Date;
import java.util.Random;

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

/**
 * Created by talm on 12/11/15.
 */
public class ECDSASignatureTest {
    public static String KEYFILE_EXAMPLE = "/certs/enduser-certs/user1-key-with-password-secret.p12";
    public static String KEYFILE_PASSWORD = "secret";

    public static String CERT1_PEM_EXAMPLE = "/certs/enduser-certs/user1.crt";
    public static String CERT2_DER_EXAMPLE = "/certs/enduser-certs/user2.der";

    public static String MSG_PLAINTEXT_EXAMPLE = "/certs/signed-messages/helloworld.txt";
    public static String MSG_SIG_EXAMPLE = "/certs/signed-messages/helloworld.txt.sha256sig";

    public static String HELLO_WORLD = "hello world!";

    public final static int REPEAT_COUNT = 10;

    Random rand = new Random(0);

    protected ECDSASignature signer;

    protected ECDSASignature getSigner() { return new ECDSASignature(); }

    @Before
    public void setup() throws Exception {
        signer = getSigner();
    }

    @Test
    public void loadSignatureKey() throws Exception {
        InputStream keyStream = getClass().getResourceAsStream(KEYFILE_EXAMPLE);
        char[] password = KEYFILE_PASSWORD.toCharArray();

        KeyStore.Builder keyStore = signer.getPKCS12KeyStoreBuilder(keyStream, password);
        signer.loadSigningCertificate(keyStore);
        keyStream.close();
    }

    @Test
    public void loadPEMVerificationKey() throws Exception {
        InputStream certStream = getClass().getResourceAsStream(CERT1_PEM_EXAMPLE);

        signer.loadVerificationCertificates(certStream);
        certStream.close();
    }

    @Test
    public void loadDERVerificationKey() throws Exception {
        InputStream certStream = getClass().getResourceAsStream(CERT2_DER_EXAMPLE);

        signer.loadVerificationCertificates(certStream);
        certStream.close();
    }


    @Test
    public void verifyValidSig() throws Exception {
        InputStream certStream = getClass().getResourceAsStream(CERT1_PEM_EXAMPLE);
        InputStream msgStream = getClass().getResourceAsStream(MSG_PLAINTEXT_EXAMPLE);
        InputStream sigStream = getClass().getResourceAsStream(MSG_SIG_EXAMPLE);

        signer.loadVerificationCertificates(certStream);
        certStream.close();

        Crypto.Signature.Builder sig = Crypto.Signature.newBuilder();
        sig.setType(Crypto.SignatureType.ECDSA);
        sig.setSignerId(signer.loadedCertificates.entrySet().iterator().next().getKey());
        sig.setData(ByteString.readFrom(sigStream));

        Crypto.Signature builtSig = sig.build();
        signer.initVerify(builtSig);
        signer.updateContent(msgStream);
        assertTrue("Signature did not verify!", signer.verify());
    }

    @Test
    public void verifyInvalidSig() throws Exception {
        InputStream certStream = getClass().getResourceAsStream(CERT1_PEM_EXAMPLE);
        InputStream msgStream = getClass().getResourceAsStream(MSG_PLAINTEXT_EXAMPLE);
        InputStream sigStream = getClass().getResourceAsStream(MSG_SIG_EXAMPLE);

        signer.loadVerificationCertificates(certStream);
        certStream.close();

        Crypto.Signature.Builder sig = Crypto.Signature.newBuilder();
        sig.setType(Crypto.SignatureType.ECDSA);
        sig.setSignerId(signer.loadedCertificates.entrySet().iterator().next().getKey());
        byte[] sigData = ByteString.readFrom(sigStream).toByteArray();
        ++sigData[0];

        sig.setData(ByteString.copyFrom(sigData));


        Crypto.Signature builtSig = sig.build();
        signer.initVerify(builtSig);
        signer.updateContent(msgStream);
        assertFalse("Bad Signature passed verification!", signer.verify());
    }


    @Test
    public void verifyInvalidMsg() throws Exception {
        InputStream certStream = getClass().getResourceAsStream(CERT1_PEM_EXAMPLE);
        InputStream msgStream = getClass().getResourceAsStream(MSG_PLAINTEXT_EXAMPLE);
        InputStream sigStream = getClass().getResourceAsStream(MSG_SIG_EXAMPLE);

        signer.loadVerificationCertificates(certStream);
        certStream.close();

        Crypto.Signature.Builder sig = Crypto.Signature.newBuilder();
        sig.setType(Crypto.SignatureType.ECDSA);
        sig.setSignerId(signer.loadedCertificates.entrySet().iterator().next().getKey());
        sig.setData(ByteString.readFrom(sigStream));
        byte[] msgData = ByteString.readFrom(msgStream).toByteArray();
        ++msgData[0];

        Crypto.Signature builtSig = sig.build();
        signer.initVerify(builtSig);
        signer.updateContent(msgStream);
        assertFalse("Signature doesn't match message but passed verification!", signer.verify());
    }



    protected void loadSigningKeys() throws Exception {
        InputStream keyStream = getClass().getResourceAsStream(KEYFILE_EXAMPLE);
        char[] password = KEYFILE_PASSWORD.toCharArray();

        KeyStore.Builder keyStore = signer.getPKCS12KeyStoreBuilder(keyStream, password);
        signer.loadSigningCertificate(keyStore);
    }

    @Test
    public void signAndVerify() throws Exception {
        loadSigningKeys();

        BigInteger rawMsg = new BigInteger(50, rand);
        Crypto.BigInteger usMsg = Crypto.BigInteger.newBuilder()
                .setData(ByteString.copyFrom(rawMsg.toByteArray())).build();

        signer.updateContent(usMsg);
        Crypto.Signature sig = signer.sign();

        signer.loadVerificationCertificates(getClass().getResourceAsStream(CERT1_PEM_EXAMPLE));

        signer.initVerify(sig);
        signer.updateContent(usMsg);
        assertTrue("Couldn't verify signature on ", signer.verify());
    }


    @Test
    public void signMultipleAndVerify() throws Exception {
        loadSigningKeys();

        Message[] msgs = new Message[REPEAT_COUNT];
        for (int i = 0; i < msgs.length; ++i) {

            BigInteger rawMsg = new BigInteger(50, rand);
            msgs[i] = Crypto.BigInteger.newBuilder()
                    .setData(ByteString.copyFrom(rawMsg.toByteArray())).build();
            signer.updateContent(msgs[i]);
        }

        Crypto.Signature sig = signer.sign();

        signer.loadVerificationCertificates(getClass().getResourceAsStream(CERT1_PEM_EXAMPLE));

        signer.initVerify(sig);
        for (int i = 0; i < msgs.length; ++i) {
            signer.updateContent(msgs[i]);
        }
        assertTrue("Couldn't verify signature on ", signer.verify());
    }

    @Test
    public void multipleSignAndVerify() throws Exception {
        loadSigningKeys();

        Message[] msgs = new Message[REPEAT_COUNT];
        Crypto.Signature[] sigs = new Crypto.Signature[REPEAT_COUNT];
        for (int i = 0; i < msgs.length; ++i) {
            BigInteger rawMsg = new BigInteger(50, rand);
            msgs[i] = Crypto.BigInteger.newBuilder()
                    .setData(ByteString.copyFrom(rawMsg.toByteArray())).build();
            signer.updateContent(msgs[i]);
            sigs[i] = signer.sign();
        }

        signer.loadVerificationCertificates(getClass().getResourceAsStream(CERT1_PEM_EXAMPLE));


        for (int i = 0; i < msgs.length; ++i) {
            signer.initVerify(sigs[i]);
            signer.updateContent(msgs[i]);
            assertTrue("Couldn't verify signature on ", signer.verify());
        }

    }


    @Test
    public void generateKeysSignAndVerify() throws Exception {
        signer.generateSigningCertificate(BigInteger.ONE, new Date(), new Date(System.currentTimeMillis() + 1000 * 60 * 60 * 24 * 365), "testing");

        Crypto.SignatureVerificationKey verificationKey = signer.getSignerPublicKey();

        signer.loadVerificationCertificate(verificationKey);

        BigInteger rawMsg = new BigInteger(50, rand);
        Crypto.BigInteger usMsg = Crypto.BigInteger.newBuilder()
                .setData(ByteString.copyFrom(rawMsg.toByteArray())).build();

        signer.updateContent(usMsg);
        Crypto.Signature sig = signer.sign();

        signer.initVerify(sig);
        signer.updateContent(usMsg);
        assertTrue("Couldn't verify signature on ", signer.verify());
    }

}
