package org.bouncycastle.pqc.crypto.test;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashMap;

import junit.framework.TestCase;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.SecretWithEncapsulation;
import org.bouncycastle.pqc.legacy.bike.BIKEKEMExtractor;
import org.bouncycastle.pqc.legacy.bike.BIKEKEMGenerator;
import org.bouncycastle.pqc.legacy.bike.BIKEKeyGenerationParameters;
import org.bouncycastle.pqc.legacy.bike.BIKEKeyPairGenerator;
import org.bouncycastle.pqc.legacy.bike.BIKEParameters;
import org.bouncycastle.pqc.legacy.bike.BIKEPrivateKeyParameters;
import org.bouncycastle.pqc.legacy.bike.BIKEPublicKeyParameters;
import org.bouncycastle.pqc.crypto.util.PrivateKeyFactory;
import org.bouncycastle.pqc.crypto.util.PrivateKeyInfoFactory;
import org.bouncycastle.pqc.crypto.util.PublicKeyFactory;
import org.bouncycastle.pqc.crypto.util.SubjectPublicKeyInfoFactory;
import org.bouncycastle.test.TestResourceFinder;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;

public class BIKETest
    extends TestCase
{

     
    public String getName()
    {
        return "BIKE Test";
    }

    public void testVectors()
        throws Exception
    {
        String[] files;
        files = new String[]{
            "PQCkemKAT_BIKE_3114.rsp",
            "PQCkemKAT_BIKE_6198.rsp",
            "PQCkemKAT_BIKE_10276.rsp"
        };

        BIKEParameters[] listParams = new BIKEParameters[]{
            BIKEParameters.bike128,
            BIKEParameters.bike192,
            BIKEParameters.bike256
        };

        for (int fileIndex = 0; fileIndex < files.length; fileIndex++)
        {
            String name = files[fileIndex];

            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/bike", name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

            String line = null;
            HashMap  buf = new HashMap ();
            TestSampler sampler = new TestSampler();
            while ((line = bin.readLine()) != null)
            {
                line = line.trim();

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        String count = (String)buf.get("count");
                        if (sampler.skipTest(count))
                        {
                            continue;
                        }

                        // System.out.println("test case: " + count);
                        byte[] seed = Hex.decode((String)buf.get("seed")); // seed for bike secure random
                        byte[] pk = Hex.decode((String)buf.get("pk"));     // public key
                        byte[] sk = Hex.decode((String)buf.get("sk"));     // private key
                        byte[] ct = Hex.decode((String)buf.get("ct"));     // ciphertext
                        byte[] ss = Hex.decode((String)buf.get("ss"));     // session key

                        NISTSecureRandom random = new NISTSecureRandom(seed, null);
                        BIKEParameters parameters = listParams[fileIndex];

                        BIKEKeyPairGenerator bikeKeyGen = new BIKEKeyPairGenerator();
                        BIKEKeyGenerationParameters genParam = new BIKEKeyGenerationParameters(random, parameters);

                        //
                        // Generate keys and test.
                        //

                        // KEM Keypair
                        bikeKeyGen.init(genParam);
                        AsymmetricCipherKeyPair pair = bikeKeyGen.generateKeyPair();

                        BIKEPublicKeyParameters generatedPk = (BIKEPublicKeyParameters)PublicKeyFactory.createKey(SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo((BIKEPublicKeyParameters)pair.getPublic()));
                        BIKEPrivateKeyParameters generatedSk = (BIKEPrivateKeyParameters)PrivateKeyFactory.createKey(PrivateKeyInfoFactory.createPrivateKeyInfo((BIKEPrivateKeyParameters)pair.getPrivate()));
                        assertTrue(name + " " + count + ": public key", Arrays.areEqual(pk, generatedPk.getEncoded()));
                        assertTrue(name + " " + count + ": secret key", Arrays.areEqual(sk, generatedSk.getEncoded()));

                        // KEM Encapsulation
                        BIKEKEMGenerator bikekemGenerator = new BIKEKEMGenerator(random);
                        SecretWithEncapsulation secretWithEnc = bikekemGenerator.generateEncapsulated(generatedPk);
                        byte[] secret = secretWithEnc.getSecret();
                        byte[] c = secretWithEnc.getEncapsulation();

                        assertTrue(name + " " + count + ": ciphertext", Arrays.areEqual(c, ct));
                        assertTrue(name + " " + count + ": kem_dec ss", Arrays.areEqual(secret, 0, secret.length, ss, 0, secret.length));

                        // KEM Decapsulation
                        BIKEKEMExtractor bikekemExtractor = new BIKEKEMExtractor(generatedSk);
                        byte[] dec_key = bikekemExtractor.extractSecret(c);

                        assertEquals(parameters.getSessionKeySize(), secret.length * 8);
                        assertTrue(name + " " + count + ": kem_dec ss", Arrays.areEqual(dec_key, 0, dec_key.length, ss, 0, dec_key.length));
                        assertTrue(name + " " + count + ": kem_dec key", Arrays.areEqual(dec_key, secret));
                    }
                    buf.clear();
                    continue;
                }
                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
            // System.out.println("Testing successful!");
        }
    }

    /**
     * A ciphertext that fails to decode must be handled by Fujisaki-Okamoto implicit rejection
     * (returning a pseudo-random shared secret), not by throwing. The decoder previously returned
     * null on a decoding failure, which surfaced as a NullPointerException out of decapsulation --
     * an uncaught crash on malformed input and a decryption-failure oracle.
     */
    public void testDecodingFailureImplicitRejection()
        throws Exception
    {
        byte[] seed = new byte[48];
        for (int i = 0; i != seed.length; i++)
        {
            seed[i] = (byte)i;
        }

        NISTSecureRandom random = new NISTSecureRandom(seed, null);
        BIKEParameters parameters = BIKEParameters.bike128;

        BIKEKeyPairGenerator bikeKeyGen = new BIKEKeyPairGenerator();
        bikeKeyGen.init(new BIKEKeyGenerationParameters(random, parameters));
        AsymmetricCipherKeyPair pair = bikeKeyGen.generateKeyPair();

        BIKEKEMGenerator bikekemGenerator = new BIKEKEMGenerator(random);
        SecretWithEncapsulation enc = bikekemGenerator.generateEncapsulated((BIKEPublicKeyParameters)pair.getPublic());
        byte[] goodSecret = enc.getSecret();

        // An all-ones ciphertext is a high-weight vector that will not decode, forcing the
        // decoder's failure path. extractSecret must still return a session key of the correct
        // length rather than throwing, and (with overwhelming probability) a different one.
        byte[] badCt = new byte[enc.getEncapsulation().length];
        Arrays.fill(badCt, (byte)0xFF);

        BIKEKEMExtractor extractor = new BIKEKEMExtractor((BIKEPrivateKeyParameters)pair.getPrivate());
        byte[] rejectKey = extractor.extractSecret(badCt);

        assertNotNull("decapsulation of a non-decodable ciphertext must not return null", rejectKey);
        assertEquals(parameters.getSessionKeySize() / 8, rejectKey.length);
        assertFalse("implicit rejection key must differ from the genuine shared secret",
            Arrays.areEqual(rejectKey, goodSecret));
    }
}
