package org.bouncycastle.pqc.crypto.test;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import junit.framework.TestCase;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusKeyGenerationParameters;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusKeyPairGenerator;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusParameters;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusPrivateKeyParameters;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusPublicKeyParameters;
import org.bouncycastle.pqc.legacy.sphincsplus.SPHINCSPlusSigner;
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;
import org.bouncycastle.util.test.FixedSecureRandom;

public class SphincsPlusTest
    extends TestCase
{
    public void testVectors()
        throws Exception
    {
        String files =
            " sha2-128f-simple.rsp sha2-192f-simple.rsp sha2-256f-simple.rsp shake-128f-simple.rsp" +
            " shake-192f-simple.rsp shake-256f-simple.rsp " +
            " sha2-128s-simple.rsp sha2-192s-simple.rsp" +
            " sha2-256s-simple.rsp shake-128s-simple.rsp shake-192s-simple.rsp shake-256s-simple.rsp" +
            " haraka-128f-simple.rsp  haraka-192s-simple.rsp" +
            " haraka-128s-simple.rsp  haraka-256f-simple.rsp" +
            " haraka-192f-simple.rsp  haraka-256s-simple.rsp";

        String[] fileList = splitOn(files, ' ');

        for (int i = 0; i != fileList.length; i++)
        {
            String name = fileList[i];

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

            String line = null;
            HashMap<String, String> buf = new HashMap<String, String>();
            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");
                        byte[] sk = Hex.decode((String)buf.get("sk"));
                        byte[] pk = Hex.decode((String)buf.get("pk"));
                        byte[] msg = Hex.decode((String)buf.get("msg"));
                        byte[] sigExpected = Hex.decode((String)buf.get("sm"));
                        byte[] oprR = Hex.decode((String)buf.get("optrand"));

                        if (sampler.skipTest(count))
                        {
                            continue;
                        }

                        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();
                        SecureRandom random = new FixedSecureRandom(sk);

                        SPHINCSPlusParameters parameters;

                        String[] nameParts = splitOn(name, '-');
                        boolean sha2 = nameParts[0].equals("sha2");
                        boolean shake = nameParts[0].equals("shake");
                        boolean haraka = nameParts[0].equals("haraka");
                        int size = Integer.parseInt(nameParts[1].substring(0, 3));
                        boolean fast = nameParts[1].endsWith("f");
                        boolean slow = nameParts[1].endsWith("s");
                        boolean simple = nameParts[2].equals("simple.rsp");
                        boolean robust = nameParts[2].equals("robust.rsp");

                        StringBuilder b = new StringBuilder();
                        if (sha2)
                        {
                            b.append("sha2");
                        }
                        else if (shake)
                        {
                            b.append("shake");
                        }
                        else if (haraka)
                        {
                            b.append("haraka");
                        }
                        else
                        {
                            throw new IllegalArgumentException("unknown digest");
                        }

                        b.append("_");
                        b.append(size);

                        if (fast)
                        {
                            b.append("f");
                        }
                        else if (slow)
                        {
                            b.append("s");
                        }
                        else
                        {
                            throw new IllegalArgumentException("unknown speed");
                        }

                        if (robust)
                        {
                            if (b.indexOf("haraka") < 0)
                            {
                                b.append("_robust");
                            }
                        }
                        else if (simple)
                        {
                            if (b.indexOf("haraka") >= 0)
                            {
                                b.append("_simple");
                            }
                        }
                        else
                        {
                            throw new IllegalArgumentException("unknown complexity");
                        }


                        parameters = (SPHINCSPlusParameters)SPHINCSPlusParameters.class.getField(b.toString()).get(null);

                        //
                        // Generate keys and test.
                        //
                        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, parameters));
                        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

                        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
                        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

                        // OID encode/decode round-trip runs for all variants including "-simple":
                        // the OIDs for the simple variants were wired correctly, so the former
                        // FIXME guard was stale. The getParameters() assertions below catch any
                        // future OID<->parameters mapping regression.
                        pubParams = (SPHINCSPlusPublicKeyParameters)PublicKeyFactory.createKey(
                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubParams));
                        privParams = (SPHINCSPlusPrivateKeyParameters)PrivateKeyFactory.createKey(
                            PrivateKeyInfoFactory.createPrivateKeyInfo(privParams));

                        assertTrue(name + " " + count + ": public key parameters", parameters == pubParams.getParameters());
                        assertTrue(name + " " + count + ": private key parameters", parameters == privParams.getParameters());

//                            System.err.println(Hex.toHexString(pubParams.getEncoded()));
//                            System.err.println(Hex.toHexString(Arrays.concatenate(pubParams.getParameters().getEncoded(), pk)));
                        assertTrue(name + " " + count + ": public key", Arrays.areEqual(pk, pubParams.getEncoded()));
                        assertTrue(name + " " + count + ": secret key", Arrays.areEqual(sk, privParams.getEncoded()));

                        //
                        // Signature test
                        //

                        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

                        signer.init(true, new ParametersWithRandom(privParams, new FixedSecureRandom(oprR)));

                        byte[] sigGenerated = signer.generateSignature(msg);
                        byte[] attachedSig = Arrays.concatenate(sigGenerated, msg);


                        signer.init(false, pubParams);

                        assertTrue(name + " " + count + ": signature verify", signer.verifySignature(msg, Arrays.copyOfRange(sigExpected, 0, sigGenerated.length)));

//                            System.err.println(Hex.toHexString(sigExpected));
//                            System.err.println(Hex.toHexString(attachedSig));
                        assertTrue(name + " " + count + ": signature gen match", Arrays.areEqual(sigExpected, attachedSig));

                    }
                    buf.clear();

                    continue;
                }

                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
            src.close();
        }
        //System.err.println(System.currentTimeMillis() - startTime);
    }

    public void testBasicKeyGeneration()
        throws IOException
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4792F267AAFA3F87CA60D01CB54F29202A3E784CCB7EBCDCFD45542B7F6AF778742E0F4479175084AA488B3B74340678AAD111491E7E52F6F1D726DAF2A4E75CAFB60D034B6E912B26BE68464B0095D60D"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.shake_256f_robust));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        assertTrue(Arrays.areEqual(Hex.decode("3e784ccb7ebcdcfd45542b7f6af778742e0f4479175084aa488b3b74340678aa6ba9430051e61cb676e8449087b938a79575b3a16736ce68a3655a28001155f5"), pubParams.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e4792f267aafa3f87ca60d01cb54f29202a3e784ccb7ebcdcfd45542b7f6af778742e0f4479175084aa488b3b74340678aa6ba9430051e61cb676e8449087b938a79575b3a16736ce68a3655a28001155f5"), privParams.getEncoded()));

        SubjectPublicKeyInfo pubInfo = SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubParams);
        PrivateKeyInfo privInfo = PrivateKeyInfoFactory.createPrivateKeyInfo(privParams);

        pubParams = (SPHINCSPlusPublicKeyParameters)PublicKeyFactory.createKey(pubInfo.getEncoded());
        privParams = (SPHINCSPlusPrivateKeyParameters)PrivateKeyFactory.createKey(privInfo.getEncoded());

        assertTrue(Arrays.areEqual(Hex.decode("3e784ccb7ebcdcfd45542b7f6af778742e0f4479175084aa488b3b74340678aa6ba9430051e61cb676e8449087b938a79575b3a16736ce68a3655a28001155f5"), pubParams.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e4792f267aafa3f87ca60d01cb54f29202a3e784ccb7ebcdcfd45542b7f6af778742e0f4479175084aa488b3b74340678aa6ba9430051e61cb676e8449087b938a79575b3a16736ce68a3655a28001155f5"), privParams.getEncoded()));
    }

    public void testBasicKeyImportSimpleSign()
    {
        SPHINCSPlusPublicKeyParameters pubParams = new SPHINCSPlusPublicKeyParameters(SPHINCSPlusParameters.sha2_128f_robust, Hex.decode("b505d7cfad1b497499323c8686325e473985e5a31e5b9a0457916c84320c2ea8"));
        SPHINCSPlusPrivateKeyParameters privParams = new SPHINCSPlusPrivateKeyParameters(SPHINCSPlusParameters.sha2_128f_robust, Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e473985e5a31e5b9a0457916c84320c2ea8"));

        assertTrue(Arrays.areEqual(Hex.decode("b505d7cfad1b497499323c8686325e473985e5a31e5b9a0457916c84320c2ea8"), pubParams.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e473985e5a31e5b9a0457916c84320c2ea8"), privParams.getEncoded()));

        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

        signer.init(true, new ParametersWithRandom(privParams, new FixedSecureRandom(Hex.decode("33b3c07507e4201748494d832b6ee2a6"))));

        byte[] sig = signer.generateSignature(msg);
        byte[] attachedSig = Arrays.concatenate(sig, msg);

        byte[] expected = Hex.decode("[encoded data, 34242 characters]");
        assertTrue(Arrays.areEqual(expected, attachedSig));

        signer.init(false, pubParams);

        assertTrue(signer.verifySignature(msg, sig));
    }

    public void testBasicSignature()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.sha2_128f_robust));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

        signer.init(true, new ParametersWithRandom(privParams, new FixedSecureRandom(Hex.decode("33b3c07507e4201748494d832b6ee2a6"))));

        byte[] sig = signer.generateSignature(msg);
        byte[] attachedSig = Arrays.concatenate(sig, msg);

        byte[] expected = Hex.decode("[encoded data, 12784 characters]" +
            "[encoded data, 21458 characters]");
        assertTrue(Arrays.areEqual(expected, attachedSig));

        signer.init(false, pubParams);

        assertTrue(signer.verifySignature(msg, sig));
    }

    public void testDeterministicSignature()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.sha2_128f_robust));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

        signer.init(true, privParams);

        byte[] sig = signer.generateSignature(msg);
        byte[] attachedSig = Arrays.concatenate(sig, msg);

        byte[] expected = Hex.decode("[encoded data, 12133 characters]" +
            "[encoded data, 22109 characters]");
        assertTrue(Arrays.areEqual(expected, attachedSig));

        signer.init(false, pubParams);

        assertTrue(signer.verifySignature(msg, sig));
    }

    public void testBasicKeyGenerationShake256128fSimple()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();

        SecureRandom random = new FixedSecureRandom(Hex.decode(
            "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4766BA69D8560A9F84846AD8B765390C84"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.shake_128f));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        assertTrue(Arrays.areEqual(Hex.decode("b505d7cfad1b497499323c8686325e47afbc007ba1e2b4a138f03aa9a6195ac8"), pubParams.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e47afbc007ba1e2b4a138f03aa9a6195ac8"), privParams.getEncoded()));
    }

    public void testBasicKeyGenerationShake256128fSimpleSign()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();

        SecureRandom random = new FixedSecureRandom(Hex.decode(
            "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4766BA69D8560A9F84846AD8B765390C84"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.shake_128f));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

        signer.init(true, new ParametersWithRandom(privParams, new FixedSecureRandom(Hex.decode("33b3c07507e4201748494d832b6ee2a6"))));

        byte[] sig = signer.generateSignature(msg);
        byte[] attachedSig = Arrays.concatenate(sig, msg);

        byte[] expected = Hex.decode("[encoded data, 34242 characters]");
        assertTrue(Arrays.areEqual(expected, attachedSig));

        signer.init(false, pubParams);

        assertTrue(signer.verifySignature(msg, sig));
    }


    public void testBasicKeyGenerationShake256128fRobust()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();

        SecureRandom random = new FixedSecureRandom(Hex.decode(
            "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E47354D75735F16E03DEC94D1F5B00C213D"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.shake_128f_robust));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        assertTrue(Arrays.areEqual(Hex.decode("b505d7cfad1b497499323c8686325e4714be46e5b92237d09a0ea8a0404033a6"), pubParams.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode("7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2db505d7cfad1b497499323c8686325e4714be46e5b92237d09a0ea8a0404033a6"), privParams.getEncoded()));
    }

    public void testBasicKeyGenerationShake256128fRobustSign()
    {
        SPHINCSPlusKeyPairGenerator kpGen = new SPHINCSPlusKeyPairGenerator();

        SecureRandom random = new FixedSecureRandom(Hex.decode(
            "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E47354D75735F16E03DEC94D1F5B00C213D"));

        kpGen.init(new SPHINCSPlusKeyGenerationParameters(random, SPHINCSPlusParameters.shake_128f_robust));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();

        SPHINCSPlusPublicKeyParameters pubParams = (SPHINCSPlusPublicKeyParameters)kp.getPublic();
        SPHINCSPlusPrivateKeyParameters privParams = (SPHINCSPlusPrivateKeyParameters)kp.getPrivate();

        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        SPHINCSPlusSigner signer = new SPHINCSPlusSigner();

        signer.init(true, new ParametersWithRandom(privParams, new FixedSecureRandom(Hex.decode("33b3c07507e4201748494d832b6ee2a6"))));

        byte[] sig = signer.generateSignature(msg);
        byte[] attachedSig = Arrays.concatenate(sig, msg);

        byte[] expected = Hex.decode("[encoded data, 34242 characters]");
        assertTrue(Arrays.areEqual(expected, attachedSig));

        signer.init(false, pubParams);

        assertTrue(signer.verifySignature(msg, sig));
    }

    private static String[] splitOn(String input, char c)
    {
        String s = input.trim();
        List l = new ArrayList();

        int idx = s.indexOf(c);
        while (idx > 0)
        {
            l.add(s.substring(0, idx));
            s = s.substring(idx + 1).trim();
            idx = s.indexOf(c);
        }

        if (s.length() > 0)
        {
            l.add(s);
        }

        return (String[])l.toArray(new String[0]);
    }
}
