package rsa;

import java.io.*;
import java.security.*;
import javax.crypto.*;

/**
 * Program wykorzystujący algorytm szyfrowania RSA.
 * Uruchamianie:
 * java RSATest -genkey public private
 * java RSATest -encrypt plaintext encrypted public
 * java RSATest -decrypt encrypted decrypted private
 * @author Cay Horstmann
 * @version 1.02 2018-05-01
 */
public class RSATest
{
   private static final int KEYSIZE = 512;

   public static void main(String[] args) 
      throws IOException, GeneralSecurityException, ClassNotFoundException
   {
      if (args[0].equals("-genkey"))
      {
         KeyPairGenerator pairgen = KeyPairGenerator.getInstance("RSA");
         var random = new SecureRandom();
         pairgen.initialize(KEYSIZE, random);
         KeyPair keyPair = pairgen.generateKeyPair();
         try (var out = new ObjectOutputStream(new FileOutputStream(args[1])))
         {
            out.writeObject(keyPair.getPublic());
         }
         try (var out = new ObjectOutputStream(new FileOutputStream(args[2])))
         {
            out.writeObject(keyPair.getPrivate());
         }
      }
      else if (args[0].equals("-encrypt"))
      {
         KeyGenerator keygen = KeyGenerator.getInstance("AES");
         var random = new SecureRandom();
         keygen.init(random);
         SecretKey key = keygen.generateKey();

         // Szyfruje klucz AES za pomocą klucza publicznego RSA
         try (var keyIn = new ObjectInputStream(new FileInputStream(args[3]));
               var out = new DataOutputStream(new FileOutputStream(args[2]));
               var in = new FileInputStream(args[1]) )
         {
            var publicKey = (Key) keyIn.readObject();
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.WRAP_MODE, publicKey);
            byte[] wrappedKey = cipher.wrap(key);
            out.writeInt(wrappedKey.length);
            out.write(wrappedKey);
         
            cipher = Cipher.getInstance("AES");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            Util.crypt(in, out, cipher);            
         }         
      }
      else
      {
         try (var in = new DataInputStream(new FileInputStream(args[1]));
               var keyIn = new ObjectInputStream(new FileInputStream(args[3]));
               var out = new FileOutputStream(args[2]))
         {
            int length = in.readInt();
            var wrappedKey = new byte[length];
            in.read(wrappedKey, 0, length);

            // Odszyfrowuje za pomocą klucza prywatnego RSA
            var privateKey = (Key) keyIn.readObject();
   
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.UNWRAP_MODE, privateKey);
            Key key = cipher.unwrap(wrappedKey, "AES", Cipher.SECRET_KEY);
   
            cipher = Cipher.getInstance("AES");
            cipher.init(Cipher.DECRYPT_MODE, key);
   
            Util.crypt(in, out, cipher);
         }
      }
   }
}
