package org.bouncycastle.pqc.jcajce.provider.test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.SecureRandom;
import java.security.Security;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import junit.framework.TestCase;
import org.bouncycastle.asn1.ASN1BitString;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.pqc.jcajce.interfaces.DilithiumKey;
import org.bouncycastle.pqc.jcajce.interfaces.DilithiumPrivateKey;
import org.bouncycastle.pqc.jcajce.spec.DilithiumParameterSpec;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Strings;
import org.bouncycastle.util.encoders.Hex;

/**
 * Dilithum now in BC provider
 */
public class DilithiumTest
    extends TestCase
{
    byte[] msg = Strings.toByteArray("Hello World!");

    public void setUp()
    {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null)
        {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    public void testPrivateKeyRecovery()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("Dilithium", "BC");

        kpg.initialize(DilithiumParameterSpec.dilithium3, new DilithiumTest.RiggedRandom());

        KeyPair kp = kpg.generateKeyPair();

        KeyFactory kFact = KeyFactory.getInstance("Dilithium", "BC");

        DilithiumKey privKey = (DilithiumKey)kFact.generatePrivate(new PKCS8EncodedKeySpec(kp.getPrivate().getEncoded()));

        assertEquals(kp.getPrivate(), privKey);
        assertEquals(kp.getPrivate().getAlgorithm(), privKey.getAlgorithm());
        assertEquals(kp.getPrivate().hashCode(), privKey.hashCode());

        assertEquals(((DilithiumPrivateKey)kp.getPrivate()).getPublicKey(), ((DilithiumPrivateKey)privKey).getPublicKey());

        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        ObjectOutputStream oOut = new ObjectOutputStream(bOut);

        oOut.writeObject(privKey);

        oOut.close();

        ObjectInputStream oIn = new ObjectInputStream(new ByteArrayInputStream(bOut.toByteArray()));

        DilithiumKey privKey2 = (DilithiumKey)oIn.readObject();

        assertEquals(privKey, privKey2);

        assertEquals(kp.getPublic(), ((DilithiumPrivateKey)privKey2).getPublicKey());
        assertEquals(kp.getPrivate().getAlgorithm(), privKey2.getAlgorithm());
        assertEquals(kp.getPrivate().hashCode(), privKey2.hashCode());

        assertEquals(((DilithiumPrivateKey)privKey).getPublicKey(), ((DilithiumPrivateKey)privKey2).getPublicKey());
    }

    public void testPublicKeyRecovery()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("Dilithium", "BC");

        kpg.initialize(DilithiumParameterSpec.dilithium5, new DilithiumTest.RiggedRandom());

        KeyPair kp = kpg.generateKeyPair();

        KeyFactory kFact = KeyFactory.getInstance("Dilithium", "BC");

        DilithiumKey pubKey = (DilithiumKey)kFact.generatePublic(new X509EncodedKeySpec(kp.getPublic().getEncoded()));

        assertEquals(kp.getPublic(), pubKey);
        assertEquals(kp.getPublic().getAlgorithm(), pubKey.getAlgorithm());
        assertEquals(kp.getPublic().hashCode(), pubKey.hashCode());

        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        ObjectOutputStream oOut = new ObjectOutputStream(bOut);

        oOut.writeObject(pubKey);

        oOut.close();

        ObjectInputStream oIn = new ObjectInputStream(new ByteArrayInputStream(bOut.toByteArray()));

        DilithiumKey pubKey2 = (DilithiumKey)oIn.readObject();

        assertEquals(pubKey, pubKey2);
        assertEquals(pubKey.getAlgorithm(), pubKey2.getAlgorithm());
        assertEquals(pubKey.hashCode(), pubKey2.hashCode());
    }

    public void testRestrictedSignature()
        throws Exception
    {
        doTestRestrictedSignature("DILITHIUM2", DilithiumParameterSpec.dilithium2, DilithiumParameterSpec.dilithium5);
        doTestRestrictedSignature("DILITHIUM3", DilithiumParameterSpec.dilithium3, DilithiumParameterSpec.dilithium5);
        doTestRestrictedSignature("DILITHIUM5", DilithiumParameterSpec.dilithium5, DilithiumParameterSpec.dilithium2);
    }

    private void doTestRestrictedSignature(String sigName, DilithiumParameterSpec spec, DilithiumParameterSpec altSpec)
        throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("Dilithium", "BC");

        kpg.initialize(spec, new SecureRandom());

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance(sigName, "BC");

        sig.initSign(kp.getPrivate(), new SecureRandom());

        sig.update(msg, 0, msg.length);

        byte[] s = sig.sign();

        sig = Signature.getInstance(sigName, "BC");

        assertEquals(sigName, sig.getAlgorithm());

        sig.initVerify(kp.getPublic());

        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));

        kpg = KeyPairGenerator.getInstance("Dilithium", "BC");

        kpg.initialize(altSpec, new SecureRandom());

        kp = kpg.generateKeyPair();

        try
        {
            sig.initVerify(kp.getPublic());
            fail("no exception");
        }
        catch (InvalidKeyException e)
        {
            assertEquals("signature configured for " + spec.getName(), e.getMessage());
        }
    }

    public void testRestrictedKeyPairGen()
        throws Exception
    {
        doTestRestrictedKeyPairGen(DilithiumParameterSpec.dilithium2, DilithiumParameterSpec.dilithium5);
        doTestRestrictedKeyPairGen(DilithiumParameterSpec.dilithium3, DilithiumParameterSpec.dilithium5);
        doTestRestrictedKeyPairGen(DilithiumParameterSpec.dilithium5, DilithiumParameterSpec.dilithium2);
    }

    private void doTestRestrictedKeyPairGen(DilithiumParameterSpec spec, DilithiumParameterSpec altSpec)
        throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance(spec.getName(), "BC");

        kpg.initialize(spec, new SecureRandom());

        KeyPair kp = kpg.generateKeyPair();

        assertEquals(spec.getName(), kpg.getAlgorithm());
        assertEquals(spec.getName(), kp.getPublic().getAlgorithm());
        assertEquals(spec.getName(), kp.getPrivate().getAlgorithm());

        kpg = KeyPairGenerator.getInstance(spec.getName(), "BC");

        try
        {
            kpg.initialize(altSpec, new SecureRandom());
            fail("no exception");
        }
        catch (InvalidAlgorithmParameterException e)
        {
            assertEquals("key pair generator locked to " + spec.getName(), e.getMessage());
        }
    }

    public void testDilithiumRandomSig()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("Dilithium", "BC");

        kpg.initialize(DilithiumParameterSpec.dilithium2, new SecureRandom());

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("Dilithium", "BC");

        sig.initSign(kp.getPrivate(), new SecureRandom());

        sig.update(msg, 0, msg.length);

        byte[] s = sig.sign();

        sig = Signature.getInstance("Dilithium", "BC");

        sig.initVerify(kp.getPublic());

        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));
    }

    /**
count = 0
seed = 061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1
mlen = 33
msg = D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8
pk = 1C0EE1111B08003F28E65E8B3BDEB037CF8F221DFCDAF5950EDB38D506D85BEF6177E3DE0D4F1EF5847735947B56D08E841DB2444FA2B729ADEB1417CA7ADF42A1490C5A097F002760C1FC419BE8325AAD0197C52CED80D3DF18E7774265B289912CECA1BE3A90D8A4FDE65C84C610864E47DEECAE3EEA4430B9909559408D11A6ABDB7DB9336DF7F96EAB4864A6579791265FA56C348CB7D2DDC90E133A95C3F6B13601429F5408BD999AA479C1018159550EC55A113C493BE648F4E036DD4F8C809E036B4FBB918C2C484AD8E1747AE05585AB433FDF461AF03C25A773700721AA05F7379FE7F5ED96175D4021076E7F52B60308EFF5D42BA6E093B3D0815EB3496646E49230A9B35C8D41900C2BB8D3B446A23127F7E096D85A1C794AD4C89277904FC6BFEC57B1CDD80DF9955030FDCA741AFBDAC827B13CCD5403588AF4644003C2265DFA4D419DBCCD2064892386518BE9D51C16498275EBECF5CDC7A820F2C29314AC4A6F08B2252AD3CFB199AA42FE0B4FB571975C1020D949E194EE1EAD937BFB550BB3BA8E357A029C29F077554602E1CA2F2289CB9169941C3AAFDB8E58C7F2AC77291FB4147C65F6B031D3EBA42F2ACFD9448A5BC22B476E07CCCEDA2306C554EC9B7AB655F1D7318C2B7E67D5F69BEDF56000FDA98986B5AB1B3A22D8DFD6681697B23A55C96E8710F3F98C044FB15F606313EE56C0F1F5CA0F512E08484FCB358E6E528FFA89F8A866CCFF3C0C5813147EC59AF0470C4AAD0141D34F101DA2E5E1BD52D0D4C9B13B3E3D87D1586105796754E7978CA1C68A7D85DF112B7AB921B359A9F03CBD27A7EAC87A9A80B0B26B4C9657ED85AD7FA2616AB345EB8226F69FC0F48183FF574BCD767B5676413ADB12EA2150A0E97683EE54243C25B7EA8A718606F86993D8D0DACE834ED341EEB724FE3D5FF0BC8B8A7B8104BA269D34133A4CF8300A2D688496B59B6FCBC61AE96062EA1D8E5B410C5671F424417ED693329CD983001FFCD10023D598859FB7AD5FD263547117100690C6CE7438956E6CC57F1B5DE53BB0DC72CE9B6DEAA85789599A70F0051F1A0E25E86D888B00DF36BDBC93EF7217C45ACE11C0790D70E9953E5B417BA2FD9A4CAF82F1FCE6F45F53E215B8355EF61D891DF1C794231C162DD24164B534A9D48467CDC323624C2F95D4402FF9D66AB1191A8124144AFA35D4E31DC86CAA797C31F68B85854CD959C4FAC5EC53B3B56D374B888A9E979A6576B6345EC8522C9606990281BF3EF7C5945D10FD21A2A1D2E5404C5CF21220641391B98BCF825398305B56E58B611FE5253203E3DF0D22466A73B3F0FBE43B9A62928091898B8A0E5B269DB586B0E4DDEF50D682A12D2C1BE824149AA254C6381BB412D77C3F9AA902B688C81715A59C839558556D35ED4FC83B4AB18181F40F73DCD76860D8D8BF94520237C2AC0E463BA09E3C9782380DC07FE4FCBA340CC2003439FD2314610638070D6C9EEA0A70BAE83B5D5D3C5D3FDE26DD01606C8C520158E7E5104020F248CEAA666457C10AEBF068F8A3BD5CE7B52C6AF0ABD5944AF1AD4752C9113976083C03B6C34E1D47ED69644CAD782C2F7D05F8A148961D965FA2E1723A8DDEBC22A90CD783DD1F4DB38FB9AE5A6714B3D946781643D317B7DD79381CF789A9588BB3E193B92A0B60D6B07D047F6984B0609EC57543C394CA8D5E5BCC2A731A79618BD1E2E0DA8704AF98F20F5F8F5452DDF646B95B341DD7F0D2CC1FA15BD9895CD5B65AA1CB94B5E2E788FDA9825B656639193D98328154A4F2C35495A38B6EA0D2FFAAA35DF92C203C7F31CBBCA7BD03C3C2302190CECD161FD49237E4F839E3F3
sk = [encoded data, 5120 characters]
smlen = 2453
sm = [encoded data, 4906 characters]
     * @throws Exception
     */
    public void testDilithiumKATSig()
            throws Exception
    {
        byte[] pubK = Hex.decode("1C0EE1111B08003F28E65E8B3BDEB037CF8F221DFCDAF5950EDB38D506D85BEF6177E3DE0D4F1EF5847735947B56D08E841DB2444FA2B729ADEB1417CA7ADF42A1490C5A097F002760C1FC419BE8325AAD0197C52CED80D3DF18E7774265B289912CECA1BE3A90D8A4FDE65C84C610864E47DEECAE3EEA4430B9909559408D11A6ABDB7DB9336DF7F96EAB4864A6579791265FA56C348CB7D2DDC90E133A95C3F6B13601429F5408BD999AA479C1018159550EC55A113C493BE648F4E036DD4F8C809E036B4FBB918C2C484AD8E1747AE05585AB433FDF461AF03C25A773700721AA05F7379FE7F5ED96175D4021076E7F52B60308EFF5D42BA6E093B3D0815EB3496646E49230A9B35C8D41900C2BB8D3B446A23127F7E096D85A1C794AD4C89277904FC6BFEC57B1CDD80DF9955030FDCA741AFBDAC827B13CCD5403588AF4644003C2265DFA4D419DBCCD2064892386518BE9D51C16498275EBECF5CDC7A820F2C29314AC4A6F08B2252AD3CFB199AA42FE0B4FB571975C1020D949E194EE1EAD937BFB550BB3BA8E357A029C29F077554602E1CA2F2289CB9169941C3AAFDB8E58C7F2AC77291FB4147C65F6B031D3EBA42F2ACFD9448A5BC22B476E07CCCEDA2306C554EC9B7AB655F1D7318C2B7E67D5F69BEDF56000FDA98986B5AB1B3A22D8DFD6681697B23A55C96E8710F3F98C044FB15F606313EE56C0F1F5CA0F512E08484FCB358E6E528FFA89F8A866CCFF3C0C5813147EC59AF0470C4AAD0141D34F101DA2E5E1BD52D0D4C9B13B3E3D87D1586105796754E7978CA1C68A7D85DF112B7AB921B359A9F03CBD27A7EAC87A9A80B0B26B4C9657ED85AD7FA2616AB345EB8226F69FC0F48183FF574BCD767B5676413ADB12EA2150A0E97683EE54243C25B7EA8A718606F86993D8D0DACE834ED341EEB724FE3D5FF0BC8B8A7B8104BA269D34133A4CF8300A2D688496B59B6FCBC61AE96062EA1D8E5B410C5671F424417ED693329CD983001FFCD10023D598859FB7AD5FD263547117100690C6CE7438956E6CC57F1B5DE53BB0DC72CE9B6DEAA85789599A70F0051F1A0E25E86D888B00DF36BDBC93EF7217C45ACE11C0790D70E9953E5B417BA2FD9A4CAF82F1FCE6F45F53E215B8355EF61D891DF1C794231C162DD24164B534A9D48467CDC323624C2F95D4402FF9D66AB1191A8124144AFA35D4E31DC86CAA797C31F68B85854CD959C4FAC5EC53B3B56D374B888A9E979A6576B6345EC8522C9606990281BF3EF7C5945D10FD21A2A1D2E5404C5CF21220641391B98BCF825398305B56E58B611FE5253203E3DF0D22466A73B3F0FBE43B9A62928091898B8A0E5B269DB586B0E4DDEF50D682A12D2C1BE824149AA254C6381BB412D77C3F9AA902B688C81715A59C839558556D35ED4FC83B4AB18181F40F73DCD76860D8D8BF94520237C2AC0E463BA09E3C9782380DC07FE4FCBA340CC2003439FD2314610638070D6C9EEA0A70BAE83B5D5D3C5D3FDE26DD01606C8C520158E7E5104020F248CEAA666457C10AEBF068F8A3BD5CE7B52C6AF0ABD5944AF1AD4752C9113976083C03B6C34E1D47ED69644CAD782C2F7D05F8A148961D965FA2E1723A8DDEBC22A90CD783DD1F4DB38FB9AE5A6714B3D946781643D317B7DD79381CF789A9588BB3E193B92A0B60D6B07D047F6984B0609EC57543C394CA8D5E5BCC2A731A79618BD1E2E0DA8704AF98F20F5F8F5452DDF646B95B341DD7F0D2CC1FA15BD9895CD5B65AA1CB94B5E2E788FDA9825B656639193D98328154A4F2C35495A38B6EA0D2FFAAA35DF92C203C7F31CBBCA7BD03C3C2302190CECD161FD49237E4F839E3F3");
        byte[] privK = Hex.decode("[encoded data, 5120 characters]");
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");
        byte[] s = Hex.decode("[encoded data, 4840 characters]");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("Dilithium", "BC");
        SecureRandom katRandom = new NISTSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"), null);

        kpg.initialize(DilithiumParameterSpec.dilithium2, katRandom);

        KeyPair kp = kpg.generateKeyPair();

        SubjectPublicKeyInfo pubInfo = SubjectPublicKeyInfo.getInstance(kp.getPublic().getEncoded());

        ASN1BitString pubSeq = pubInfo.getPublicKeyData();
        assertTrue(Arrays.areEqual(pubSeq.getOctets(), pubK));

        PrivateKeyInfo privInfo = PrivateKeyInfo.getInstance(kp.getPrivate().getEncoded());
        ASN1OctetString seq = ASN1OctetString.getInstance(privInfo.parsePrivateKey());

        assertTrue(Arrays.areEqual(seq.getOctets(), privK));

        Signature sig = Signature.getInstance("Dilithium", "BC");

        sig.initSign(kp.getPrivate());

        sig.update(msg, 0, msg.length);

        byte[] genS = sig.sign();

        assertTrue(Arrays.areEqual(s, genS));

        sig = Signature.getInstance("Dilithium", "BC");

        sig.initVerify(kp.getPublic());

        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));

        // check randomisation

        sig.initSign(kp.getPrivate(), new SecureRandom());

        sig.update(msg, 0, msg.length);

        genS = sig.sign();

        assertFalse(Arrays.areEqual(s, genS));

        sig = Signature.getInstance("Dilithium", "BC");

        sig.initVerify(kp.getPublic());

        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));
    }

    private static class RiggedRandom
            extends SecureRandom
    {
        public void nextBytes(byte[] bytes)
        {
            for (int i = 0; i != bytes.length; i++)
            {
                bytes[i] = (byte)(i & 0xff);
            }
        }
    }
}
