package org.bouncycastle.crypto.test;

import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.digests.CSHAKEDigest;
import org.bouncycastle.crypto.digests.SHAKEDigest;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Strings;
import org.bouncycastle.util.encoders.Hex;

/**
 * CSHAKE test vectors from:
 *
 * https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Standards-and-Guidelines/documents/examples/cSHAKE_samples.pdf
 */
public class CSHAKETest
    extends DigestTest
{
    private static String[] messages =
    {
        "",
        "a",
        "abc",
        "abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"
    };

    private static String[] digests =
    {
        "28a0e58d342a45ef1522d69cd41748beee29c188364df18c84ed4ab8bed0cc85",
        "cb0a67f6ff4a3b64497a757a85bda4a275cf11970ff226abd2bf2bbcba5890ea",
        "346c136daea11c436d8d9e668e08888bd4e341dae05da4cb8773f74402c5bdbc",
        "64602dc88c880bdfb6d0c9163a72b2e3653ab6114e4f4e25d7aaf5b8d441e36f"
    };

    public CSHAKETest()
    {
        super(new CSHAKEDigest(128, new byte[1], new byte[1]), messages, digests);
    }

    public String getName()
    {
        return "CSHAKE";
    }

    protected Digest cloneDigest(Digest digest)
    {
        return new CSHAKEDigest((CSHAKEDigest)digest);
    }

    protected Digest cloneDigest(byte[] encodedState)
    {
        return new CSHAKEDigest(encodedState);
    }

    public void performTest()
    {
        super.performTest();

        CSHAKEDigest cshake = new CSHAKEDigest(128, new byte[0], Strings.toByteArray("Email Signature"));

        cshake.update(Hex.decode("00010203"), 0, 4);

        byte[] res = new byte[32];

        cshake.doOutput(res, 0, res.length);

        isTrue("oops!", Arrays.areEqual(Hex.decode("c1c36925b6409a04f1b504fcbca9d82b4017277cb5ed2b2065fc1d3814d5aaf5"), res));

        cshake = new CSHAKEDigest(128, new byte[0], Strings.toByteArray("Email Signature"));

        cshake.update(Hex.decode(
            "000102030405060708090A0B0C0D0E0F" +
                "101112131415161718191A1B1C1D1E1F" +
                "202122232425262728292A2B2C2D2E2F" +
                "303132333435363738393A3B3C3D3E3F" +
                "404142434445464748494A4B4C4D4E4F" +
                "505152535455565758595A5B5C5D5E5F" +
                "606162636465666768696A6B6C6D6E6F" +
                "707172737475767778797A7B7C7D7E7F" +
                "808182838485868788898A8B8C8D8E8F" +
                "909192939495969798999A9B9C9D9E9F" +
                "A0A1A2A3A4A5A6A7A8A9AAABACADAEAF" +
                "B0B1B2B3B4B5B6B7B8B9BABBBCBDBEBF" +
                "C0C1C2C3C4C5C6C7"), 0, 1600 / 8);

        res = new byte[32];

        cshake.doOutput(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode("C5221D50E4F822D96A2E8881A961420F294B7B24FE3D2094BAED2C6524CC166B "), res));

        cshake = new CSHAKEDigest(256, new byte[0], Strings.toByteArray("Email Signature"));

        cshake.update(Hex.decode("00010203"), 0, 4);

        res = new byte[64];

        cshake.doOutput(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode(
        "D008828E2B80AC9D2218FFEE1D070C48"+
            "B8E4C87BFF32C9699D5B6896EEE0EDD1"+
            "64020E2BE0560858D9C00C037E34A969"+
            "37C561A74C412BB4C746469527281C8C"),res));

        cshake = new CSHAKEDigest(256, new byte[0], Strings.toByteArray("Email Signature"));

        cshake.update(Hex.decode(
            "000102030405060708090A0B0C0D0E0F" +
                "101112131415161718191A1B1C1D1E1F" +
                "202122232425262728292A2B2C2D2E2F" +
                "303132333435363738393A3B3C3D3E3F" +
                "404142434445464748494A4B4C4D4E4F" +
                "505152535455565758595A5B5C5D5E5F" +
                "606162636465666768696A6B6C6D6E6F" +
                "707172737475767778797A7B7C7D7E7F" +
                "808182838485868788898A8B8C8D8E8F" +
                "909192939495969798999A9B9C9D9E9F" +
                "A0A1A2A3A4A5A6A7A8A9AAABACADAEAF" +
                "B0B1B2B3B4B5B6B7B8B9BABBBCBDBEBF" +
                "C0C1C2C3C4C5C6C7"), 0, 1600 / 8);

        res = new byte[64];

        cshake.doOutput(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode(
                "07DC27B11E51FBAC75BC7B3C1D983E8B"+
                    "4B85FB1DEFAF218912AC864302730917"+
                    "27F42B17ED1DF63E8EC118F04B23633C"+
                    "1DFB1574C8FB55CB45DA8E25AFB092BB"), res));

        doFinalTest();
        longBlockTest();
        checkZeroPadZ();
        
        checkSHAKE(128, new CSHAKEDigest(128, new byte[0], new byte[0]), Hex.decode("eeaabeef"));
        checkSHAKE(256, new CSHAKEDigest(256, new byte[0], null), Hex.decode("eeaabeef"));
        checkSHAKE(128, new CSHAKEDigest(128, null, new byte[0]), Hex.decode("eeaabeef"));
        checkSHAKE(128, new CSHAKEDigest(128, null, null), Hex.decode("eeaabeef"));
        checkSHAKE(256, new CSHAKEDigest(256, null, null), Hex.decode("eeaabeef"));
    }

    private void checkZeroPadZ()
    {
        byte[] buf = new byte[20];

        CSHAKEDigest cshake1 = new CSHAKEDigest(256, new byte[0], new byte[265]);
        cshake1.doOutput(buf, 0, buf.length);
        isTrue(areEqual(Hex.decode("6e393540387004f087c4180db008acf6825190cf"), buf));

        CSHAKEDigest cshake2 = new CSHAKEDigest(128, new byte[0], new byte[329]);
        cshake2.doOutput(buf, 0, buf.length);
        isTrue(areEqual(Hex.decode("309bd7c285fcf8b839c9686b2cc00bd578947bee"), buf));

        cshake2 = new CSHAKEDigest(128, new byte[29], new byte[300]);
        cshake2.doOutput(buf, 0, buf.length);
        isTrue(areEqual(Hex.decode("ff6aafd83b8d22fc3e2e9b9948b581967ed9c5e7"), buf));
    }
    
    private void doFinalTest()
    {
        CSHAKEDigest cshake = new CSHAKEDigest(128, new byte[0], Strings.toByteArray("Email Signature"));

        cshake.update(Hex.decode("00010203"), 0, 4);

        byte[] res = new byte[32];

        cshake.doOutput(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode("c1c36925b6409a04f1b504fcbca9d82b4017277cb5ed2b2065fc1d3814d5aaf5"), res));

        cshake.doOutput(res, 0, res.length);

        isTrue(!Arrays.areEqual(Hex.decode("c1c36925b6409a04f1b504fcbca9d82b4017277cb5ed2b2065fc1d3814d5aaf5"), res));

        cshake.doFinal(res, 0, res.length);

        cshake.update(Hex.decode("00010203"), 0, 4);

        cshake.doFinal(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode("c1c36925b6409a04f1b504fcbca9d82b4017277cb5ed2b2065fc1d3814d5aaf5"), res));

        cshake.update(Hex.decode("00010203"), 0, 4);

        cshake.doOutput(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode("c1c36925b6409a04f1b504fcbca9d82b4017277cb5ed2b2065fc1d3814d5aaf5"), res));
        
        cshake.doFinal(res, 0, res.length);

        isTrue(Arrays.areEqual(Hex.decode("9cbce830079c452abdeb875366a49ebfe75b89ef17396e34898e904830b0e136"), res));
    }

    private void longBlockTest()
    {
        byte[] data = new byte[16000];
        byte[] res = new byte[32];

        for (int i = 0; i != data.length; i++)
        {
            data[i] = (byte)i;
        }

        byte[] hex0123 = Hex.decodeStrict("00010203");
        for (int i = 10000; i != data.length; i++)
        {
            CSHAKEDigest cshake = new CSHAKEDigest(128, new byte[0], Arrays.copyOfRange(data, 0, i));

            cshake.update(hex0123, 0, 4);

            cshake.doFinal(res, 0, 16);
        }

        CSHAKEDigest cshake = new CSHAKEDigest(256, new byte[0], new byte[200]);

        cshake.update(data, 0, 200);

        cshake.doFinal(res, 0, 32);

        isTrue(Arrays.areEqual(Hex.decode("4a899b5be460d85a9789215bc17f88b8f8ac049bd3b519f561e7b5d3870dafa3"), res));
    }

    private void checkSHAKE(int bitSize, CSHAKEDigest cshake, byte[] msg)
    {
        SHAKEDigest ref = new SHAKEDigest(bitSize);

        ref.update(msg, 0, msg.length);
        cshake.update(msg, 0, msg.length);

        byte[] res1 = new byte[32];
        byte[] res2 = new byte[32];

        ref.doFinal(res1, 0, res1.length);
        cshake.doFinal(res2, 0, res2.length);

        isTrue(Arrays.areEqual(res1, res2));
    }
    public static void main(
        String[] args)
    {
        runTest(new CSHAKETest());
    }
}
