package org.bouncycastle.pqc.crypto.test;

import java.io.BufferedReader;
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
{
    /*
count = 0
seed = 061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1
mlen = 33
msg = D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8
pk = 1C0EE1111B08003F28E65E8B3BDEB037CF8F221DFCDAF5950EDB38D506D85BEFD9FDE3A496F75819F0A20D0441DC7830B4AA1CB8ECFC91BA0EEC3AFB6744E477B4E6EC3FDAE75048FFEBAABEA8E822117D5787F79070EA88287CE3CD5011FD8D93AB7E8B51F26116BF9B6D21C03F88BFEC488876F4D075A142D4E784D734407511F992069353F1DB67ACF73034A468A118588062111D320E00BCFF6DC63573FCED1E96AAEBA6452E3C7ACD19181F9B814BA19D39B4BAB5496DC055426E7EA461AF55D5B9FE97F9DF7E253203C1F9E152E96D75F9D9A84F5C263EC8C250440ADC986F4E36414C703B3E05426B28B7065950DA6D0E0B2C60AC3672DB6F3C78447DB7C20915770EA6FCE81DAB5339C1D5AF82A5D3324099DF56516A07DB7C0FC64383805C65F2B02FBCFCE63E93C4BF09409F9F0F77E73DA3B0019F2057E4CD7CFF0E5745EF18C3FD766E01747A64D415FC9789ABFA62284E11C7FF05D0548D973F679559A6A3AAD77ED5132D0150C014C3EC3A395F017E7ACFE3EABFCA44910CA06FF33542ECCE6241974742357D37F5C284BF0FE1A74B50C073551372133AF2DD41E21BAFC9C590EE6EBC4ACE731EF566156CA03755DC493C137028AF3B3DE5B00BD6CB3D9A87D0151F887C6768BC6CA02A94FB2086551A0F89BA26154E9D4506AD9FAF39F5723E234E06CFDED69D4EE4146B73E5DC1E4152A2A3159D73DBC833D3D417CD5CF7FB3DC7745CEED4DC0F5B1C6D6B69C1764157EA43DF9DBB442EFA39D1D0162E87C2D30C5012FD16D869C8A1FCBB45EDCC8E1813B2B190A961F9FC86591D3ABC5388AF678FF03DA78B7CC0F6185721C0DF33CC906435225DF2611002DF120E83566532292DEA3D8ACD109A0DFFAB3B0B43012796DB5B50683FB4C2D250DAB76AAE35A48E8C8D4A5CC154759745F0A1230F6CA9DD9C99E2F80EDC83304CE01E98F6C9489529A822F90033C228315EB2FCC8DBA382ED4301E07607A5B076C725F124994F18A997D2C5BBF9A324605265108ACBF4610FA1C3374408850A0864E2B61017EBEC1FBAB89DE3AB1B93CE4918B9E2C9E3FE456758062A9F882B283318271F4B9552FCF32624A9FDAA44C65C60E2B3648BEF1F17D0B7C74869EE0B53C4A62A24845DCEA5BCBF93B92E4C26648584E33479282E6C8B1D8FE21181BD9CF75F8A961724D4C4309779F1F1B775D254F70BD1769CC7C0EDD2A95FE5C9D84B16F7C54D85CCE4C8A182810809ED81E97D074884EEDF401CCACDAEAD82C14D06B68AEA6CE14B861B0CFD16090CBBF469C5E084314C0D8D3960EA06A3426D8B3FE762E00D09BDA374F3AE2CBEDE2838FF89D81DEB3013090E44199AED604963EAF919914CE04F207AC82CD4351FEF7B2D94393066FE4D44E3CC5952E75EB6F3714058915DE0EE184D8C55300F576A8B82A863E81AF33417BD4CFC94E7A61263B39F01F6E2E70748B6E5E59CF6CA01B0028C93BBBCEBC548F987F10755BF33CA585CB41CF578DF5FFE37924E3C2C072ED1DAC9162176972971E79B62FB208F1A73BF0361E2993DCCCD3110C34D839D18DD43A5E8F0D941E99ADCF441405F32107671B2D8B2244F7BA92DCED587A210FE8FF43C616ACB5E766E6AF2CEB03599BA3DE376EB5735EF16143953D1FDDB7E9F2874B0D6083DD7EC4386AE003F51CCF2D21EF6059163C5152174423F57119D0FCE627D763D81C10AA1329F74C8D445437BA6718A33DB6E79375172B2AE3591821978D520824E2D2FF898B7F4C867FF462722BC07EADAD389A910B6F65429DA129735FE049E3ECB3889F6047CF2BD2A88D50A651B3235D2480E1DA5A35247FA76C831736399D37E8D033C1D051C9B6A99AB80B1313FA24C5C59766E6C51A38FE9F1186A767EEBD0D88001AE0246CD4EBE2C979DE82C30BBDB98B4744F11F9E639EDDD8C194D7911201A8FA745991B4D8A5709B62A21B63B9762913D36CE995C2D6B79151E8D83838CD1F38840A9417255DD166B7A3584499003FB625611404C95B960DF0DB1BCF1574B0965DBD834EE148117D5E05A7CC7CC1A865618A2BE4854DB8935CDA1E68BD8D09E72F0AC9053C882C4ABA4004A614D10505300B6176CA1F324E22E7824299F9C40755B71D82B679547F06AD48BE66D68072C9390233C933F80A14F8D4A6B0B4E1970E1ACC1BEA7F5D3BE224448F857BAB68AEFA6D8CB819B64294A12997916CDBF56E9A8D002DD065F12C61823F4FC214508232E431F0B6898475BB5DD0D7D528E840C22809AF7E15363724A613ACCFBE2B37438C159CE14CB0C98BFD499C08DAC0CF45D821CC2FA47319B6FB4CED7E5985EC8274DE09071D3C10DA5BF9E522B01CE91D66B91795D3D22C00483454275DD2BBDD7C2DCC4A167E5D7FCDBB9F6208CD4C9A485FAAEB809A7711DAC2865CED4306474B22B4448F85DF33417F3FACE1C05D42703ED313042A05DE0362740130188ECB445BB255DC76EE8443F733117F8351F17603175554FEB00B7FF54D80786F305CDE18CD5EC56EC0962A3E04482DCE3622D040D24C40F2E8A14A447659D6C561F2FFEE68F8D3DE511B23E8B172A01A3EDA4D3780E74C677244330E9AEFF019FE07BE3D33F322F9CE2214B9D9CFF99D05A59E47551432AE76F4CD4F8DD51520FFE811B4B93CD6219C81B63B1D627785C2A0FC22E3AEA86CEEE1F7FBC4EFCB46DDFBCD88A02F3B4E67C5FF2E8DC68BF16C74699BBB628902F72C3DEBC8BF5DF706D47A605A107DAA0014139CE40F0D46D8D6DC7
sk = [encoded data, 8064 characters]
smlen = 3326
sm = [encoded data, 6652 characters]
    */
    public void testDilithium()
    {
        String seed = "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1";
        String pk = "1C0EE1111B08003F28E65E8B3BDEB037CF8F221DFCDAF5950EDB38D506D85BEFD9FDE3A496F75819F0A20D0441DC7830B4AA1CB8ECFC91BA0EEC3AFB6744E477B4E6EC3FDAE75048FFEBAABEA8E822117D5787F79070EA88287CE3CD5011FD8D93AB7E8B51F26116BF9B6D21C03F88BFEC488876F4D075A142D4E784D734407511F992069353F1DB67ACF73034A468A118588062111D320E00BCFF6DC63573FCED1E96AAEBA6452E3C7ACD19181F9B814BA19D39B4BAB5496DC055426E7EA461AF55D5B9FE97F9DF7E253203C1F9E152E96D75F9D9A84F5C263EC8C250440ADC986F4E36414C703B3E05426B28B7065950DA6D0E0B2C60AC3672DB6F3C78447DB7C20915770EA6FCE81DAB5339C1D5AF82A5D3324099DF56516A07DB7C0FC64383805C65F2B02FBCFCE63E93C4BF09409F9F0F77E73DA3B0019F2057E4CD7CFF0E5745EF18C3FD766E01747A64D415FC9789ABFA62284E11C7FF05D0548D973F679559A6A3AAD77ED5132D0150C014C3EC3A395F017E7ACFE3EABFCA44910CA06FF33542ECCE6241974742357D37F5C284BF0FE1A74B50C073551372133AF2DD41E21BAFC9C590EE6EBC4ACE731EF566156CA03755DC493C137028AF3B3DE5B00BD6CB3D9A87D0151F887C6768BC6CA02A94FB2086551A0F89BA26154E9D4506AD9FAF39F5723E234E06CFDED69D4EE4146B73E5DC1E4152A2A3159D73DBC833D3D417CD5CF7FB3DC7745CEED4DC0F5B1C6D6B69C1764157EA43DF9DBB442EFA39D1D0162E87C2D30C5012FD16D869C8A1FCBB45EDCC8E1813B2B190A961F9FC86591D3ABC5388AF678FF03DA78B7CC0F6185721C0DF33CC906435225DF2611002DF120E83566532292DEA3D8ACD109A0DFFAB3B0B43012796DB5B50683FB4C2D250DAB76AAE35A48E8C8D4A5CC154759745F0A1230F6CA9DD9C99E2F80EDC83304CE01E98F6C9489529A822F90033C228315EB2FCC8DBA382ED4301E07607A5B076C725F124994F18A997D2C5BBF9A324605265108ACBF4610FA1C3374408850A0864E2B61017EBEC1FBAB89DE3AB1B93CE4918B9E2C9E3FE456758062A9F882B283318271F4B9552FCF32624A9FDAA44C65C60E2B3648BEF1F17D0B7C74869EE0B53C4A62A24845DCEA5BCBF93B92E4C26648584E33479282E6C8B1D8FE21181BD9CF75F8A961724D4C4309779F1F1B775D254F70BD1769CC7C0EDD2A95FE5C9D84B16F7C54D85CCE4C8A182810809ED81E97D074884EEDF401CCACDAEAD82C14D06B68AEA6CE14B861B0CFD16090CBBF469C5E084314C0D8D3960EA06A3426D8B3FE762E00D09BDA374F3AE2CBEDE2838FF89D81DEB3013090E44199AED604963EAF919914CE04F207AC82CD4351FEF7B2D94393066FE4D44E3CC5952E75EB6F3714058915DE0EE184D8C55300F576A8B82A863E81AF33417BD4CFC94E7A61263B39F01F6E2E70748B6E5E59CF6CA01B0028C93BBBCEBC548F987F10755BF33CA585CB41CF578DF5FFE37924E3C2C072ED1DAC9162176972971E79B62FB208F1A73BF0361E2993DCCCD3110C34D839D18DD43A5E8F0D941E99ADCF441405F32107671B2D8B2244F7BA92DCED587A210FE8FF43C616ACB5E766E6AF2CEB03599BA3DE376EB5735EF16143953D1FDDB7E9F2874B0D6083DD7EC4386AE003F51CCF2D21EF6059163C5152174423F57119D0FCE627D763D81C10AA1329F74C8D445437BA6718A33DB6E79375172B2AE3591821978D520824E2D2FF898B7F4C867FF462722BC07EADAD389A910B6F65429DA129735FE049E3ECB3889F6047CF2BD2A88D50A651B3235D2480E1DA5A35247FA76C831736399D37E8D033C1D051C9B6A99AB80B1313FA24C5C59766E6C51A38FE9F1186A767EEBD0D88001AE0246CD4EBE2C979DE82C30BBDB98B4744F11F9E639EDDD8C194D7911201A8FA745991B4D8A5709B62A21B63B9762913D36CE995C2D6B79151E8D83838CD1F38840A9417255DD166B7A3584499003FB625611404C95B960DF0DB1BCF1574B0965DBD834EE148117D5E05A7CC7CC1A865618A2BE4854DB8935CDA1E68BD8D09E72F0AC9053C882C4ABA4004A614D10505300B6176CA1F324E22E7824299F9C40755B71D82B679547F06AD48BE66D68072C9390233C933F80A14F8D4A6B0B4E1970E1ACC1BEA7F5D3BE224448F857BAB68AEFA6D8CB819B64294A12997916CDBF56E9A8D002DD065F12C61823F4FC214508232E431F0B6898475BB5DD0D7D528E840C22809AF7E15363724A613ACCFBE2B37438C159CE14CB0C98BFD499C08DAC0CF45D821CC2FA47319B6FB4CED7E5985EC8274DE09071D3C10DA5BF9E522B01CE91D66B91795D3D22C00483454275DD2BBDD7C2DCC4A167E5D7FCDBB9F6208CD4C9A485FAAEB809A7711DAC2865CED4306474B22B4448F85DF33417F3FACE1C05D42703ED313042A05DE0362740130188ECB445BB255DC76EE8443F733117F8351F17603175554FEB00B7FF54D80786F305CDE18CD5EC56EC0962A3E04482DCE3622D040D24C40F2E8A14A447659D6C561F2FFEE68F8D3DE511B23E8B172A01A3EDA4D3780E74C677244330E9AEFF019FE07BE3D33F322F9CE2214B9D9CFF99D05A59E47551432AE76F4CD4F8DD51520FFE811B4B93CD6219C81B63B1D627785C2A0FC22E3AEA86CEEE1F7FBC4EFCB46DDFBCD88A02F3B4E67C5FF2E8DC68BF16C74699BBB628902F72C3DEBC8BF5DF706D47A605A107DAA0014139CE40F0D46D8D6DC7";
        String sk = "[encoded data, 8064 characters]";

        NISTSecureRandom random = new NISTSecureRandom(Hex.decode(seed), null);

        DilithiumKeyPairGenerator keyGen = new DilithiumKeyPairGenerator();

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

        AsymmetricCipherKeyPair keyPair = keyGen.generateKeyPair();

        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
    {
        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));
        }
    }
}
