package org.bouncycastle.crypto.test;

import java.security.SecureRandom;

import org.bouncycastle.crypto.BlockCipher;
import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.engines.DSTU7624Engine;
import org.bouncycastle.crypto.modes.KXTSBlockCipher;
import org.bouncycastle.crypto.modes.NISTCTSBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.SimpleTest;

/**
 * CTS tester
 */
public class NISTCTSTest
    extends SimpleTest
{
    private static KeyParameter key = new KeyParameter(Hex.decode("000102030405060708090a0b0c0d0e0f"));
    private static byte[] iv = Hex.decode("101112131415161718191a1b1c1d1e1f");

    private static byte[] singleBlock = Hex.decode("4920616d206f6e6520626c6f636b2e2e");
    private static byte[] singleOut = Hex.decode("8aad2098847a2d74ac87de22745d2537");

    private static byte[] twoBlock = Hex.decode("4920616d206174206c656173742074776f20626c6f636b73206c6f6e672e2e2e");

    private static byte[] cs1TwoBlockOut = Hex.decode("3f07fd5816c3b96349eb9f6a074909d67237eb8aa9a7467b8a388c61d0e8f35a");
    private static byte[] cs2TwoBlockOut = Hex.decode("3f07fd5816c3b96349eb9f6a074909d67237eb8aa9a7467b8a388c61d0e8f35a");
    private static byte[] cs3TwoBlockOut = Hex.decode("7237eb8aa9a7467b8a388c61d0e8f35a3f07fd5816c3b96349eb9f6a074909d6");

    private static byte[] notQuiteTwo = Hex.decode("4920616d206e6f742071756974652074776f2e2e2e");

    private static byte[] cs1NotQuiteTwoBlockOut = Hex.decode("22ecf2ac77f098097ca69b72e3a46e9ca21bb5ebbc");
    private static byte[] cs2NotQuiteTwoBlockOut = Hex.decode("f098097ca69b72e3a46e9ca21bb5ebbc22ecf2ac77");
    private static byte[] cs3NotQuiteTwoBlockOut = Hex.decode("f098097ca69b72e3a46e9ca21bb5ebbc22ecf2ac77");

    static byte[] in1 = Hex.decode("4e6f7720697320746865207420");
    static byte[] in2 = Hex.decode("000102030405060708090a0b0c0d0e0fff0102030405060708090a0b0c0d0e0f0aaa");
    static byte[] out1 = Hex.decode("9952f131588465033fa40e8a98");
    static byte[] out2 = Hex.decode("358f84d01eb42988dc34efb994");
    static byte[] out3 = Hex.decode("170171cfad3f04530c509b0c1f0be0aefbd45a8e3755a873bff5ea198504b71683c6");

    private void testCTS(
        int id,
        int type,
        BlockCipher cipher,
        CipherParameters params,
        byte[] input,
        byte[] output)
        throws Exception
    {
        byte[] out = new byte[input.length];
        BufferedBlockCipher engine = new NISTCTSBlockCipher(type, cipher);

        engine.init(true, params);

        int len = engine.processBytes(input, 0, input.length, out, 0);

        engine.doFinal(out, len);

        if (!areEqual(output, out))
        {
            fail(id + " failed encryption expected " + new String(Hex.encode(output)) + " got " + new String(Hex.encode(out)));
        }

        engine.init(false, params);

        len = engine.processBytes(output, 0, output.length, out, 0);

        engine.doFinal(out, len);

        if (!areEqual(input, out))
        {
            fail(id + " failed decryption expected " + new String(Hex.encode(input)) + " got " + new String(Hex.encode(out)));
        }
    }

    private void testExceptions()
        throws InvalidCipherTextException
    {
        BufferedBlockCipher engine = new NISTCTSBlockCipher(NISTCTSBlockCipher.CS1, AESEngine.newInstance());
        CipherParameters params = new KeyParameter(new byte[engine.getBlockSize()]);
        engine.init(true, params);

        byte[] out = new byte[engine.getOutputSize(engine.getBlockSize())];

        engine.processBytes(new byte[engine.getBlockSize() - 1], 0, engine.getBlockSize() - 1, out, 0);
        try
        {
            engine.doFinal(out, 0);
            fail("Expected CTS encrypt error on < 1 block input");
        }
        catch (DataLengthException e)
        {
            // Expected
        }

        engine.init(true, params);
        engine.processBytes(new byte[engine.getBlockSize()], 0, engine.getBlockSize(), out, 0);
        try
        {
            engine.doFinal(out, 0);
        }
        catch (DataLengthException e)
        {
            fail("Unexpected CTS encrypt error on == 1 block input");
        }

        engine.init(false, params);
        engine.processBytes(new byte[engine.getBlockSize() - 1], 0, engine.getBlockSize() - 1, out, 0);
        try
        {
            engine.doFinal(out, 0);
            fail("Expected CTS decrypt error on < 1 block input");
        }
        catch (DataLengthException e)
        {
            // Expected
        }

        engine.init(false, params);
        engine.processBytes(new byte[engine.getBlockSize()], 0, engine.getBlockSize(), out, 0);
        try
        {
            engine.doFinal(out, 0);
        }
        catch (DataLengthException e)
        {
            fail("Unexpected CTS decrypt error on == 1 block input");
        }

    }

    private void testOverlapping()
        throws Exception
    {
        SecureRandom random = new SecureRandom();
        byte[] keyBytes = new byte[16];
        byte[] iv = new byte[16];
        random.nextBytes(keyBytes);
        BufferedBlockCipher bc = new NISTCTSBlockCipher(NISTCTSBlockCipher.CS1, AESEngine.newInstance());
        ParametersWithIV param = new ParametersWithIV(new KeyParameter(keyBytes), iv);

        int offset = 1 + random.nextInt(bc.getBlockSize() - 1) + bc.getBlockSize();
        byte[] data = new byte[bc.getBlockSize() * 4 + offset];
        byte[] expected = new byte[bc.getOutputSize(bc.getBlockSize() * 3)];
        random.nextBytes(data);

        bc.init(true, param);
        int len = bc.processBytes(data, 0, expected.length, expected, 0);
        bc.doFinal(expected, len);
        bc.init(true, param);
        len = bc.processBytes(data, 0, expected.length, data, offset);
        bc.doFinal(data, offset + len);

        if (!areEqual(expected, Arrays.copyOfRange(data, offset, offset + expected.length)))
        {
            fail("failed for overlapping encryption");
        }

        bc.init(false, param);
        bc.processBytes(data, 0, expected.length, expected, 0);
        bc.init(false, param);
        bc.processBytes(data, 0, expected.length, data, offset);

        if (!areEqual(expected, Arrays.copyOfRange(data, offset, offset + expected.length)))
        {
            fail("failed for overlapping decryption");
        }
    }

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

    public void performTest()
        throws Exception
    {
        testCTS(1, NISTCTSBlockCipher.CS1, AESEngine.newInstance(), new ParametersWithIV(key, iv), singleBlock, singleOut);
        testCTS(2, NISTCTSBlockCipher.CS2, AESEngine.newInstance(), new ParametersWithIV(key, iv), singleBlock, singleOut);
        testCTS(3, NISTCTSBlockCipher.CS3, AESEngine.newInstance(), new ParametersWithIV(key, iv), singleBlock, singleOut);

        testCTS(4, NISTCTSBlockCipher.CS1, AESEngine.newInstance(), new ParametersWithIV(key, iv), twoBlock, cs1TwoBlockOut);
        testCTS(5, NISTCTSBlockCipher.CS2, AESEngine.newInstance(), new ParametersWithIV(key, iv), twoBlock, cs2TwoBlockOut);
        testCTS(6, NISTCTSBlockCipher.CS3, AESEngine.newInstance(), new ParametersWithIV(key, iv), twoBlock, cs3TwoBlockOut);

        testCTS(7, NISTCTSBlockCipher.CS1, AESEngine.newInstance(), new ParametersWithIV(key, iv), notQuiteTwo, cs1NotQuiteTwoBlockOut);
        testCTS(8, NISTCTSBlockCipher.CS2, AESEngine.newInstance(), new ParametersWithIV(key, iv), notQuiteTwo, cs2NotQuiteTwoBlockOut);
        testCTS(9, NISTCTSBlockCipher.CS3, AESEngine.newInstance(), new ParametersWithIV(key, iv), notQuiteTwo, cs3NotQuiteTwoBlockOut);

        byte[] aes128b = Hex.decode("aafd12f659cae63489b479e5076ddec2f06cb58faafd12f6");
        byte[] aesIn1b = Hex.decode("000102030405060708090a0b0c0d0e0fff0102030405060708090a0b0c0d0e0f");
        byte[] aesOut1b = Hex.decode("6db2f802d99e1ef0a5940f306079e083cf87f4d8bb9d1abb36cdd9f44ead7d04");

        testCTS(10, NISTCTSBlockCipher.CS3, AESEngine.newInstance(), new ParametersWithIV(new KeyParameter(aes128b), Hex.decode("aafd12f659cae63489b479e5076ddec2")), aesIn1b, aesOut1b);

        byte[] aes128c = Hex.decode("aafd12f659cae63489b479e5076ddec2");
        byte[] aesOut1c = Hex.decode("0af33c005a337af55a5149effc5108eaa1ea87de8a8556e8786b8f230da64e56");

        testCTS(11, NISTCTSBlockCipher.CS3, AESEngine.newInstance(), new ParametersWithIV(new KeyParameter(aes128c), Hex.decode("aafd12f659cae63489b479e5076ddec2")), aesIn1b, aesOut1c);

        testExceptions();
        testOverlapping();
    }

    public static void main(
        String[] args)
    {
        runTest(new NISTCTSTest());
    }
}
