// FW_RSAapp
// Benne de Weger - TU/e - February-April 2021

// FAECTOR Workshop Cryptographic Programming
// Assignment 9, 10 - application

// desired functionalities:
// - generate a key pair (always 2048 bits, e = Fermat-4)
// Assignment 9:
// - read a file and a public key, encrypt, write the encrypted file and the encrypted session key
// - read a file, a private key and an encrypted session key, decrypt, write the decrypted file
// Assignment 10:
// - read a file and a private key, sign, write the signature to a file
// - read a file, a private key and a signature, verify the signature

// command line options:
// -g id: generate a key pair with identity id
// -k fn: read key from file with name fn
//        (this can be used for both symmetric and asymmetric key files)
// -z fn: read signature from file with name fn
// -e fn: encrypt file fn to fn.enc
// -d fn: decrypt file fn to fn.dec
// -s fn: sign file fn to fn.sig
// -v fn: verify fn.sig on file fn

// usage:
// to generate a key:
//     java -cp .;crypto.jar FW_RSAapp -g testkey
//         (this produces two files: testkey_pub.txt, testkey_pri.txt)
// to encrypt:
//     java -cp .;crypto.jar FW_RSAapp -k testkey_pub.txt -e example.txt
//         (this produces two files: example.txt.enc, example.txt.key)
// to decrypt:
//     java -cp .;crypto.jar FW_RSAapp -k testkey_pri.txt -k example.txt.key -d example.txt.enc
//         (this produces one file: example.txt.enc.dec)
// to sign:
//     java -cp .;crypto.jar FW_RSAapp -k testkey_pri.txt -s example.txt
//         (this produces one file: example.txt.sig)
// to verify:
//     java -cp .;crypto.jar FW_RSAapp -k testkey_pub.txt -z example.txt.sig -v example.txt
//         (this produces no file)
// make sure the -k, -z parameters come before the -e, -d, -s, -v parameter
// when wrong keys are presented, strange error messages may result

import net.deweger.crypto.*;
import java.nio.file.*;

public class FW_RSAapp
{
    // read data from a file to a byte array
    private static byte[] readFile(String fileName) throws Exception
    {
        byte[] data = Files.readAllBytes(Paths.get(fileName));
        return data;
    }

    // write data to a file from a byte array
    private static void writeFile(String fileName, byte[] data) throws Exception
    {
        Files.write(Paths.get(fileName), data);
    }

    public static void main(String[] args)
    {
        try
        {
            FW_RSA rsa = new FW_RSA();
            AES aes = new AES();
            byte[] IV0 = new byte[16];
            byte[] encryptedSessionKey = new byte[0];
            byte[] signature = new byte[0];

            String identity = "";
            String fileName = "";

            // parse arguments and perform actions
            for (int i = 0; i < args.length; i++)
            {
                if (args[i].equals("-g") && i < args.length-1)
                {
                    // generate key pair
                    i++;
                    identity = args[i];
                    rsa.setId(identity);
                    rsa.generateRSAKeyPair(2048, true);
                    rsa.writeRSAPublicKey(identity + "_pub.txt");
                    rsa.writeRSAPrivateKey(identity + "_pri.txt");
                }
                else
                if (args[i].equals("-k") && i < args.length-1)
                {
                    // read key
                    i++;
                    fileName = args[i];
                    if (fileName.endsWith("_pub.txt"))
                        rsa.readRSAPublicKey(fileName);
                    if (fileName.endsWith("_pri.txt"))
                        rsa.readRSAPrivateKey(fileName);
                    if (fileName.endsWith(".key"))
                        encryptedSessionKey = readFile(fileName);
                }
                else
                if (args[i].equals("-z") && i < args.length-1)
                {
                    // read signature
                    i++;
                    fileName = args[i];
                    if (fileName.endsWith(".sig"))
                        signature = readFile(fileName);
                }
                else
                if (args[i].equals("-e") && i < args.length-1)
                {
                    // encrypt
                    i++;
                    fileName = args[i];
                    // read file with plaintext
                    byte[] plaintext = readFile(fileName);
                    // generate session key
                    byte[] sessionKey = aes.generateKey();
                    aes.setIv(IV0);
                    aes.setMode(AES.MODE_CTR);
                    aes.setPadding(AES.PAD_PKCS7);
                    // encrypt plaintext
                    byte[] ciphertext = aes.encrypt(plaintext);
                    // write encrypted data
                    writeFile(fileName + ".enc", ciphertext);
                    // encrypt session key
                    encryptedSessionKey = rsa.encryptKey(sessionKey);
                    // write encrypted session key
                    writeFile(fileName + ".key", encryptedSessionKey);
                }
                else
                if (args[i].equals("-d") && i < args.length-1)
                {
                    // decrypt
                    i++;
                    fileName = args[i];
                    // read file with ciphertext
                    byte[] ciphertext = readFile(fileName);
                    // decrypt session key
                    byte[] sessionKey = rsa.decryptKey(encryptedSessionKey);
                    aes.setIv(IV0);
                    aes.setMode(AES.MODE_CTR);
                    aes.setPadding(AES.PAD_PKCS7);
                    aes.setKey(sessionKey);
                    // decrypt ciphertext
                    byte[] plaintext = aes.decrypt(ciphertext);
                    // write decrypted ciphertext
                    writeFile(fileName + ".dec", plaintext);
                }
                else
                if (args[i].equals("-s") && i < args.length-1)
                {
                    // sign
                    i++;
                    fileName = args[i];
                    // read file with text
                    byte[] text = readFile(fileName);
                    // sign text
                    signature = rsa.generateSignature(text);
                    // write signature
                    writeFile(fileName + ".sig", signature);
                }
                else
                if (args[i].equals("-v") && i < args.length-1)
                {
                    // verify
                    i++;
                    fileName = args[i];
                    // read file with text
                    byte[] text = readFile(fileName);
                    // verify signature
                    boolean vf = rsa.verifySignature(text, signature);
                    if (vf)
                        System.out.println("signature is correct");
                    else
                        System.out.println("signature is incorrect");
                }
                else
                    throw new Exception("incorrect command line arguments");
            }
        }
        catch(Exception exc)
        {
            System.err.println("ERROR: FW_RSAapp: " + exc.getMessage());
        }
    }
}