package org.bouncycastle.jcacje.provider.test;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Security;

import javax.crypto.KEM;
import javax.crypto.SecretKey;

import junit.framework.TestCase;
import org.bouncycastle.jcajce.spec.KEMParameterSpec;
import org.bouncycastle.jcajce.spec.KTSParameterSpec;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.pqc.jcajce.spec.SNTRUPrimeParameterSpec;
import org.bouncycastle.pqc.jcajce.interfaces.SNTRUPrimeKey;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.FixedSecureRandom;


public class SNTRUPrimeKEMTest
    extends TestCase
{
    public void setUp()
    {
        if (Security.getProvider(BouncyCastlePQCProvider.PROVIDER_NAME) == null)
        {
            Security.addProvider(new BouncyCastlePQCProvider());
        }
    }

    // Test for no KDF
    public void testKEMVec()
        throws Exception
    {
            Security.addProvider(new BouncyCastlePQCProvider());

        // count 0
//        byte[] seed = Hex.decode("061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1");
//        byte[] pk = Hex.decode("01031FE0C4328B956E9AFA7B5D86DE2CBF1F0728A1236F685DD2FD221ACDD05CAB64B01AF7A5F3C2837A5B7FE5FE31E366F7F04C5B77C1D6BD3B25A765A84CF9AA08FD9ABA3BB6A4B5C3504064D0803687A04F06B2BB5F7BB95AABF6CB823ACC6E74E98FC468C3E07BCAE284797221A8C8A16168A91A03D066F0EC93F26C8BB6E948BECB9919A72E0FB86F380F1F2C1A2EFCEF871D6ACE4F5E0AD226A80117F7CEEFE566193D872A9ABADB9873D3E9ABB0119CF4133D06F0DA9F0ECB6319CEE796F2FF5B85B4A6754AD4C94346E325BAFF9CC922561CF02B827E2A487ECE6A4AD9ABBDE1D81C218CD5D50A66B8E3EB9870C9F68457BC813A2D827EA7C87951439DD394DA1ED071E8AAEE175BDBFC5CE2ACDE12F78C9B82D9653B970F035A130A1DFBE5B2D4C186C909B9ADF4FB821B2FA21EC03D8D8E095D68F4CC12CEDBE1DE7428A201B04A9575BE758D58340E73602C3467876A4B18218C34A3126DBCC877144E4A42A45005217D9C50F38DBAFAC6477011DA1A1FC749F26B6C747DDADF1F1E4707440A34C22046A6145D3033408212254FBEFAE6A4FCF5AA41BA14EF3C6CA23F3E04CA53A5C2AA03C717A04E928E5DD392ADFBA921528E54944B4183ACBED4D21DFFA3F33D35B9679068B3711AC438F7CAF3042436A9C0BB296ED1C7891B7B782EAABFCE080BDC54CB4EAE3E75536171BF1235397ED37A4394D05D31D7274A1640C1B42912CD9AEBB6B323C370130A854E5BB2253D6F3CA7AFA061C241646C28A8D1EA7AACA73A028920C2BA5DC9EC21F1645B92B35FE810E024CB764CE53F611E61DFC62CEDD8BCD445870D50EBD4DD8021924CC366FADE330AC3826DBF43A328254C367B15718967D72EB4F7EC3BFB692B5BEF357BEF0974739400113F766922F3171964E303D255FE2453BE85AA7B9B9B6408D1655D475BE7EF0351273BC084C9E71B794FAC5C256E87CFEDE5FD8E03EC272B24675B947BB5C1711D894285E026221A152D19DDE65ED293710D195D31968D18E9E2ACF3A0BF5C9F4760B1C20DDCB49FD9A24A3027025A5C090C708E97CD0F273E89DD43180ABEAC4665478732683325406CFBF20BF3059AE57756FAF185612EC199424EC64EC444BC190F0FAA6B2E9A2D96E7814E2FC3BF673B87B7CEC8D6F3A814B5774CF95490E258F0B10269E1ADD8C1D1C4BDD5346ABD921CE3E02A2D051A95E56DFE9A0C655D926FCA45D445170498F6D0870BC8D3F444982E55DE23D59A385E1F18732F7D7F6526289C6659D4363009EBCDF2066411E49E3A8E3D6B312DDCC49169BBF9B13C827A88ACFD5B3E61A9116916F41052A3AAF50ABDA2E7CAA9DB7EA816F44C0F315CB86700F62E25E05C90294FBD55342D62BAFA8BA55BEE7B532D50CD947065E704E625");
//        byte[] sk = Hex.decode("5515559645A956544158655545514489451621815999295514A25949655418096985596A011900615556598565956A6866A916555455214551995455460640548011661525945242549500A515565165659455A52545A6A5445050A2054599A4659591541906682494418654A68A969155864655419119405555105595842994525415591649A569515155058551859A565A8854655454848415444169519555149656009A284AA0664A42991A8809015AA6555A54AA551010460084421102809A028A940502189698410446514812188888644650A081288196648158A61261A94A50908918A5810846011A9925260685A244404A0A585A0289066611A85A5A12408668A6A6591565185058559A18658821104418281602A994A6505012A82A662412A40255690122299A0954295A946110669299104014952500529294624041461249AA142A0201031FE0C4328B956E9AFA7B5D86DE2CBF1F0728A1236F685DD2FD221ACDD05CAB64B01AF7A5F3C2837A5B7FE5FE31E366F7F04C5B77C1D6BD3B25A765A84CF9AA08FD9ABA3BB6A4B5C3504064D0803687A04F06B2BB5F7BB95AABF6CB823ACC6E74E98FC468C3E07BCAE284797221A8C8A16168A91A03D066F0EC93F26C8BB6E948BECB9919A72E0FB86F380F1F2C1A2EFCEF871D6ACE4F5E0AD226A80117F7CEEFE566193D872A9ABADB9873D3E9ABB0119CF4133D06F0DA9F0ECB6319CEE796F2FF5B85B4A6754AD4C94346E325BAFF9CC922561CF02B827E2A487ECE6A4AD9ABBDE1D81C218CD5D50A66B8E3EB9870C9F68457BC813A2D827EA7C87951439DD394DA1ED071E8AAEE175BDBFC5CE2ACDE12F78C9B82D9653B970F035A130A1DFBE5B2D4C186C909B9ADF4FB821B2FA21EC03D8D8E095D68F4CC12CEDBE1DE7428A201B04A9575BE758D58340E73602C3467876A4B18218C34A3126DBCC877144E4A42A45005217D9C50F38DBAFAC6477011DA1A1FC749F26B6C747DDADF1F1E4707440A34C22046A6145D3033408212254FBEFAE6A4FCF5AA41BA14EF3C6CA23F3E04CA53A5C2AA03C717A04E928E5DD392ADFBA921528E54944B4183ACBED4D21DFFA3F33D35B9679068B3711AC438F7CAF3042436A9C0BB296ED1C7891B7B782EAABFCE080BDC54CB4EAE3E75536171BF1235397ED37A4394D05D31D7274A1640C1B42912CD9AEBB6B323C370130A854E5BB2253D6F3CA7AFA061C241646C28A8D1EA7AACA73A028920C2BA5DC9EC21F1645B92B35FE810E024CB764CE53F611E61DFC62CEDD8BCD445870D50EBD4DD8021924CC366FADE330AC3826DBF43A328254C367B15718967D72EB4F7EC3BFB692B5BEF357BEF0974739400113F766922F3171964E303D255FE2453BE85AA7B9B9B6408D1655D475BE7EF0351273BC084C9E71B794FAC5C256E87CFEDE5FD8E03EC272B24675B947BB5C1711D894285E026221A152D19DDE65ED293710D195D31968D18E9E2ACF3A0BF5C9F4760B1C20DDCB49FD9A24A3027025A5C090C708E97CD0F273E89DD43180ABEAC4665478732683325406CFBF20BF3059AE57756FAF185612EC199424EC64EC444BC190F0FAA6B2E9A2D96E7814E2FC3BF673B87B7CEC8D6F3A814B5774CF95490E258F0B10269E1ADD8C1D1C4BDD5346ABD921CE3E02A2D051A95E56DFE9A0C655D926FCA45D445170498F6D0870BC8D3F444982E55DE23D59A385E1F18732F7D7F6526289C6659D4363009EBCDF2066411E49E3A8E3D6B312DDCC49169BBF9B13C827A88ACFD5B3E61A9116916F41052A3AAF50ABDA2E7CAA9DB7EA816F44C0F315CB86700F62E25E05C90294FBD55342D62BAFA8BA55BEE7B532D50CD947065E704E62500D57230D15E762228DA98178A193D7D95284B20E82D74228146FDF68D59B37C5E7F78C0E14E7C40BD4F4D9CF0B189B69983DD39E29AA6439ECAB0A5B294D2A1216CE31CAC6D1B2358D2B0476C4D8002519B6D63639B2D4564E2458C8E06BDEEBE82262570777780F43B110A96547415A69A38EF0ECE0C2AF16B9CF11FF8326E2DE6DD0333B2EDE445AA3A1057824478BFC71FD00DEE601938734810D816F4DADD35F4C152B10CF13957945F8CC47FEF0CE3BDD1F0C5D8CD24B07E3F6D2A5D01717CFD06");
        byte[] ct = Hex.decode("BC60F537CD5FE3038DACE613B8819A0DC920B8FE092474F763A1BD05F69137A1A084AA1A45A5BF1055F9340368D60CE415217C2C382E7EFFCD289D62A0A9E1E7FA2EFD07EF4B75B02DD436535AED897AECD520763AC7F8DFAFEBC122594388F210DF28DE472E748CFB640E8899B07CB4514401DEF9E9C542DC89733F20F9605F641DD1DD6F332DC051C2E5B2C391ACE70FD00FBC251EE1EA2AB27D165F966DA338F8AA970AAF66931F7CC68E2846EE5AD7EAF8E15B9CA333AD1260163E2BE2324C03094E249D3418164A655BD469E506F5B8420BF4FF6080DA230373B18D74AF03F11C9B545E3791CCC9B3B9927B2222983B7A6DE620BE8520B215ECDA45302695C0229C2A9607260BFD3C6EB9D7EAFAB5E47FD5760FA23067BAC1E6A980C4727B0C187A5B6397AC43D3BE24C42A1D2FCFE43512494A7BFC5CC19455DF6AA6F2698A2903831CA2E5F4F84442E6A74BC9D3CC3D6FBEE97D5CF4B6C58D3DC9ADB359FB56CF35FE21F96E885624498ACDF0BAF3A52B7F2564D1EA384DDA7FD32167786C5E010C3BCF5D2E7164BD6ECE815366887250D184F8061E57C3935214F191B38E982132EF4262E91808FF1CFD902B5248F6E7C031A98234DE0D578D1234B7453F2F575A88B622B0F11902C2146F45912E9BDAFA0846AF7B6789E621ABE3C65A4D990A96488B6307BE210887E82BBADFB026B60CA7DF3C9D429D8FEFCCC93A82E345F3B17756F6341F5DB3151974B1B9C7AC4AA33E80923475FE38FFC0AE7F47529360787AC283008971D87FB07369BAD111F1B99FAA59F2E1B5C088A5F31A47ED6A8E30FF6465A7DC79967EA78891C8EBD462D2566A4EA30EDFB8B7BEDD3C4E2ABFEEE7B366125A19B6372CEE2EB2F85514648AE920D292F444E6F90346DF87FD7C107848796219027BA581B3D87704FA8605108CA88ACF262C5C07D9BFB42635A1944634AF0633B471D0D2B2E4D6284894E0E48D58AE5D07E95130C9FBA77A0F124F9348E1A570BA0D009B7B46E03CEB201B17114B935AA91691A7F830346636FD34ECE2CC295D0124AD8CA76CC81E942C1607699F0DFFD3BDB509392FAF8212363F562554B05F756FF252765E6FCAFB04C6F2CF6A76F2E59E14F5BD4EA8E2ECF50B17E632F346D697FFB7AEB82ADDEF5B4DEC0DEE73D2D96B1A33C6AB4BCF7A160B79114A48F224FA037315ADF0485AB63DC2BF09F02606DAD80F13F0ED0241E6F53B21BBBE584794F77D0B783A7093F53AC916A848FF0F9FA3FB08A053C4");
        byte[] ss = Hex.decode("2FFCE8C89E97BBA86A7ECD68088FB116A20601B88B8E88A3783CF2D2ED439769");

        byte[] fixedRandomBytes = Hex.decode("[encoded data, 16000 characters]");
        FixedSecureRandom fixedRandom = new FixedSecureRandom(fixedRandomBytes);

        // Receiver side
        KeyPairGenerator g = KeyPairGenerator.getInstance("SNTRUPrime");
        g.initialize(SNTRUPrimeParameterSpec.sntrup653, fixedRandom);
        KeyPair kp = g.generateKeyPair();
        SNTRUPrimeKey pkR = (SNTRUPrimeKey)kp.getPublic();

        // Sender side
        KEM kemS = KEM.getInstance("SNTRUPrime"); //Should the name be "SNTRUPrime-KEM" ?
        KTSParameterSpec ktsSpec = null;
        KEM.Encapsulator e = kemS.newEncapsulator((PublicKey)pkR, ktsSpec, fixedRandom);
        KEM.Encapsulated enc = e.encapsulate();
        SecretKey secS = enc.key();
        byte[] em = enc.encapsulation();
        byte[] params = enc.params();

        assertTrue(Arrays.areEqual(em, ct));
        assertTrue(Arrays.areEqual(enc.key().getEncoded(), ss));

        // Receiver side
        KEM kemR = KEM.getInstance("SNTRUPrime");
        KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsSpec);
        SecretKey secR = d.decapsulate(em);

        // secS and secR will be identical
        assertEquals(secS.getAlgorithm(), secR.getAlgorithm());
        assertTrue(Arrays.areEqual(secS.getEncoded(), secR.getEncoded()));

    }
    public void testKEM()
            throws Exception
    {
        // Receiver side
        KeyPairGenerator g = KeyPairGenerator.getInstance("SNTRUPrime");

        g.initialize(SNTRUPrimeParameterSpec.sntrup653, new SecureRandom());

        KeyPair kp = g.generateKeyPair();
        PublicKey pkR = kp.getPublic();

        // Sender side
        KEM kemS = KEM.getInstance("SNTRUPrime"); //Should the name be "SNTRUPrime-KEM" ?
        KTSParameterSpec ktsSpec = null;
        KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsSpec, null);
        KEM.Encapsulated enc = e.encapsulate();
        SecretKey secS = enc.key();
        byte[] em = enc.encapsulation();
        byte[] params = enc.params();

        // Receiver side
        KEM kemR = KEM.getInstance("SNTRUPrime");
//        AlgorithmParameters algParams = AlgorithmParameters.getInstance("SNTRUPrime");
//        algParams.init(params);
//        SNTRUPrimeParameterSpec specR = algParams.getParameterSpec(SNTRUPrimeParameterSpec.class);
        KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsSpec);
        SecretKey secR = d.decapsulate(em);

        // secS and secR will be identical
        assertEquals(secS.getAlgorithm(), secR.getAlgorithm());
        assertTrue(Arrays.areEqual(secS.getEncoded(), secR.getEncoded()));
    }

    public void testBasicKEMAES()
            throws Exception
    {
        if (Security.getProvider(BouncyCastlePQCProvider.PROVIDER_NAME) == null)
        {
            Security.addProvider(new BouncyCastlePQCProvider());
        }
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SNTRUPrime", "BCPQC");
        kpg.initialize(SNTRUPrimeParameterSpec.sntrup653, new SecureRandom());

        performKEM(kpg.generateKeyPair(), new KEMParameterSpec("AES"));
        performKEM(kpg.generateKeyPair(),0, 16, "AES", new KEMParameterSpec("AES"));
        performKEM(kpg.generateKeyPair(), new KEMParameterSpec("AES-KWP"));

        try
        {
            performKEM(kpg.generateKeyPair(),0, 16, "AES-KWP", new KEMParameterSpec("AES"));
            fail();
        }
        catch (Exception ex)
        {
        }

        kpg.initialize(SNTRUPrimeParameterSpec.sntrup1013, new SecureRandom());
        performKEM(kpg.generateKeyPair(), new KEMParameterSpec("AES"));



    }

    public void testBasicKEMCamellia()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SNTRUPrime", "BCPQC");
        kpg.initialize(SNTRUPrimeParameterSpec.sntrup653, new SecureRandom());

        performKEM(kpg.generateKeyPair(), new KTSParameterSpec.Builder("Camellia", 256).build());
        performKEM(kpg.generateKeyPair(), new KTSParameterSpec.Builder("Camellia-KWP", 256).build());
    }

    public void testBasicKEMSEED()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SNTRUPrime", "BCPQC");
        kpg.initialize(SNTRUPrimeParameterSpec.sntrup653, new SecureRandom());

        performKEM(kpg.generateKeyPair(), new KTSParameterSpec.Builder("SEED", 128).build());
    }

    public void testBasicKEMARIA()
            throws Exception
    {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("SNTRUPrime", "BCPQC");
        kpg.initialize(SNTRUPrimeParameterSpec.sntrup653, new SecureRandom());

        performKEM(kpg.generateKeyPair(), new KEMParameterSpec("ARIA"));
        performKEM(kpg.generateKeyPair(), new KEMParameterSpec("ARIA-KWP"));
    }

    private void performKEM(KeyPair kp, int from, int to, String algorithm, KTSParameterSpec ktsParameterSpec)
            throws Exception
    {
        PublicKey pkR = kp.getPublic();

        // Sender side
        KEM kemS = KEM.getInstance("SNTRUPrime");
        KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsParameterSpec, null);
        KEM.Encapsulated enc = e.encapsulate(from, to, algorithm);
        SecretKey secS = enc.key();
        byte[] em = enc.encapsulation();

        // Receiver side
        KEM kemR = KEM.getInstance("SNTRUPrime");
        KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsParameterSpec);
        SecretKey secR = d.decapsulate(em, from, to, algorithm);

        // secS and secR will be identical
        assertEquals(secS.getAlgorithm(), secR.getAlgorithm());
        assertTrue(Arrays.areEqual(secS.getEncoded(), secR.getEncoded()));
    }

    private void performKEM(KeyPair kp, KTSParameterSpec ktsParameterSpec)
            throws Exception
    {
        PublicKey pkR = kp.getPublic();

        // Sender side
        KEM kemS = KEM.getInstance("SNTRUPrime");
        KEM.Encapsulator e = kemS.newEncapsulator(pkR, ktsParameterSpec, null);
        KEM.Encapsulated enc = e.encapsulate();
        SecretKey secS = enc.key();
        byte[] em = enc.encapsulation();

        // Receiver side
        KEM kemR = KEM.getInstance("SNTRUPrime");
//        KTSParameterSpec RktsParameterSpec = new KTSParameterSpec.Builder(
//                ktsParameterSpec.getKeyAlgorithmName(),
//                enc.key().getEncoded().length
//        ).withParameterSpec(ktsParameterSpec).build();
        KEM.Decapsulator d = kemR.newDecapsulator(kp.getPrivate(), ktsParameterSpec);
        SecretKey secR = d.decapsulate(em);

        // secS and secR will be identical
        assertEquals(secS.getAlgorithm(), secR.getAlgorithm());
        assertTrue(Arrays.areEqual(secS.getEncoded(), secR.getEncoded()));
    }
}
