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.AlgorithmParameters;
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.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import junit.framework.TestCase;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.jcajce.interfaces.SLHDSAKey;
import org.bouncycastle.jcajce.interfaces.SLHDSAPrivateKey;
import org.bouncycastle.jcajce.spec.ContextParameterSpec;
import org.bouncycastle.jcajce.spec.SLHDSAParameterSpec;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.pqc.crypto.slhdsa.SLHDSAParameters;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Strings;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.FixedSecureRandom;

/**
 * Test cases for the use of SLH-DSA with the provider.
 */
public class SLHDSATest
    extends TestCase
{
    // test vector courtesy the "Yawning Angel" GO implementation and the SUPERCOP reference implementation.
    static private final byte[] msg = Strings.toByteArray("Cthulhu Fthagn --What a wonderful phrase!Cthulhu Fthagn --Say it and you're crazed!");

    static private final String[] names = new String[] {
        "SLH-DSA-SHA2-128F",
        "SLH-DSA-SHA2-128S",
        "SLH-DSA-SHA2-192F",
        "SLH-DSA-SHA2-192S",
        "SLH-DSA-SHA2-256F",
        "SLH-DSA-SHA2-256S",
        "SLH-DSA-SHAKE-128F",
        "SLH-DSA-SHAKE-128S",
        "SLH-DSA-SHAKE-192F",
        "SLH-DSA-SHAKE-192S",
        "SLH-DSA-SHAKE-256F",
        "SLH-DSA-SHAKE-256S",
        "SLH-DSA-SHA2-128F-WITH-SHA256",
        "SLH-DSA-SHA2-128S-WITH-SHA256",
        "SLH-DSA-SHA2-192F-WITH-SHA512",
        "SLH-DSA-SHA2-192S-WITH-SHA512",
        "SLH-DSA-SHA2-256F-WITH-SHA512",
        "SLH-DSA-SHA2-256S-WITH-SHA512",
        "SLH-DSA-SHAKE-128F-WITH-SHAKE128",
        "SLH-DSA-SHAKE-128S-WITH-SHAKE128",
        "SLH-DSA-SHAKE-192F-WITH-SHAKE256",
        "SLH-DSA-SHAKE-192S-WITH-SHAKE256",
        "SLH-DSA-SHAKE-256F-WITH-SHAKE256",
        "SLH-DSA-SHAKE-256S-WITH-SHAKE256",
    };

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

    public void testParametersAndParamSpecs()
        throws Exception
    {
        for (int i = 0; i != names.length; i++)
        {
            assertEquals(names[i], SLHDSAParameterSpec.fromName(names[i]).getName());
        }

        SLHDSAParameters slhdsaParameters[] = new SLHDSAParameters[]
            {
                SLHDSAParameters.sha2_128f,
                SLHDSAParameters.sha2_128s,
                SLHDSAParameters.sha2_192f,
                SLHDSAParameters.sha2_192s,
                SLHDSAParameters.sha2_256f,
                SLHDSAParameters.sha2_256s,
                SLHDSAParameters.shake_128f,
                SLHDSAParameters.shake_128s,
                SLHDSAParameters.shake_192f,
                SLHDSAParameters.shake_192s,
                SLHDSAParameters.shake_256f,
                SLHDSAParameters.shake_256s,
                SLHDSAParameters.sha2_128f_with_sha256,
                SLHDSAParameters.sha2_128s_with_sha256,
                SLHDSAParameters.sha2_192f_with_sha512,
                SLHDSAParameters.sha2_192s_with_sha512,
                SLHDSAParameters.sha2_256f_with_sha512,
                SLHDSAParameters.sha2_256s_with_sha512,
                SLHDSAParameters.shake_128f_with_shake128,
                SLHDSAParameters.shake_128s_with_shake128,
                SLHDSAParameters.shake_192f_with_shake256,
                SLHDSAParameters.shake_192s_with_shake256,
                SLHDSAParameters.shake_256f_with_shake256,
                SLHDSAParameters.shake_256s_with_shake256
            };

        for (int i = 0; i != names.length; i++)
        {
            assertEquals(names[i], SLHDSAParameterSpec.fromName(slhdsaParameters[i].getName()).getName());
        }
    }

    public void testKeyFactory()
        throws Exception
    {
        KeyPairGenerator kpGen44 = KeyPairGenerator.getInstance("ML-DSA-44");
        KeyPair kp44 = kpGen44.generateKeyPair();

        KeyFactory kFact = KeyFactory.getInstance("HASH-SLH-DSA", "BC");

        ASN1ObjectIdentifier[] oids = new ASN1ObjectIdentifier[] {
            NISTObjectIdentifiers.id_slh_dsa_sha2_128f,
            NISTObjectIdentifiers.id_slh_dsa_sha2_128s,
            NISTObjectIdentifiers.id_slh_dsa_sha2_192f,
            NISTObjectIdentifiers.id_slh_dsa_sha2_192s,
            NISTObjectIdentifiers.id_slh_dsa_sha2_256f,
            NISTObjectIdentifiers.id_slh_dsa_sha2_256s,
            NISTObjectIdentifiers.id_slh_dsa_shake_128f,
            NISTObjectIdentifiers.id_slh_dsa_shake_128s,
            NISTObjectIdentifiers.id_slh_dsa_shake_192f,
            NISTObjectIdentifiers.id_slh_dsa_shake_192s,
            NISTObjectIdentifiers.id_slh_dsa_shake_256f,
            NISTObjectIdentifiers.id_slh_dsa_shake_256s,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_128f_with_sha256,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_128s_with_sha256,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_192f_with_sha512,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_192s_with_sha512,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_256f_with_sha512,
            NISTObjectIdentifiers.id_hash_slh_dsa_sha2_256s_with_sha512,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_128f_with_shake128,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_128s_with_shake128,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_192f_with_shake256,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_192s_with_shake256,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_256f_with_shake256,
            NISTObjectIdentifiers.id_hash_slh_dsa_shake_256s_with_shake256,
        };
        
        for (int i = 0; i != names.length; i++)
        {
            KeyPairGenerator kpGen = KeyPairGenerator.getInstance(names[i]);
            KeyPair kp = kpGen.generateKeyPair();

            tryKeyFact(KeyFactory.getInstance(names[i], "BC"), kp, kp44, "2.16.840.1.101.3.4.3.17");
            tryKeyFact(KeyFactory.getInstance(oids[i].toString(), "BC"), kp, kp44, "2.16.840.1.101.3.4.3.17");
        }
    }

    private void tryKeyFact(KeyFactory kFact, KeyPair kpValid, KeyPair kpInvalid, String oid)
        throws Exception
    {
        kFact.generatePrivate(new PKCS8EncodedKeySpec(kpValid.getPrivate().getEncoded()));
        kFact.generatePublic(new X509EncodedKeySpec(kpValid.getPublic().getEncoded()));

        try
        {
            kFact.generatePrivate(new PKCS8EncodedKeySpec(kpInvalid.getPrivate().getEncoded()));
            fail("no exception");
        }
        catch (InvalidKeySpecException e)
        {
            assertEquals("incorrect algorithm OID for key: " + oid, e.getMessage());
        }
        try
        {
            kFact.generatePublic(new X509EncodedKeySpec(kpInvalid.getPublic().getEncoded()));
            fail("no exception");
        }
        catch (InvalidKeySpecException e)
        {
            assertEquals("incorrect algorithm OID for key: " + oid, e.getMessage());
        }
    }

//    public void testSphincsDefaultKeyGen()
//        throws Exception
//    {
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");
//
//        kpg.initialize(new SLHDSAKeyGenParameterSpec(), new RiggedRandom());
//
//        KeyPair kp = kpg.generateKeyPair();
//
//        SLHDSAKey pub = (SLHDSAKey)kp.getPublic();
//
//        assertTrue(Arrays.areEqual(expSha2Pub, pub.getKeyData()));
//
//        SLHDSAKey priv = (SLHDSAKey)kp.getPrivate();
//
//        assertTrue(Arrays.areEqual(expSha2Priv, priv.getKeyData()));
//
//        KeyFactory keyFact = KeyFactory.getInstance("SLH-DSA", "BC");
//
//        SLHDSAKey pub2 = (SLHDSAKey)keyFact.generatePublic(new X509EncodedKeySpec(pub.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha2Pub, pub2.getKeyData()));
//
//        SLHDSAKey priv2 = (SLHDSAKey)keyFact.generatePrivate(new PKCS8EncodedKeySpec(priv.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha2Priv, priv2.getKeyData()));
//    }

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

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, new RiggedRandom());

        KeyPair kp = kpg.generateKeyPair();

        KeyFactory kFact = KeyFactory.getInstance("SLH-DSA", "BC");

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

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

        oOut.writeObject(privKey);

        oOut.close();

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

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

        assertEquals(privKey, privKey2);

        assertEquals(kp.getPublic(), ((SLHDSAPrivateKey)privKey2).getPublicKey());
        assertEquals(((SLHDSAPrivateKey)privKey).getPublicKey(), ((SLHDSAPrivateKey)privKey2).getPublicKey());
    }

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

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, new RiggedRandom());

        KeyPair kp = kpg.generateKeyPair();

        KeyFactory kFact = KeyFactory.getInstance("SLH-DSA", "BC");

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

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

        oOut.writeObject(pubKey);

        oOut.close();

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

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

        assertEquals(pubKey, pubKey2);
    }

//    public void testSphincsDefaultSha2KeyGen()
//        throws Exception
//    {
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");
//
//        kpg.initialize(new SLHDSAKeyGenParameterSpec(SLHDSAKeyGenParameterSpec.SHA512_256), new RiggedRandom());
//
//        KeyPair kp = kpg.generateKeyPair();
//
//        SLHDSAKey pub = (SLHDSAKey)kp.getPublic();
//
//        assertTrue(Arrays.areEqual(expSha2Pub, pub.getKeyData()));
//
//        SLHDSAKey priv = (SLHDSAKey)kp.getPrivate();
//
//        assertTrue(Arrays.areEqual(expSha2Priv, priv.getKeyData()));
//
//        KeyFactory keyFact = KeyFactory.getInstance("SLH-DSA", "BC");
//
//        SLHDSAKey pub2 = (SLHDSAKey)keyFact.generatePublic(new X509EncodedKeySpec(pub.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha2Pub, pub2.getKeyData()));
//
//        SubjectPublicKeyInfo pkInfo = SubjectPublicKeyInfo.getInstance(pub2.getEncoded());
//
//        assertEquals(new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha512_256), SLHDSAKeyParams.getInstance(pkInfo.getAlgorithm().getParameters()).getTreeDigest());
//
//        SLHDSAKey priv2 = (SLHDSAKey)keyFact.generatePrivate(new PKCS8EncodedKeySpec(priv.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha2Priv, priv2.getKeyData()));
//    }
//
//    public void testSphincsDefaultSha3KeyGen()
//        throws Exception
//    {
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");
//
//        kpg.initialize(new SLHDSAKeyGenParameterSpec(SLHDSAKeyGenParameterSpec.SHA3_256), new RiggedRandom());
//
//        KeyPair kp = kpg.generateKeyPair();
//
//        SLHDSAKey pub = (SLHDSAKey)kp.getPublic();
//
//        assertTrue(Arrays.areEqual(expSha3Pub, pub.getKeyData()));
//
//        SLHDSAKey priv = (SLHDSAKey)kp.getPrivate();
//
//        assertTrue(Arrays.areEqual(expSha3Priv, priv.getKeyData()));
//
//        KeyFactory keyFact = KeyFactory.getInstance("SLH-DSA", "BC");
//
//        SLHDSAKey pub2 = (SLHDSAKey)keyFact.generatePublic(new X509EncodedKeySpec(pub.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha3Pub, pub2.getKeyData()));
//
//        SubjectPublicKeyInfo pkInfo = SubjectPublicKeyInfo.getInstance(pub2.getEncoded());
//
//        assertEquals(new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha3_256), SLHDSAKeyParams.getInstance(pkInfo.getAlgorithm().getParameters()).getTreeDigest());
//
//        SLHDSAKey priv2 = (SLHDSAKey)keyFact.generatePrivate(new PKCS8EncodedKeySpec(priv.getEncoded()));
//
//        assertTrue(Arrays.areEqual(expSha3Priv, priv2.getKeyData()));
//    }
//
//    public void testSphincsSha2Signature()
//        throws Exception
//    {
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");
//
//        kpg.initialize(new SLHDSAKeyGenParameterSpec(SLHDSAKeyGenParameterSpec.SHA512_256), new RiggedRandom());
//
//        KeyPair kp = kpg.generateKeyPair();
//
//        Signature sig = Signature.getInstance("SHA512withSPHINCSPlus", "BC");
//
//        sig.initSign(kp.getPrivate());
//
//        sig.update(msg, 0, msg.length);
//
//        byte[] s = sig.sign();
//
//        assertTrue(Arrays.areEqual(expSha2Sig, s));
//    }
//
//    public void testSphincsSha3Signature()
//        throws Exception
//    {
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");
//
//        kpg.initialize(new SLHDSAKeyGenParameterSpec(SLHDSAKeyGenParameterSpec.SHA3_256), new RiggedRandom());
//
//        KeyPair kp = kpg.generateKeyPair();
//
//        Signature sig = Signature.getInstance("SHA3-512withSPHINCSPlus", "BC");
//
//        sig.initSign(kp.getPrivate());
//
//        sig.update(msg, 0, msg.length);
//
//        byte[] s = sig.sign();
//
//        assertTrue(Arrays.areEqual(expSha3Sig, s));
//    }
//


    public void testSphincsRandomSigSHA2()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initSign(kp.getPrivate(), new FixedSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1")));

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

        byte[] s = sig.sign();
        byte[] expected = Hex.decode("[encoded data, 34176 characters]\n");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

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

        assertTrue(sig.verify(s));
    }

    public void testSphincsRandomSigSHA2WithContext()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initSign(kp.getPrivate(), new FixedSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1")));

        sig.setParameter(new ContextParameterSpec(Strings.toByteArray("Hello, world!")));

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

        byte[] s = sig.sign();

        AlgorithmParameters sp = sig.getParameters();

        ContextParameterSpec sspec = sp.getParameterSpec(ContextParameterSpec.class);

        assertTrue(Arrays.areEqual(Strings.toByteArray("Hello, world!"), sspec.getContext()));
        
        byte[] expected = Hex.decode("[encoded data, 34176 characters]");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

        sig.setParameter(new ContextParameterSpec(Strings.toByteArray("Hello, world!")));
        
        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));

        AlgorithmParameters vp = sig.getParameters();

        ContextParameterSpec vspec = vp.getParameterSpec(ContextParameterSpec.class);

        assertTrue(Arrays.areEqual(Strings.toByteArray("Hello, world!"), vspec.getContext()));
    }

    public void testSLHDSARandomSigSHA2()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initSign(kp.getPrivate(), new FixedSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1")));

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

        byte[] s = sig.sign();
        byte[] expected = Hex.decode("[encoded data, 34176 characters]\n");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

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

        assertTrue(sig.verify(s));
    }

    public void testSLHDSARandomPrehashSigSHA2()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("HASH-SLH-DSA", "BC");

        sig.initSign(kp.getPrivate(), new FixedSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1")));

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

        byte[] s = sig.sign();
        byte[] expected = Hex.decode("[encoded data, 34176 characters]");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("HASH-SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

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

        assertTrue(sig.verify(s));
    }

    public void testSLHDSARandomPrehashSigSHA2WithContext()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("HASH-SLH-DSA", "BC");

        sig.initSign(kp.getPrivate(), new FixedSecureRandom(Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1")));

        sig.setParameter(new ContextParameterSpec(Strings.toByteArray("Hello, world!")));

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

        byte[] s = sig.sign();
        byte[] expected = Hex.decode("[encoded data, 34176 characters]");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("HASH-SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

        sig.setParameter(new ContextParameterSpec(Strings.toByteArray("Hello, world!")));
        
        sig.update(msg, 0, msg.length);

        assertTrue(sig.verify(s));
    }

    // TODO
    /*
    public void testSphincsDeterministicSigSHA2()
        throws Exception
    {
        SecureRandom random = new FixedSecureRandom(Hex.decode("7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E4711E95F8A383854BA16A5DD3E25FF71D3"
            + "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1"));
        byte[] msg = Hex.decode("D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8");

        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_sha2_128f, random);

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initSign(kp.getPrivate());

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

        byte[] s = sig.sign();

        byte[] expected = Hex.decode("[encoded data, 34176 characters]");
        assertTrue(Hex.toHexString(s), Arrays.areEqual(expected, s));

        sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

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

        assertTrue(sig.verify(s));
    }
      */
    public void testSLHDSARandomSigSHAKE()
        throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_shake_256f, new SecureRandom());

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("SLH-DSA", "BC");

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

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

        byte[] s = sig.sign();

        sig = Signature.getInstance("SLH-DSA", "BC");

        sig.initVerify(kp.getPublic());

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

        assertTrue(sig.verify(s));
    }

    public void testSLHDSARandomPrehashSigSHAKE()
        throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("HASH-SLH-DSA", "BC");

        kpg.initialize(SLHDSAParameterSpec.slh_dsa_shake_256f_with_shake256, new SecureRandom());

        KeyPair kp = kpg.generateKeyPair();

        Signature sig = Signature.getInstance("HASH-SLH-DSA", "BC");

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

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

        byte[] s = sig.sign();

        sig = Signature.getInstance("HASH-SLH-DSA", "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);
            }
        }
    }
}

