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.HashMap;

import junit.framework.TestCase;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumSigner;
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.Strings;
import org.bouncycastle.util.encoders.Hex;

public class CrystalsDilithiumTest
    extends TestCase
{
    public void testKeyGen() throws IOException
    {
        String[] files = new String[]{
                "keyGen_ML-DSA-44.txt",
                "keyGen_ML-DSA-65.txt",
                "keyGen_ML-DSA-87.txt",
        };

        DilithiumParameters[] params = new DilithiumParameters[]{
                DilithiumParameters.dilithium2,
                DilithiumParameters.dilithium3,
                DilithiumParameters.dilithium5,
        };

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

            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/dilithium/acvp", 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)
                    {
                        byte[] seed = Hex.decode((String)buf.get("seed"));
                        byte[] pk = Hex.decode((String)buf.get("pk"));
                        byte[] sk = Hex.decode((String)buf.get("sk"));

                        DilithiumParameters parameters = params[fileIndex];

                        DilithiumKeyPairGenerator kpGen = new DilithiumKeyPairGenerator();
                        DilithiumKeyGenerationParameters genParam = new DilithiumKeyGenerationParameters(new SecureRandom(), parameters);
                        //
                        // Generate keys and test.
                        //
                        kpGen.init(genParam);
                        AsymmetricCipherKeyPair kp = kpGen.internalGenerateKeyPair(seed);

                        DilithiumPublicKeyParameters pubParams = (DilithiumPublicKeyParameters)PublicKeyFactory.createKey(
                                SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo((DilithiumPublicKeyParameters)kp.getPublic()));
                        DilithiumPrivateKeyParameters privParams = (DilithiumPrivateKeyParameters)PrivateKeyFactory.createKey(
                                PrivateKeyInfoFactory.createPrivateKeyInfo((DilithiumPrivateKeyParameters)kp.getPrivate()));

                        assertTrue(name + ": public key", Arrays.areEqual(pk, pubParams.getEncoded()));
                        assertTrue(name + ": secret key", Arrays.areEqual(sk, privParams.getEncoded()));

                    }
                    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!");
        }
    }
    public void testSigGen() throws IOException
    {
        String[] files = new String[]{
                "sigGen_ML-DSA-44.txt",
                "sigGen_ML-DSA-65.txt",
                "sigGen_ML-DSA-87.txt",
        };

        DilithiumParameters[] params = new DilithiumParameters[]{
                DilithiumParameters.dilithium2,
                DilithiumParameters.dilithium3,
                DilithiumParameters.dilithium5,
        };

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

            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/dilithium/acvp", 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)
                    {
                        boolean deterministic = !buf.containsKey("rnd");
                        byte[] sk = Hex.decode((String)buf.get("sk"));
                        byte[] message = Hex.decode((String)buf.get("message"));
                        byte[] signature = Hex.decode((String)buf.get("signature"));
                        byte[] rnd = null;
                        if(!deterministic)
                        {
                            rnd = Hex.decode((String)buf.get("rnd"));
                        }

                        DilithiumParameters parameters = params[fileIndex];

                        DilithiumPrivateKeyParameters privParams = new DilithiumPrivateKeyParameters(parameters, sk, null);

                        // sign
                        DilithiumSigner signer = new DilithiumSigner();

                        signer.init(true, privParams);
                        byte[] sigGenerated;
                        if(!deterministic)
                        {
                            sigGenerated = signer.internalGenerateSignature(message, rnd);
                        }
                        else
                        {
                            sigGenerated = signer.generateSignature(message);
                        }
                        assertTrue(Arrays.areEqual(sigGenerated, signature));
                    }
                    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!");
        }
    }
    public void testSigVer() throws IOException
    {
        String[] files = new String[]{
                "sigVer_ML-DSA-44.txt",
                "sigVer_ML-DSA-65.txt",
                "sigVer_ML-DSA-87.txt",
        };

        DilithiumParameters[] params = new DilithiumParameters[]{
                DilithiumParameters.dilithium2,
                DilithiumParameters.dilithium3,
                DilithiumParameters.dilithium5,
        };

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

            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/dilithium/acvp", 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)
                    {
                        boolean testPassed = TestUtils.parseBoolean((String)buf.get("testPassed"));
                        String reason = (String)buf.get("reason");
                        byte[] pk = Hex.decode((String)buf.get("pk"));
                        byte[] sk = Hex.decode((String)buf.get("sk"));
                        byte[] message = Hex.decode((String)buf.get("message"));
                        byte[] signature = Hex.decode((String)buf.get("signature"));

                        DilithiumParameters parameters = params[fileIndex];

                        DilithiumPublicKeyParameters pubParams = new DilithiumPublicKeyParameters(parameters, pk);
                        DilithiumPrivateKeyParameters privParams = new DilithiumPrivateKeyParameters(parameters, sk, null);


                        DilithiumSigner verifier = new DilithiumSigner();
                        verifier.init(false, pubParams);

                        boolean ver = verifier.verifySignature(message, signature);
                        assertEquals("expected " + testPassed + " " + reason, ver, testPassed);
                    }
                    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!");
        }
    }

    public void testDilithium()
    {
        String seed = "70CEFB9AED5B68E018B079DA8284B9D5CAD5499ED9C265FF73588005D85C225C";
        String pk = "D2FD03F3A1B7F635AF9F34D580A98F524C735BD5BA2355DC6E035BD21765580CBB111923F194A7CC8A7BB2EBC5C0E71AA637CC800E6103B850A539B2A39E1B6D713E5DB8314C9AE1F8BF8A38F06AFB9D73B161B0FFE3A4891706AE26D54FFB496DF8DC0F1983509500C9ABBD28E59B3FCDABBDADABD45EC31499378BDE849E7C1F19B7044D67E05106D7136D95380D5605D4465D877557065DF0A75D3C28542F40FEED42EC7E280637B083D988BCA5F6394E02396C4676184FB63318DAFAF5BBDDE00E308FE84019C2340A3F3E1C0865624970711283356AE14BD6B94D1C9AE188DE1A8A2CA824A8EAE2FE6AFB38D83A2D99996AB21FE3E84C0BE6B6DA08879B677374FA7C691B13D40FA9D4CC26B2288D5A8C9A43724381004D61B0D57FF400314C8E30EE796AF10F7EE21BF13D08180465ABC72EDDB080C6A07184E3EEDC47C19AA7F09D1F3309E183A2BD9B0573DDE474A81BA4F78D0C523D0C04F90060FD571A35C037E079C5E210D7390DF568F2E2F03CE44420C82F3FE69EB9B48EE90962D6B0F24440648F71EDB241EE6566FC1A64CABF66BE6FECBCB1387C82A7BC202D9E367998E2A291AF0CD1570677FE8D63A3285A2EA6EB29AF9DC1AEC1C36C4706B12BAA20839692F286A6E0321468F7479345C4D52FBDB2F06725B554B89E2492612681ACEBC6C7BADA9225818DBC35D64C22C48BFF80A730D0716DFAC99DFD5B8992611D0C93EE90BDB260022AFE25D913E06EFFB59CB1F8A60CBFA5AB2F459A16F467E989525E0A37EBE56E833FDE55DB9D1530ADCF45846DF281E47CAA1E0A27EFDE2107D354CEA0F6A454692F04CD838EBDD46E191E5D9C11839A2C3F488A4FC7CD265A7B5D32B08CBDBFAB9D2CCD76222C8EE37DDCBD2AA063ED861473A6454CAEA377850B1A2B9DDBBCB374FAB5B12F351C8E5888872E5CD1F60A4FAE1FF837D192C22BEB41EE6FA392FCDF4550FF46B5CE906D017EF3077DF132300D8BBFA9BB03C75E79E2F04C284AD06A44399649C3E2A2A8D1EFE9B7A4E0C271047AB75908BFF7DF9E30ECA547745BAE23A86FF9A8B58C2538B88B866401076902DC5F0BD761687B49EAFE36D350CBEDFDD36C121CF23786BFCF7E47076496EAB6BBDA774049C2EBABE2DE99C4C24F2DB73684015B373977496760CF9AC23D8B623133DB2DE10D73FA6AD1C6DAC8434F28C6E251CE7293CFF3F3B61EFCB5A435123670F29846A13DF3EE712604461F1BAB8F4EBC836DE058978AE734396A98081B35CC98188A86949C99270D4709854C5B35B17F48A373134C814CC8A0F3E2FA807F2A918530907864778282D75E03A41B2504EED816A417A3AC6BA16080C39B7310192002A728F7F20395009A9E16767CE1971F5DE7D229A50613369E4382045A8E81901F4DBA8102F3D413FE35B326A874F233B719A7137600D35D33AEB6B7259624083AA968730C8F78292AD28F14EEABE660835984FE69EF23DEC8C327C0EB0B882D587E1EC433DA85C9FD1E0A34994DEA240C854452D18C30F496E49EC904B602E0F5062EDCDA03280A53B4313574CC2C0D5471BC9613BDFD6641F5BD127BAB5B5EB3D499A33114048220E819F8EE12CA922C8F17D9C9F51AD5BD6883B10E6AA2483BA49DC547DA7686151344F4E9099B38E430B5226B059832CF03DB48FB02DBA4E61593DC4576360491890E53EC0E6AC73CF32B25D823B38456E286505A541E5AEEE96B1914F5F76687CE2B0160227ABED77993594BCD831366206D75714082F1C46F1F4439AC81A57AF31C81C555307A070FFA94E0479B784BBD88A60CD4C7CFD94E6AFE02F6B21F72AF0DCD6609D40C965C14E5F2389183E53DE930F7DE1D44215CF49144844E8B87F78A7F132AEFE22BE80B4E3A05EE3A68CCF609EF44047402E4493046E6F9C767FF8A75E28B3CE077FDE7E7EED313B5BF7E460127CA8182E9BC794C0DFA730FB920080575A751B5CAEC85A109B4422BA266743F0D032BDA8F1CA6248CDB917530DF1302A5F8C18DC642D52478C98C12A3F16EF2B62B4F59EA1BB58DE7B65B3C7153CE6DA5E4950746F80E087A0E3586D097791BF36DEF865D68591D39D0903773EEA962147F34704138B54DF7924CDD8C333DB5E1A409CCB2B34E2C3C8C7FDD3FD8D012CBF382AAA85E83A12F235A2D147D035B7B28B34B6F57949F322482A7D4D3B15045C420D5ADDC7F0E69B4DC1CBA58B01D872480B06A260D827D891B13C4C5CA50C748DE3C771BE61E9AA170165CB01F4BF5DA27A7791D3AD3F6267B4CB4E61B28FA1708418D932DFC4161880C5D3B17A9663A9061FA8F1804315850FE4E7306C882B38227E867F80872CDC1944D472615EA4900EF7D270B881D4130F56C5CC980D92A47ADA6657EB6F37A385D2D8CC993E1442EB05281853636991E34AADC68954D04E7ADEF76BF880F059B0CBB55D915A4B123E2F1339A073CBFBC409BEFF6400AE096D5AE18EC42CFFAD5B4980FA35BF03413ADB5D7E6876AC355D1C9ED70CA2B973954D12B3CDD76AC6835DB96003ED8C4E288B71FD77DBAA7635720E12AE0A317DE808C664E317F55275791F3245CA4FE5D4D41077FC150A6E403D5A208E46EADBE8F2CFB8AF472F4A0CEAC015219478E6B86C958CF86525B7485C1734C7EF00E90683FFF5DBD0A7D413A855021026A1B32013A4616CBCD3700ACBC705BE3EFBA625C69A025267BCE9D135E3F5B5CC8C43956407E84B6663103E29C242035551AE797F56C6374BE0C798C0CF398F1ED";
        String sk = "[encoded data, 8064 characters]";

        SecureRandom random = new SecureRandom();

        DilithiumKeyPairGenerator keyGen = new DilithiumKeyPairGenerator();

        keyGen.init(new DilithiumKeyGenerationParameters(random, DilithiumParameters.dilithium3));

        AsymmetricCipherKeyPair keyPair = keyGen.internalGenerateKeyPair(Hex.decode(seed));

        assertTrue(Arrays.areEqual(Hex.decode(sk), ((DilithiumPrivateKeyParameters)keyPair.getPrivate()).getEncoded()));
        DilithiumPublicKeyParameters dPub = (DilithiumPublicKeyParameters)keyPair.getPublic();
        assertTrue(Arrays.areEqual(Hex.decode(pk), dPub.getEncoded()));
    }

    public void testRNG()
    {
        String temp = "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1";
        byte[] seed = Hex.decode(temp);

        NISTSecureRandom r = new NISTSecureRandom(seed, null);

        String testBytesString = "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2D";
        byte[] testBytes = Hex.decode(testBytesString);

        byte[] randBytes = new byte[testBytes.length];
        r.nextBytes(randBytes);

        assertTrue(Arrays.areEqual(randBytes, testBytes));
    }

    public void testVectors()
        throws Exception
    {
        //Disabled OLD KATs

//        String[] files = new String[]{
//            "PQCsignKAT_Dilithium2.rsp",
//            "PQCsignKAT_Dilithium3.rsp",
//            "PQCsignKAT_Dilithium5.rsp",
//        };
//
//        DilithiumParameters[] parameters = new DilithiumParameters[]{
//            DilithiumParameters.dilithium2,
//            DilithiumParameters.dilithium3,
//            DilithiumParameters.dilithium5,
//        };
//
//        TestSampler sampler = new TestSampler();
//        for (int fileindex = 0; fileindex < files.length; fileindex++)
//        {
//            String name = files[fileindex];
//            // System.out.println("testing: " + name);
//            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/dilithium", name);
//            BufferedReader bin = new BufferedReader(new InputStreamReader(src));
//            String line = null;
//            HashMap<String, String> buf = new HashMap<String, String>();
//            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 Dilithium secure random
//                        byte[] pk = Hex.decode((String)buf.get("pk"));     // public key
//                        byte[] sk = Hex.decode((String)buf.get("sk"));     // private key
//                        byte[] sm = Hex.decode((String)buf.get("sm"));     // signed message
//                        int sm_len = Integer.parseInt((String)buf.get("smlen"));
//                        byte[] msg = Hex.decode((String)buf.get("msg")); // message
//                        int m_len = Integer.parseInt((String)buf.get("mlen"));
//
//                        NISTSecureRandom random = new NISTSecureRandom(seed, null);
//
//                        // keygen
//                        DilithiumKeyGenerationParameters kparam = new DilithiumKeyGenerationParameters(random, parameters[fileindex]);
//                        DilithiumKeyPairGenerator kpg = new DilithiumKeyPairGenerator();
//                        kpg.init(kparam);
//
//                        AsymmetricCipherKeyPair kp = kpg.generateKeyPair();
//
//                        DilithiumPublicKeyParameters pubParams = (DilithiumPublicKeyParameters)PublicKeyFactory.createKey(
//                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo((DilithiumPublicKeyParameters)kp.getPublic()));
//                        DilithiumPrivateKeyParameters privParams = (DilithiumPrivateKeyParameters)PrivateKeyFactory.createKey(
//                            PrivateKeyInfoFactory.createPrivateKeyInfo((DilithiumPrivateKeyParameters)kp.getPrivate()));
//
//                        assertTrue(name + " " + count + " public key", Arrays.areEqual(pk, pubParams.getEncoded()));
//                        assertTrue(name + " " + count + " secret key", Arrays.areEqual(sk, privParams.getEncoded()));
//
//                        // sign
//                        DilithiumSigner signer = new DilithiumSigner();
//                        DilithiumPrivateKeyParameters skparam = (DilithiumPrivateKeyParameters)kp.getPrivate();
//
//                        signer.init(true, skparam);
//
//                        byte[] sigGenerated = signer.generateSignature(msg);
//                        byte[] attachedSig = Arrays.concatenate(sigGenerated, msg);
//
//                        // verify
//                        DilithiumSigner verifier = new DilithiumSigner();
//                        DilithiumPublicKeyParameters pkparam = pubParams;
//                        verifier.init(false, pkparam);
//
//                        boolean vrfyrespass = verifier.verifySignature(msg, sigGenerated);
//                        sigGenerated[3]++; // changing the signature by 1 byte should cause it to fail
//                        boolean vrfyresfail = verifier.verifySignature(msg, sigGenerated);
//
//                        // print results
//                            /*
//                            // System.out.println("--Keygen");
//                            boolean kgenpass = true;
//                            if (!Arrays.areEqual(respk, pk)) {
//                                // System.out.println("  == Keygen: pk do not match");
//                                kgenpass = false;
//                            }
//                            if (!Arrays.areEqual(ressk, sk)) {
//                                // System.out.println("  == Keygen: sk do not match");
//                                kgenpass = false;
//                            }
//                            if (kgenpass) {
//                                // System.out.println("  ++ Keygen pass");
//                            } else {
//                                // System.out.println("  == Keygen failed");
//                                return;
//                            }
//
//                            // System.out.println("--Sign");
//                            boolean spass = true;
//                            if (!Arrays.areEqual(ressm, sm)) {
//                                // System.out.println("  == Sign: signature do not match");
//                                spass = false;
//                            }
//                            if (spass) {
//                                // System.out.println("  ++ Sign pass");
//                            } else {
//                                // System.out.println("  == Sign failed");
//                                return;
//                            }
//
//                            // System.out.println("--Verify");
//                            if (vrfyrespass && !vrfyresfail) {
//                                // System.out.println("  ++ Verify pass");
//                            } else {
//                                // System.out.println("  == Verify failed");
//                                return;
//                            }
//                             */
//                        // AssertTrue
//
//                        //sign
//                        // // System.out.println("attached Sig = ");
//                        // Helper.printByteArray(attachedSig);
//                        // // System.out.println("sm = ");
//                        // Helper.printByteArray(sm);
//
//                        assertTrue(name + " " + count + " signature", Arrays.areEqual(attachedSig, sm));
//
//                        //verify
//                        assertTrue(name + " " + count + " verify failed when should pass", vrfyrespass);
//                        assertFalse(name + " " + count + " verify passed when should fail", vrfyresfail);
//
//                    }
//                    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!");
//        }
    }

    public void testDilithiumRandom()
    {
        byte[] msg = Strings.toByteArray("Hello World!");
        DilithiumKeyPairGenerator keyGen = new DilithiumKeyPairGenerator();

        SecureRandom random = new SecureRandom();

        keyGen.init(new DilithiumKeyGenerationParameters(random, DilithiumParameters.dilithium3));

        for (int i = 0; i != 1000; i++)
        {
            AsymmetricCipherKeyPair keyPair = keyGen.generateKeyPair();

            // sign
            DilithiumSigner signer = new DilithiumSigner();
            DilithiumPrivateKeyParameters skparam = (DilithiumPrivateKeyParameters)keyPair.getPrivate();
            ParametersWithRandom skwrand = new ParametersWithRandom(skparam, random);
            signer.init(true, skwrand);

            byte[] sigGenerated = signer.generateSignature(msg);

            // verify
            DilithiumSigner verifier = new DilithiumSigner();
            DilithiumPublicKeyParameters pkparam = (DilithiumPublicKeyParameters)keyPair.getPublic();
            verifier.init(false, pkparam);

            assertTrue("count = " + i, verifier.verifySignature(msg, sigGenerated));
        }
    }
}
