// CSE 142 Exploration Session
// Modular arithmetic
// Gets a message from the user and encrypts/decrypts it using RSA.
// Here we're using industrial strength RSA and letting Java abstract
// away all the prime numbers, exponentiation, and modular arithmetic.

import java.security.*; // for KeyPairGenerator, KeyPair, Key
import java.util.*; // for Scanner, Arrays
import javax.crypto.*; // for Cipher

public class IndustrialRSA {
	public static final String ALGORITHM = "RSA";
	public static final int BITS = 2048;

	public static void main(String[] args) {
		KeyPairGenerator kpg = null;
		try {
			kpg = KeyPairGenerator.getInstance(ALGORITHM);
		} catch (Exception e) {
			bail();
		}
		kpg.initialize(BITS);
		KeyPair kp = kpg.genKeyPair();
		Key publicKey = kp.getPublic();
		Key privateKey = kp.getPrivate();
		
		Scanner console = new Scanner(System.in);
		System.out.print("Message to encrypt/decrypt: ");
		byte[] message = console.nextLine().getBytes();
		System.out.println("Message: " + new String(message));
		byte[] encrypted = encrypt(message, privateKey);
		System.out.println("Encrypted: " + new String(encrypted));
		byte[] decrypted = decrypt(encrypted, publicKey);
		System.out.println("Decrypted: " + new String(decrypted));
		System.out.println("Decrypted matches original: " + Arrays.equals(decrypted, message));
	}
	
	// Return data encrypted with key.
	public static byte[] encrypt(byte[] data, Key key) {
		return convert(data, key, Cipher.ENCRYPT_MODE);
	}
	
	// Return data decrypted with key.
	public static byte[] decrypt(byte[] data, Key key) {
		return convert(data, key, Cipher.DECRYPT_MODE);
	}
	
	// Use the passed key and mode to convert the data to a different format.
	public static byte[] convert(byte[] data, Key key, int mode) {
		try {
			Cipher cipher = Cipher.getInstance(ALGORITHM);
			cipher.init(mode, key);
			return cipher.doFinal(data);
		} catch (Exception e) {
			bail();
		}
		return new byte[0]; // We'll never actually get here.
	}
	
	// Print error message and exit program.
	public static void bail() {
		System.out.println("Something bad happened ):");
		System.exit(1);
	}
}