package meerkat.mixer.main;

import com.google.protobuf.StringValue;
import meerkat.crypto.concrete.ECElGamalEncryption;
import meerkat.crypto.concrete.Util;
import meerkat.protobuf.Crypto;
import org.factcenter.qilin.primitives.concrete.ECElGamal;
import org.junit.Before;
import org.junit.Test;

import java.io.*;
import java.util.Random;

import static org.junit.Assert.*;

/**
 * Created by talm on 21/01/17.
 */
public class MixTest {
    public Mix mix;
    Random rand = new Random(1);

    @Before
    public void setUp() throws Exception {
        mix = new Mix();
    }

    @Test
    public void testKeygen() throws Exception {
        File tmpKeys = File.createTempFile("elgamal", "key");
        mix.createKeypair(tmpKeys);

        ECElGamal.SK secretKey = mix.secretKey;
        assertNotNull(secretKey);
        assertNotNull(mix.serializedPk);

        Mix newMix = new Mix();

        newMix.loadKeypair(tmpKeys);
        assertEquals(mix.serializedPk, newMix.serializedPk);

        StringValue tst = StringValue.newBuilder().setValue("test").build();

        Crypto.RerandomizableEncryptedMessage cipher = mix.enc.encrypt(tst, mix.enc.generateRandomness(rand));
        StringValue newTst = Util.decrypt(StringValue.class, newMix.secretKey, newMix.enc.getGroup(), cipher);

        assertEquals(tst, newTst);

    }

    File createPlaintexts(int n) throws Exception {
        File tmpData = File.createTempFile("plaintext", "txt");
        PrintStream out = new PrintStream(tmpData);

        for (int i = 0; i < n; ++i) {
            out.println("Line " + i);
        }

        out.close();

        return tmpData;
    }

    @Test
    public void testEncryptDecrypt() throws Exception {
        File tmpKeys = File.createTempFile("elgamal", "key");
        mix.createKeypair(tmpKeys);

        int n = 35;
        File plaintexts = createPlaintexts(n);
        File ciphertexts = File.createTempFile("ciphertexts", "bin");
        mix.encrypt(plaintexts, ciphertexts);

        Mix newMix = new Mix();
        newMix.loadKeypair(tmpKeys);
        newMix.loadMix(ciphertexts);
        File newPlaintexts = File.createTempFile("plaintexts2", "txt");
        newMix.decrypt(newPlaintexts);

        BufferedReader in1 = new BufferedReader(new FileReader(plaintexts));
        BufferedReader in2 = new BufferedReader(new FileReader(newPlaintexts));

        int n2 = 1 << newMix.inMix.getLogN();
        assert(n2 >= n);

        for (int i = 0; i < n; ++i) {
            String line1 = in1.readLine();
            String line2 = in2.readLine();
            assertEquals(line1, line2);
        }

        for (int i = n; i < n2; ++i) {
            String line2 = in2.readLine();
            assertEquals("", line2);
        }
    }

    @Test
    public void testMixVerify() throws Exception {
        File tmpKeys = File.createTempFile("elgamal", "key");
        mix.createKeypair(tmpKeys);

        int n = 35;
        File plaintexts = createPlaintexts(n);
        File ciphertexts = File.createTempFile("ciphertexts", "enc");
        mix.encrypt(plaintexts, ciphertexts);

        File mixedFile = File.createTempFile("mixed", "enc");
        Mix newMix1 = new Mix();
        newMix1.loadKeypair(tmpKeys);
        newMix1.loadMix(ciphertexts);
        newMix1.mix(mixedFile);

        Mix newMix2 = new Mix();
        newMix2.loadKeypair(tmpKeys);
        newMix2.loadMix(mixedFile);
        boolean ok = newMix2.verify(true);
        assertTrue(ok);
    }
}