import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.KeyGenerator;

public class TestCipherOutputStream {

    public static byte[] encryptCipherOutputStream(Key key, byte[] input) throws GeneralSecurityException, IOException {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();

        Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");

        cipher.init(Cipher.ENCRYPT_MODE, key);
        CipherOutputStream cos = new CipherOutputStream(bos, cipher);
        cos.write(input);
        // cos.close should throw an IOException on BadPaddingException
        cos.close();
        return bos.toByteArray();

    }

    public static byte[] encryptCipherInputStream(Key key, byte[] input) throws GeneralSecurityException, IOException {
        Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, key);
        CipherInputStream cis = new CipherInputStream(new ByteArrayInputStream(input), cipher);
        byte[] outBuffer = new byte[1024];
        int bytesRead = cis.read(outBuffer);
        cis.close();
        byte[] ret = new byte[bytesRead];
        System.arraycopy(outBuffer, 0, ret, 0, bytesRead);
        return ret;

    }

    public static void main(String[] args) throws Exception {
        KeyGenerator kgen = KeyGenerator.getInstance("AES");
        Key key = kgen.generateKey();

        byte[] toEncryptOk = new byte[16];

        System.out.println("Cipher of cipher.getBlockSize byte will pass");
        byte[] encOk1 = encryptCipherOutputStream(key, toEncryptOk);
        byte[] encOk2 = encryptCipherInputStream(key, toEncryptOk);
        assert (Arrays.equals(encOk1, encOk2));

        // Problematic behavior
        byte[] toEncryptShouldAlwaysFailed = new byte[15];

        byte[] encKo1 = null, encKo2 = null;
        try {
            System.out.println("Cipher of less than cipher.getBlockSize should fail...");
            encKo1 = encryptCipherOutputStream(key, toEncryptShouldAlwaysFailed);
            System.out.println("Drop silently IllegalBlockSizeException");
        } catch (IOException ex) {
            System.out.println("Never append");
        }
        try {
            System.out.println("Cipher of less than cipher.getBlockSize will fail...");
            encKo2 = encryptCipherInputStream(key, toEncryptShouldAlwaysFailed);
            System.out.println("Never append");
        } catch (IOException ex) {
            System.out.println("and fail... Different way of handling IllegalBlockSizeException and BadPaddingException");
        }

        assert (Arrays.equals(encKo1, encKo2));
    }
}