/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package csc479_hw4_gui;

import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Formatter;
import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

/**
 *
 * @author ksiegfried
 */
public class PKE {

    static final String symmetricAlgorithm = "Blowfish";
    static final String asymmetricAlgorithm = "RSA";
    static final String pkInstance = "RSA/ECB/PKCS1Padding";
    static final int RSASIZE = 2048;
    static final int BFSIZE = 128;
    static final int CIPHERTEXTSIZE = 256;
    static final boolean append = true;
    public static final int GENERATION_SUCCESS = 0;
    public static final int GENERATION_FAILURE = 1;

    public static PrivateKey recoverPrivateKey(String fileName) throws Exception {
        try {
            BufferedReader bReader = new BufferedReader(new FileReader(fileName));
            byte[] privKeyBytes = hexStringToByteArray(bReader.readLine());
            PrivateKey privKey = KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(privKeyBytes));
            bReader.close();
            return privKey;
        } catch (Exception e) {
            e.printStackTrace();
            throw e;
        }

    }

    public static PublicKey recoverPublicKey(String fileName) throws Exception {
        try {
            BufferedReader bReader = new BufferedReader(new FileReader(fileName));
            byte[] pubKeyBytes = hexStringToByteArray(bReader.readLine());
            PublicKey pubKey = KeyFactory.getInstance("RSA").generatePublic(new X509EncodedKeySpec(pubKeyBytes));
            bReader.close();
            return pubKey;
        } catch (Exception e) {
            e.printStackTrace();
            throw e;
        }

    }

    private static void writeKey(String key, File f) throws IOException {
        PrintWriter pr = new PrintWriter(new FileOutputStream(f));
        pr.println(key);
        pr.flush();
        pr.close();
    }

    public static int generateKeyPair(String pathName) {
        File pubFile, privFile;
        PrivateKey privKey;
        PublicKey pubKey;
        try {
            KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
            kpg.initialize(2048);
            KeyPair keypair = kpg.generateKeyPair();

            pubKey = keypair.getPublic();
            System.out.println("Public Key Bytes: " + pubKey.getEncoded().length);
            privKey = keypair.getPrivate();
            System.out.println("Private Key Bytes: " + privKey.getEncoded().length);

            pubFile = new File(pathName, "MyRSAKeys.public");
            privFile = new File(pathName, "MyRSAKeys.private");

            writeKey(byteArray2Hex(pubKey.getEncoded()), pubFile);
            writeKey(byteArray2Hex(privKey.getEncoded()), privFile);
        } catch (Exception e) {
            System.out.println(e);
            return PKE.GENERATION_FAILURE;
        }

        return PKE.GENERATION_SUCCESS;
    }

    public static boolean encrypt(PublicKey pubKey, String fileName) {
        try {
            /*Set outfile*/
            File outFile = new File(fileName + ".encrypted");

            /*Initialize symmetric key crytpo*/
            SecretKey symmetricKey;
            KeyGenerator bkg = KeyGenerator.getInstance(symmetricAlgorithm);
            bkg.init(BFSIZE);
            symmetricKey = bkg.generateKey();
            byte[] secretBytes = symmetricKey.getEncoded();
            System.out.println("Secret Bytes Length: " + secretBytes.length);
            System.out.println("Secret Key Bytes: " + byteArray2Hex(secretBytes));
            Cipher BlowfishCipher = Cipher.getInstance(symmetricAlgorithm);
            BlowfishCipher.init(Cipher.ENCRYPT_MODE, symmetricKey);

            /*Set up RSA Cipher*/
            Cipher c = Cipher.getInstance(pkInstance);
            c.init(Cipher.ENCRYPT_MODE, pubKey);
            byte[] cipherText = c.doFinal(secretBytes);

            /*Save Encrypted Key*/
            FileOutputStream fOS = new FileOutputStream(outFile);
            fOS.write(cipherText);
            fOS.flush();
            fOS.close();

            /*Append encrypted data to key*/
            fOS = new FileOutputStream(outFile, append);
            FileInputStream fIS = new FileInputStream(fileName);
            CipherOutputStream cOS = new CipherOutputStream(fOS, BlowfishCipher);
            DataInputStream dIS = new DataInputStream(fIS);
            int count;
            while ((count = dIS.read()) != -1) {
                cOS.write(count);
                cOS.flush();
            }
            cOS.close();
            fOS.close();
            dIS.close();


        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }

    public static boolean decrypt(PrivateKey privKey, String fileName, String outFileName) {
        try {
            /*Set up public key cipher*/
            byte[] recoveredBytes = new byte[256];
            Cipher c = Cipher.getInstance(pkInstance);
            c.init(Cipher.DECRYPT_MODE, privKey);

            /*Recover symmetric key*/
            FileInputStream fIS = new FileInputStream(new File(fileName));
            System.out.println("Read bytes: " + fIS.read(recoveredBytes));
            System.out.println("Recovered Text: " + byteArray2Hex(recoveredBytes));
            recoveredBytes = c.doFinal(recoveredBytes);
            System.out.println("Recovered Bytes: " + byteArray2Hex(recoveredBytes));

            /*Initialize symmetric cipher for decryption*/
            SecretKeySpec BFKeySpec = new SecretKeySpec(recoveredBytes, "Blowfish");
            Cipher symmetricCipher = Cipher.getInstance("Blowfish");
            symmetricCipher.init(Cipher.DECRYPT_MODE, BFKeySpec);

            /*Skip past secret key block*/
            fIS = new FileInputStream(new File(fileName));
            fIS.skip(256);

            /*Read, and decrypt file. Store plaintext in output file*/
            CipherInputStream cIS = new CipherInputStream(fIS, symmetricCipher);
            FileOutputStream fOS = new FileOutputStream(new File(outFileName));
            int read;
            while ((read = cIS.read()) != -1) {
                fOS.write(read);
                fOS.flush();
            }
            fOS.close();
            cIS.close();


        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }

    private static String byteArray2Hex(byte[] hash) {
        Formatter formatter = new Formatter();
        for (byte b : hash) {
            formatter.format("%02x", b);
        }
        return formatter.toString();
    }

    private static byte[] hexStringToByteArray(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                    + Character.digit(s.charAt(i + 1), 16));
        }
        return data;
    }
}