package elte.java2_utikalauz5.crypto;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.ByteBuffer;
import java.security.KeyStore;
import java.security.KeyStore.PrivateKeyEntry;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

/**
Variációk rejtjelezésre.

@link.forrásfájl {@docRoot}/../data/crypto/src RSA.java
@link.letöltés {@docRoot}/../data/crypto RSA.jar
@since Java 2 Útikalauz programozóknak 5.0
*/
public class RSA {
    private static final byte[] IV = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08};

    private final FileInputStream   inFile;
    private final FileOutputStream  outFile;
    private final KeyStore.PrivateKeyEntry rsaKeyPair;
    private final boolean isEncrypt;

    public RSA(boolean isEncrypt, String keyId, String keyPassword, String keyStorePassword,
                       String inFileName, String outFileName) throws Exception {
        this.isEncrypt  = isEncrypt;

        this.inFile     = new FileInputStream(new File(inFileName));
        File outf   = new File(outFileName);
        outf.createNewFile();
        this.outFile    = new FileOutputStream(outf);

        KeyStore ks = KeyStore.getInstance(KeyStore.getDefaultType());
        ks.load(new FileInputStream(".keystore"), keyStorePassword.toCharArray());
        KeyStore.PasswordProtection keyPwd = new KeyStore.PasswordProtection(keyPassword.toCharArray());
        this.rsaKeyPair = (PrivateKeyEntry) ks.getEntry(keyId, keyPwd);
    }

    private void close() throws IOException {
        try {
            this.inFile.close();
        } finally {
            this.outFile.close();
        }
    }

    private void copy() throws Exception {
        if( this.isEncrypt ) {
            encrypt();
        } else {
            decrypt();
        }
    }

    private void encrypt() throws Exception {
        Cipher keyCipher = Cipher.getInstance("RSA");
        KeyGenerator keyGen = KeyGenerator.getInstance("Blowfish");
        SecretKey bulkKey = keyGen.generateKey();
        keyCipher.init(Cipher.WRAP_MODE,this.rsaKeyPair.getCertificate());
        byte[] wrappedBulkKey = keyCipher.wrap(bulkKey);
        ByteBuffer keyLength = ByteBuffer.allocate(4);
        keyLength.putInt(wrappedBulkKey.length);
        this.outFile.write(keyLength.array());
        this.outFile.write(wrappedBulkKey);

        System.err.println("Key type: " + bulkKey.getAlgorithm());
        printKeyData(wrappedBulkKey);
        sanity(wrappedBulkKey, bulkKey);

        Cipher bulkCipher = Cipher.getInstance("Blowfish/CBC/PKCS5Padding");
        AlgorithmParameterSpec bulkParams = new IvParameterSpec(IV);
        bulkCipher.init(Cipher.ENCRYPT_MODE, bulkKey, bulkParams);

        copyData(bulkCipher);
    }

    private void sanity(byte[] wrappedBulkKey, SecretKey bulkKey) throws Exception {
        Cipher deRsa = Cipher.getInstance("RSA");
        deRsa.init(Cipher.UNWRAP_MODE,this.rsaKeyPair.getPrivateKey());
        SecretKey sk = (SecretKey) deRsa.unwrap(wrappedBulkKey,bulkKey.getAlgorithm(),Cipher.SECRET_KEY);
        System.out.println("Sanity check: " + sk.equals(bulkKey));
    }

    private static void printKeyData(byte[] wrappedBulkKey) {
        System.err.println("Key length: " + wrappedBulkKey.length);
        System.err.print("Key data: ");
        hexPrint(wrappedBulkKey, System.err);
        System.err.println();
    }

    private static final char[] hexDigit = "0123456789ABCDEF".toCharArray();
    private static void hexPrint(byte[] data, PrintStream out) {
       for(byte b:data) {
           out.print(" 0x");
           out.print(hexDigit[(b&0xF0)>>8]);
           out.print(hexDigit[b&0x0F]);
       }

    }

    private void copyData(Cipher bulkCipher) throws Exception {
        byte[] buffer = new byte[2048];
        int len;
        while( (len = this.inFile.read(buffer)) >= 0 ) {
            this.outFile.write(bulkCipher.update(buffer,0,len));
        }
        this.outFile.write(bulkCipher.doFinal());
    }

    private void decrypt() throws Exception {
        ByteBuffer keyLengthBuffer = ByteBuffer.allocate(4);
        if( 4 != this.inFile.read(keyLengthBuffer.array()) ) {
            throw new IOException("Unexpected end of file");
        }
        int keyLength = keyLengthBuffer.getInt();
        byte[] bulkKeyData = new byte[keyLength];
        if( keyLength != this.inFile.read(bulkKeyData) ) {
            throw new IOException("Unexpected end of file");
        }

        printKeyData(bulkKeyData);

        Cipher keyCipher = Cipher.getInstance("RSA");
        keyCipher.init(Cipher.UNWRAP_MODE,this.rsaKeyPair.getPrivateKey());
        SecretKey bulkKey = (SecretKey) keyCipher.unwrap(bulkKeyData,"Blowfish",Cipher.SECRET_KEY);

        Cipher bulkCipher = Cipher.getInstance("Blowfish/CBC/PKCS5Padding");
        AlgorithmParameterSpec bulkParams = new IvParameterSpec(IV);
        bulkCipher.init(Cipher.DECRYPT_MODE, bulkKey, bulkParams);

        copyData(bulkCipher);
    }

    /**
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {
        RSA rsa = parseArgs(args);
        try {
            rsa.copy();
        } finally {
            rsa.close();
        }
    }

    private static final String KEYID_SWITCH = "--keyid:";
    private static final String KEYPASSWORD_SWITCH = "--keypassword:";
    private static final String KEYSTOREPASSWORD_SWITCH = "--keystorepassword:";
    private static final String DECRYPT_SWITCH = "-d";
    private static RSA parseArgs(String[] args) throws Exception {
        String keyId = null;
        String inFileName = null;
        String outFileName = null;
        String keyPassword = null;
        String keyStorePassword = null;
        boolean isEncrypt  = true;

        for( String arg: args ) {
            if( arg.startsWith("-")) {
                if( arg.startsWith(KEYID_SWITCH)) {
                    keyId = arg.substring(KEYID_SWITCH.length());
                } else if( arg.startsWith(KEYSTOREPASSWORD_SWITCH)) {
                    keyStorePassword = arg.substring(KEYSTOREPASSWORD_SWITCH.length());
                } else if( arg.startsWith(KEYPASSWORD_SWITCH)) {
                    keyPassword = arg.substring(KEYPASSWORD_SWITCH.length());
                } else if( arg.equals(DECRYPT_SWITCH)) {
                    isEncrypt = false;
                } else {
                    error("Unrecognized option: " + arg, -1);
                }
            } else if( null == inFileName ) {
                inFileName = arg;
            } else if( null == outFileName ) {
                outFileName = arg;
            } else {
                error("Too many file names: " + arg, -1);
            }
        }
        if( (null == keyId) || (null == inFileName) || (null == outFileName) || (null == keyPassword)
                || (null == keyStorePassword )) {
            error("Missing argument!", -1);
        }

        return new RSA(isEncrypt, keyId, keyPassword, keyStorePassword, inFileName, outFileName);
    }
    private static void error(String message, int exitCode) {
        if( null != message ) {
            System.err.println(message+"\n");
        }
        System.err.println("Usage:");
        System.err.println("\tRSA [-d] --keyid:<keyid> --keystorepassword:<password> --keypassword:<password> <inputfile> <outputfile>");
        System.exit(exitCode);
    }

}