package cz.muni.fi.pv079.exercises01;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.Key;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AESEncrypter {

	private static void encrypt(Cipher cipher, InputStream in, OutputStream out) {

		byte[] buf = new byte[1024];

		try {
			// Bytes written to out will be encrypted
			out = new CipherOutputStream(out, cipher);

			// Read in the cleartext bytes and write to out to encrypt
			int numRead = 0;
			while ((numRead = in.read(buf)) >= 0) {
				out.write(buf, 0, numRead);
			}
			out.close();
		} catch (IOException e) {
		}
	}

	private static void decrypt(Cipher cipher, InputStream in, OutputStream out) {

		byte[] buf = new byte[1024];

		try {
			// Bytes read from in will be decrypted
			in = new CipherInputStream(in, cipher);

			// Read in the decrypted bytes and write the cleartext to out
			int numRead = 0;
			while ((numRead = in.read(buf)) >= 0) {
				out.write(buf, 0, numRead);
			}
			out.close();
		} catch (IOException e) {
		}
	}

        private static void aescbc(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher ecipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
	  ecipher.init(Cipher.ENCRYPT_MODE, key, paramSpec);
	  encrypt(ecipher, new FileInputStream("lena.bmp"), new FileOutputStream("lena.aescbc"));
        }
        
        private static void aescfb(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher ecipher = Cipher.getInstance("AES/CFB64/PKCS5Padding"); // CFB64 set feedback size on 64 bits
	  ecipher.init(Cipher.ENCRYPT_MODE, key, paramSpec);
	  encrypt(ecipher, new FileInputStream("lena.bmp"), new FileOutputStream("lena.aescfb"));
        }
        
        private static void aesofb(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher ecipher = Cipher.getInstance("AES/OFB64/PKCS5Padding"); // CFB64 set feedback size on 64 bits
	  ecipher.init(Cipher.ENCRYPT_MODE, key, paramSpec);
	  encrypt(ecipher, new FileInputStream("lena.bmp"), new FileOutputStream("lena.aesofb"));
        }
        
        private static void aesctr(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher ecipher = Cipher.getInstance("AES/CTR/PKCS5Padding"); // CFB64 set feedback size on 64 bits
	  ecipher.init(Cipher.ENCRYPT_MODE, key, paramSpec);
	  encrypt(ecipher, new FileInputStream("lena.bmp"), new FileOutputStream("lena.aesctr"));
        }
        
        private static void daescbc(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher dcipher = Cipher.getInstance("AES/cbc/PKCS5Padding");
	  dcipher.init(Cipher.DECRYPT_MODE, key, paramSpec);
	  decrypt(dcipher, new FileInputStream("lena.aescbc"), new FileOutputStream("lena-cbc.bmp"));
        }
        
        private static void daescfb(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher dcipher = Cipher.getInstance("AES/CFB64/PKCS5Padding");
	  dcipher.init(Cipher.DECRYPT_MODE, key, paramSpec);
	  decrypt(dcipher, new FileInputStream("lena.aescfb"), new FileOutputStream("lena-cfb.bmp"));
        }
        
        private static void daesofb(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher dcipher = Cipher.getInstance("AES/OFB64/PKCS5Padding");
	  dcipher.init(Cipher.DECRYPT_MODE, key, paramSpec);
	  decrypt(dcipher, new FileInputStream("lena.aesofb"), new FileOutputStream("lena-ofb.bmp"));
        }
        
        private static void daesctr(AlgorithmParameterSpec paramSpec, Key key) throws Exception {
          Cipher dcipher = Cipher.getInstance("AES/CTR/PKCS5Padding");
	  dcipher.init(Cipher.DECRYPT_MODE, key, paramSpec);
	  decrypt(dcipher, new FileInputStream("lena.aesctr"), new FileOutputStream("lena-ctr.bmp"));
        }
        
        
	public static void main(String[] args) throws Exception {
		
		// Check parameter
		if (args.length != 1 || !(args[0].equals("e") || args[0].equals("d"))) {
			System.err.println("Run application with one parameter: -e (encryption), -d (decryption)");
			System.exit(2);
		}

		// Check file
		File file = new File("lena.bmp");
		if (!file.exists() || !file.canRead()) {
			System.err.println("File 'lena.bmp' not found!");
			System.exit(1);
		}

		// Create an 16-byte initialization vector
		byte[] iv = new byte[] { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
		AlgorithmParameterSpec paramSpec = new IvParameterSpec(iv);

		// Set a secret key
		// Using UTF-8 encoding
		byte[] keyValue = ("PV07900000000000000000000").getBytes("UTF-8");
		keyValue = Arrays.copyOf(keyValue, 16); // use only first 128 bit (16 * 8 bits = 128 bits)
		Key key = new SecretKeySpec(keyValue, "AES");

		// Encryption
                // You might want to note that Java calls PKCS#7-padding "PKCS5Padding"
		if (args[0].equals("e")) {

                        aescbc(paramSpec, key);
                        aescfb(paramSpec, key);
			aesofb(paramSpec, key);
                        aesctr(paramSpec, key);
			
		} else {
			
			// Check if are created encrypted files
			file = new File("lena.aescbc");
			if (!file.exists() || !file.canRead()) {
				System.err.println("File 'lena.aescbc' not found!");
				System.exit(1);
			}
			file = new File("lena.aescfb");
			if (!file.exists() || !file.canRead()) {
				System.err.println("File 'lena.aescfb' not found!");
				System.exit(1);
			}
			file = new File("lena.aesofb");
			if (!file.exists() || !file.canRead()) {
				System.err.println("File 'lena.aesofb' not found!");
				System.exit(1);
			}
			file = new File("lena.aesctr");
			if (!file.exists() || !file.canRead()) {
				System.err.println("File 'lena.aesctr' not found!");
				System.exit(1);
			}
			
			daescbc(paramSpec, key);
                        daescfb(paramSpec, key);
			daesofb(paramSpec, key);
                        daesctr(paramSpec, key);
		}			
	}
}