package org.bouncycastle.pqc.crypto.test;

import java.io.IOException;
import java.security.SecureRandom;
import java.text.ParseException;
import java.util.List;

import junit.framework.TestCase;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.crypto.digests.SHA256Digest;
import org.bouncycastle.crypto.digests.SHA512Digest;
import org.bouncycastle.pqc.crypto.xmss.XMSS;
import org.bouncycastle.pqc.crypto.xmss.XMSSNode;
import org.bouncycastle.pqc.crypto.xmss.XMSSParameters;
import org.bouncycastle.pqc.crypto.xmss.XMSSSignature;
import org.bouncycastle.pqc.crypto.xmss.XMSSUtil;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;

/**
 * Test cases for XMSS class.
 */
public class XMSSTest
    extends TestCase
{

    public void testAuthPath()
    {
        XMSSParameters xmssParams = new XMSSParameters(10, new SHA256Digest());
        XMSS xmss = new XMSS(xmssParams, new NullPRNG());
        xmss.generateKeys();
        String[] expectedAuthPath = {
            "e0e81477f50e2b8719d1d9a3c1a55d5460711b88eb4f2c8df8d17d38832de9d4",
            "251142aa2cdb5a9d81a0c7a98e2c3c95bf96e73c68c78d467186eeac1dffce38",
            "53a6a544a1331bf650758d64bb8fab1104170824894fdd8a284d281aec9d6270",
            "5683d8f8de11a5845b9d1b5cedb7fd0bd76da3341e75ddd05f261520228664d0",
            "16b265f881e4b7bd303b7c38b8adcf31a7c211bcddb5260080105c1b70c41ee9",
            "9d421eefaf1e51d762f245853cc36250bfbb3f547770303dc8a38d03dec2d0a7",
            "233b634f83a2456f279126c7fcb47f9301cb2d5e8db69f4721f68000d78cffb3",
            "9317371748b85e824bd82284875b7bb65604c94ff0b8816cef677ade31963bd5",
            "14ba59e3d998735a963d018b840a80be11a52ccdadaceb7ed45f0763e6cd49ca",
            "6af97f6bde85e4569178192bab4fd59d8c2b1fec181f27a8fa2c17552c16e3b5"
        };
        byte[] signature = xmss.sign(new byte[1024]);
        XMSSSignature sig = new XMSSSignature.Builder(xmssParams).withSignature(signature).build();

        assertEquals(NISTObjectIdentifiers.id_sha256, xmssParams.getTreeDigestOID());

        List<XMSSNode> authPath = sig.getAuthPath();
        for (int i = 0; i < authPath.size(); i++)
        {
            assertEquals(expectedAuthPath[i], Hex.toHexString(authPath.get(i).getValue()));
        }
    }

    public void testGenKeyPairSHA256()
        throws Exception
    {
        XMSSParameters xmssParams = new XMSSParameters(10, new SHA256Digest());
        XMSS xmss = new XMSS(xmssParams, new NullPRNG());
        xmss.generateKeys();
        byte[] privateKey = xmss.exportPrivateKey().getEncoded();
        byte[] publicKey = xmss.exportPublicKey().getEncoded();
        String expectedPrivateKey = "0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000073c3fc6de1195aa5d69f9dafc9db8504aa8059115e8319ca15cf58a1c83c0de3";
        String expectedPublicKey = "0000000173c3fc6de1195aa5d69f9dafc9db8504aa8059115e8319ca15cf58a1c83c0de30000000000000000000000000000000000000000000000000000000000000000";
        byte[] strippedPrivateKey = XMSSUtil.extractBytesAtOffset(privateKey, 0, (Hex.decode(expectedPrivateKey).length));
  
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedPrivateKey), strippedPrivateKey));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedPublicKey), publicKey));
    }

    public void testGenKeyPairSHA512()
        throws Exception
    {
        XMSSParameters xmssParams = new XMSSParameters(10, new SHA512Digest());
        XMSS xmss = new XMSS(xmssParams, new NullPRNG());
        xmss.generateKeys();
        byte[] privateKey = xmss.exportPrivateKey().getEncoded();
        byte[] publicKey = xmss.exportPublicKey().getEncoded();
        String expectedPrivateKey = "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000f48df90f8e217076d8af6327691321bdcf63668c4bd28d021d49f2334eca845fa3073991049286c0eef5dc7f23ec0b31f5c1bd1e5b8edb2403ae02f292f6f30e";
        String expectedPublicKey = "00000004f48df90f8e217076d8af6327691321bdcf63668c4bd28d021d49f2334eca845fa3073991049286c0eef5dc7f23ec0b31f5c1bd1e5b8edb2403ae02f292f6f30e00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000";
        byte[] strippedPrivateKey = XMSSUtil.extractBytesAtOffset(privateKey, 0, (Hex.decode(expectedPrivateKey).length));

        assertEquals(true, Arrays.areEqual(Hex.decode(expectedPrivateKey), strippedPrivateKey));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedPublicKey), publicKey));
    }

    public void testSignSHA256()
    {
        XMSSParameters params = new XMSSParameters(10, new SHA256Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        byte[] message = new byte[1024];
        byte[] sig1 = xmss.sign(message);
        byte[] sig2 = xmss.sign(message);
        byte[] sig3 = xmss.sign(message);
        String expectedSig1 = "[encoded data, 5000 characters]";
        String expectedSig2 = "[encoded data, 5000 characters]";
        String expectedSig3 = "[encoded data, 5000 characters]";

        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig1), sig1));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig2), sig2));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig3), sig3));
    }

    public void testSignSHA256CompleteEvenHeight1()
    {
        final String[] signatures = {
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]",
            "[encoded data, 4616 characters]"};
        int height = 4;
        XMSSParameters params = new XMSSParameters(height, new SHA256Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        for (int i = 0; i < (1 << height); i++)
        {
            byte[] signature = xmss.sign(new byte[1024]);
            assertEquals(signatures[i], Hex.toHexString(signature));
        }
        try
        {
            xmss.sign(new byte[1024]);
            fail();
        }
        catch (Exception ex)
        {
        }
    }

    public void testSignSHA256CompleteEvenHeight2()
    {
        final String[] signatures = {
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]",
            "[encoded data, 5000 characters]"};
        int height = 10;
        XMSSParameters params = new XMSSParameters(height, new SHA256Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        for (int i = 0; i < (1 << height); i++)
        {
            byte[] signature = xmss.sign(new byte[1024]);
            switch (i)
            {
            case 0x005b:
                assertEquals(signatures[0], Hex.toHexString(signature));
                break;
            case 0x0822:
                assertEquals(signatures[1], Hex.toHexString(signature));
                break;
            case 0x00b0:
                assertEquals(signatures[2], Hex.toHexString(signature));
                break;
            case 0x016d:
                assertEquals(signatures[3], Hex.toHexString(signature));
                break;
            case 0x0189:
                assertEquals(signatures[4], Hex.toHexString(signature));
                break;
            case 0x01f9:
                assertEquals(signatures[5], Hex.toHexString(signature));
                break;
            case 0x02a1:
                assertEquals(signatures[6], Hex.toHexString(signature));
                break;
            case 0x038b:
                assertEquals(signatures[7], Hex.toHexString(signature));
                break;
            case 0x03fe:
                assertEquals(signatures[8], Hex.toHexString(signature));
                break;
            case 0x03ff:
                assertEquals(signatures[9], Hex.toHexString(signature));
                break;
            }
        }
        try
        {
            xmss.sign(new byte[1024]);
            fail();
        }
        catch (Exception ex)
        {
        }
    }

    public void testSignSHA512()
    {
        XMSSParameters params = new XMSSParameters(10, new SHA512Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        byte[] message = new byte[1024];
        byte[] sig1 = xmss.sign(message);
        byte[] sig2 = xmss.sign(message);
        byte[] sig3 = xmss.sign(message);
        String expectedSig1 = "[encoded data, 18184 characters]";
        String expectedSig2 = "[encoded data, 18184 characters]";
        String expectedSig3 = "[encoded data, 18184 characters]";
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig1), sig1));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig2), sig2));
        assertEquals(true, Arrays.areEqual(Hex.decode(expectedSig3), sig3));
    }

    public void testVerifySignatureSHA256()
        throws Exception
    {
        XMSSParameters params = new XMSSParameters(4, new SHA256Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        byte[] msg1 = new byte[1024];

        for (int i = 0; i < 3; i++)
        {
            byte[] publicKey = xmss.exportPublicKey().getEncoded();
            xmss.sign(msg1);
            byte[] signature = xmss.sign(msg1);
            try
            {
                assertEquals(true, xmss.verifySignature(msg1, signature, publicKey));
            }
            catch (ParseException ex)
            {
                ex.printStackTrace();
                fail();
            }
            byte[] msg2 = new byte[1024];
            msg2[0] = 0x01;
            try
            {
                assertEquals(false, xmss.verifySignature(msg2, signature, publicKey));
            }
            catch (ParseException ex)
            {
                ex.printStackTrace();
                fail();
            }
        }
    }

    public void testVerifySignatureSHA512()
        throws Exception
    {
        XMSSParameters params = new XMSSParameters(4, new SHA512Digest());
        XMSS xmss = new XMSS(params, new NullPRNG());
        xmss.generateKeys();
        byte[] msg1 = new byte[1024];

        for (int i = 0; i < 3; i++)
        {
            byte[] publicKey = xmss.exportPublicKey().getEncoded();
            xmss.sign(msg1);
            byte[] signature = xmss.sign(msg1);
            try
            {
                assertEquals(true, xmss.verifySignature(msg1, signature, publicKey));
            }
            catch (ParseException ex)
            {
                ex.printStackTrace();
                fail();
            }
            byte[] msg2 = new byte[1024];
            msg2[0] = 0x01;
            try
            {
                assertEquals(false, xmss.verifySignature(msg2, signature, publicKey));
            }
            catch (ParseException ex)
            {
                ex.printStackTrace();
                fail();
            }
        }
    }

    public void testImportStateSHA256()
        throws IOException, ClassNotFoundException
    {
        XMSSParameters params = new XMSSParameters(4, new SHA256Digest());
        XMSS xmss1 = new XMSS(params, new NullPRNG());
        xmss1.generateKeys();
        byte[] msg1 = new byte[1024];
        byte[] msg2 = new byte[2048];
        byte[] msg3 = new byte[3096];
        Arrays.fill(msg1, (byte)0xaa);
        Arrays.fill(msg2, (byte)0xbb);
        Arrays.fill(msg3, (byte)0xcc);
        byte[] signature1 = xmss1.sign(msg1);
        byte[] signature2 = xmss1.sign(msg2);
        byte[] exportedPrivateKey = xmss1.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss1.exportPublicKey().getEncoded();
        byte[] signature3 = xmss1.sign(msg3);

        XMSS xmss2 = new XMSS(params, new NullPRNG());

        xmss2.importState(exportedPrivateKey, exportedPublicKey);

        byte[] signature4 = xmss2.sign(msg3);
        assertEquals(true, Arrays.areEqual(signature3, signature4));
        xmss2.generateKeys();
        try
        {
            assertEquals(true, xmss2.verifySignature(msg1, signature1, exportedPublicKey));
            assertEquals(true, xmss2.verifySignature(msg2, signature2, exportedPublicKey));
            assertEquals(true, xmss2.verifySignature(msg3, signature3, exportedPublicKey));
            assertEquals(false, xmss2.verifySignature(msg1, signature3, exportedPublicKey));
            assertEquals(false, xmss2.verifySignature(msg2, signature3, exportedPublicKey));
        }
        catch (ParseException ex)
        {
            ex.printStackTrace();
            fail();
        }
    }

    public void testImportKeysSHA512()
        throws IOException, ClassNotFoundException
    {
        XMSSParameters params = new XMSSParameters(4, new SHA512Digest());
        XMSS xmss1 = new XMSS(params, new NullPRNG());
        xmss1.generateKeys();
        byte[] msg1 = new byte[1024];
        byte[] msg2 = new byte[2048];
        byte[] msg3 = new byte[3096];
        Arrays.fill(msg1, (byte)0xaa);
        Arrays.fill(msg2, (byte)0xbb);
        Arrays.fill(msg3, (byte)0xcc);
        byte[] signature1 = xmss1.sign(msg1);
        byte[] signature2 = xmss1.sign(msg2);
        byte[] exportedPrivateKey = xmss1.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss1.exportPublicKey().getEncoded();
        byte[] signature3 = xmss1.sign(msg3);

        XMSS xmss2 = new XMSS(params, new NullPRNG());

            xmss2.importState(exportedPrivateKey, exportedPublicKey);

        byte[] signature4 = xmss2.sign(msg3);
        assertEquals(true, Arrays.areEqual(signature3, signature4));
        xmss2.generateKeys();
        try
        {
            assertEquals(true, xmss2.verifySignature(msg1, signature1, exportedPublicKey));
            assertEquals(true, xmss2.verifySignature(msg2, signature2, exportedPublicKey));
            assertEquals(true, xmss2.verifySignature(msg3, signature3, exportedPublicKey));
            assertEquals(false, xmss2.verifySignature(msg1, signature3, exportedPublicKey));
            assertEquals(false, xmss2.verifySignature(msg2, signature3, exportedPublicKey));
        }
        catch (ParseException ex)
        {
            ex.printStackTrace();
            fail();
        }
    }

    public void testRandom()
        throws IOException, ClassNotFoundException
    {
        XMSSParameters params = new XMSSParameters(4, new SHA512Digest());
        XMSS xmss1 = new XMSS(params, new SecureRandom());
        xmss1.generateKeys();
        byte[] publicKey = xmss1.exportPublicKey().getEncoded();
        byte[] message = new byte[1024];

        for (int i = 0; i < 5; i++)
        {
            xmss1.sign(message);
        }
        byte[] signature = xmss1.sign(message);
        assertTrue(Arrays.areEqual(publicKey, xmss1.exportPublicKey().getEncoded()));
        try
        {
            xmss1.verifySignature(message, signature, publicKey);
        }
        catch (ParseException e)
        {
            e.printStackTrace();
            fail();
        }
        assertTrue(Arrays.areEqual(publicKey, xmss1.exportPublicKey().getEncoded()));
        xmss1.sign(message);
        byte[] privateKey7 = xmss1.exportPrivateKey().getEncoded();
        try
        {
            xmss1.verifySignature(message, signature, publicKey);
        }
        catch (ParseException e)
        {
            e.printStackTrace();
            fail();
        }
        assertTrue(Arrays.areEqual(privateKey7, xmss1.exportPrivateKey().getEncoded()));
        byte[] signature7 = xmss1.sign(message);

        xmss1.importState(privateKey7, publicKey);

        byte[] signature7AfterImport = xmss1.sign(message);
        assertTrue(Arrays.areEqual(signature7AfterImport, signature7));

        XMSSParameters params2 = new XMSSParameters(4, new SHA512Digest());
        XMSS xmss2 = new XMSS(params2, new SecureRandom());
        try
        {
            boolean valid = xmss2.verifySignature(message, signature7, publicKey);
            assertTrue(valid);
            valid = xmss2.verifySignature(message, signature, publicKey);
            assertTrue(valid);
        }
        catch (ParseException e)
        {
            e.printStackTrace();
            fail();
        }

        XMSS xmss3 = new XMSS(params, new NullPRNG());

            xmss3.importState(privateKey7, publicKey);

        byte[] signatureAgain = xmss3.sign(message);
        assertTrue(Arrays.areEqual(signatureAgain, signature7));
    }

    public void testBDSImport()
        throws IOException, ClassNotFoundException
    {
        XMSSParameters params = new XMSSParameters(4, new SHA256Digest());
        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();
        byte[] exportedPrivateKey = xmss.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss.exportPublicKey().getEncoded();

        xmss.importState(exportedPrivateKey, exportedPublicKey);

        byte[] sig1 = xmss.sign(new byte[1024]);

        xmss.importState(exportedPrivateKey, exportedPublicKey);

        byte[] sig2 = xmss.sign(new byte[1024]);
        assertEquals(true, Arrays.areEqual(sig1, sig2));
        try
        {
            xmss.importState(exportedPrivateKey, exportedPublicKey);
        }
        catch (Exception ex)
        {
        }
        xmss.sign(new byte[1024]);
        byte[] sig3 = xmss.sign(new byte[1024]);
        assertEquals(false, Arrays.areEqual(sig1, sig3));
        try
        {
            xmss.importState(null, exportedPublicKey);
            fail();
        }
        catch (Exception ex)
        {
        }
        try
        {
            xmss.importState(exportedPrivateKey, null);
            fail();
        }
        catch (Exception ex)
        {
        }
    }

    public void testSignVerifySP800208_SHA256_192()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x0000000d);
        assertNotNull(params);
        assertEquals(24, params.getTreeDigestSize());
        assertEquals(10, params.getHeight());

        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] msg = new byte[1024];
        new SecureRandom().nextBytes(msg);

        byte[] publicKey = xmss.exportPublicKey().getEncoded();
        byte[] signature = xmss.sign(msg);

        assertEquals(true, xmss.verifySignature(msg, signature, publicKey));

        msg[0] ^= 0xff;
        assertEquals(false, xmss.verifySignature(msg, signature, publicKey));
    }

    public void testSignVerifySP800208_SHAKE256_256()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x00000010);
        assertNotNull(params);
        assertEquals(32, params.getTreeDigestSize());
        assertEquals(10, params.getHeight());

        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] msg = new byte[1024];
        new SecureRandom().nextBytes(msg);

        byte[] publicKey = xmss.exportPublicKey().getEncoded();
        byte[] signature = xmss.sign(msg);

        assertEquals(true, xmss.verifySignature(msg, signature, publicKey));

        msg[0] ^= 0xff;
        assertEquals(false, xmss.verifySignature(msg, signature, publicKey));
    }

    public void testSignVerifySP800208_SHAKE256_192()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x00000013);
        assertNotNull(params);
        assertEquals(24, params.getTreeDigestSize());
        assertEquals(10, params.getHeight());

        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] msg = new byte[1024];
        new SecureRandom().nextBytes(msg);

        byte[] publicKey = xmss.exportPublicKey().getEncoded();
        byte[] signature = xmss.sign(msg);

        assertEquals(true, xmss.verifySignature(msg, signature, publicKey));

        msg[0] ^= 0xff;
        assertEquals(false, xmss.verifySignature(msg, signature, publicKey));
    }

    public void testKeyImportSP800208_SHA256_192()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x0000000d);
        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] exportedPrivateKey = xmss.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss.exportPublicKey().getEncoded();

        byte[] msg = new byte[1024];
        byte[] sig1 = xmss.sign(msg);

        xmss.importState(exportedPrivateKey, exportedPublicKey);
        byte[] sig2 = xmss.sign(msg);

        assertEquals(true, Arrays.areEqual(sig1, sig2));
    }

    public void testKeyImportSP800208_SHAKE256_256()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x00000010);
        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] exportedPrivateKey = xmss.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss.exportPublicKey().getEncoded();

        byte[] msg = new byte[1024];
        byte[] sig1 = xmss.sign(msg);

        xmss.importState(exportedPrivateKey, exportedPublicKey);
        byte[] sig2 = xmss.sign(msg);

        assertEquals(true, Arrays.areEqual(sig1, sig2));
    }

    public void testKeyImportSP800208_SHAKE256_192()
        throws Exception
    {
        XMSSParameters params = XMSSParameters.lookupByOID(0x00000013);
        XMSS xmss = new XMSS(params, new SecureRandom());
        xmss.generateKeys();

        byte[] exportedPrivateKey = xmss.exportPrivateKey().getEncoded();
        byte[] exportedPublicKey = xmss.exportPublicKey().getEncoded();

        byte[] msg = new byte[1024];
        byte[] sig1 = xmss.sign(msg);

        xmss.importState(exportedPrivateKey, exportedPublicKey);
        byte[] sig2 = xmss.sign(msg);

        assertEquals(true, Arrays.areEqual(sig1, sig2));
    }
}
