import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

// The working one
public class RSAMain {


  public static final String PRIVATE_KEY_FILE = "private.key";
  public static final String PUBLIC_KEY_FILE = "public.key";

  /**
   * Generate key which contains a pair of private and public key using 1024
   * bytes. Store the set of keys in Private.key and Public.key files.
   */
  public static void generateKeys() {
    try {
      KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
      keyGen.initialize(1024);
      KeyPair key = keyGen.generateKeyPair();

      File privateKeyFile = new File(PRIVATE_KEY_FILE);
      File publicKeyFile = new File(PUBLIC_KEY_FILE);

      privateKeyFile.createNewFile();
      publicKeyFile.createNewFile();

      // Saving the Public key in a file
      ObjectOutputStream publicKeyOS = new ObjectOutputStream(
          new FileOutputStream(publicKeyFile));
      publicKeyOS.writeObject(key.getPublic());
      publicKeyOS.close();

      // Saving the Private key in a file
      ObjectOutputStream privateKeyOS = new ObjectOutputStream(
          new FileOutputStream(privateKeyFile));
      privateKeyOS.writeObject(key.getPrivate());
      privateKeyOS.close();
    } catch (Exception e) {
      e.printStackTrace();
    }

  }

  // The method checks if the pair of public and private key has been generated.
  public static boolean areKeysPresent() {

    File privateKey = new File(PRIVATE_KEY_FILE);
    File publicKey = new File(PUBLIC_KEY_FILE);

    if (privateKey.exists() && publicKey.exists()) {
      return true;
    }
    else return false;
  }

  
  public static byte[] encrypt(String text, PublicKey key) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException 
  {
    byte[] cipherText = null;
    // get an RSA cipher object and print the provider
    final Cipher cipher = Cipher.getInstance("RSA");
    // encrypt the plain text using the public key
    cipher.init(Cipher.ENCRYPT_MODE, key);
    cipherText = cipher.doFinal(text.getBytes());
    return cipherText;
  }

  
  public static String decrypt(byte[] text, PrivateKey key) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException 
  {
    byte[] decryptedText = null;
    // get an RSA cipher object and print the provider
    final Cipher cipher = Cipher.getInstance("RSA");

    // decrypt the text using the private key
    cipher.init(Cipher.DECRYPT_MODE, key);
    decryptedText = cipher.doFinal(text);

    return new String(decryptedText);
  }


  
  public static void main(String[] args) {

    try {

      // Check if the pair of keys are present else generate those.
      if (!areKeysPresent()) {
        generateKeys();
      }
      

      final String originalText = "Text to be encrypted ";
      ObjectInputStream inputStream = null;

      // Encrypt the string using the public key
      // Load public key from file
      inputStream = new ObjectInputStream(new FileInputStream(PUBLIC_KEY_FILE));
      final PublicKey publicKey = (PublicKey) inputStream.readObject();
      final byte[] cipherText = encrypt(originalText, publicKey);

      // Decrypt the cipher text using the private key.
      // Load private key from file
      inputStream = new ObjectInputStream(new FileInputStream(PRIVATE_KEY_FILE));
      final PrivateKey privateKey = (PrivateKey) inputStream.readObject();
      final String plainText = decrypt(cipherText, privateKey);

      // Printing the Original, Encrypted and Decrypted Text
      System.out.println("Original Text: " + originalText);
      System.out.println("Encrypted Text: " + cipherText.toString());
      System.out.println("Decrypted Text: " + plainText);

      printKeys(privateKey,publicKey);
      
      inputStream.close();
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  private static void printKeys(PrivateKey privateKey, PublicKey publicKey) throws InvalidKeySpecException, NoSuchAlgorithmException {
	  KeyFactory fact = KeyFactory.getInstance("RSA");
		RSAPublicKeySpec pub = fact.getKeySpec(publicKey,
		  RSAPublicKeySpec.class);
		RSAPrivateKeySpec priv = fact.getKeySpec(privateKey,
		  RSAPrivateKeySpec.class);
		System.out.println("Printing keys...");
		System.out.println("  Public exponent: " + pub.getPublicExponent().toString());
		System.out.println("  Public modulus: " + pub.getModulus().toString());
		System.out.println("  Private exponent: " + priv.getPrivateExponent().toString());
		System.out.println("  Private modulus: " + priv.getModulus().toString());
  }
}