#include <openssl/aes.h>
#include <openssl/rand.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>
#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/err.h>
#include "modelDecrypt.h"

int padding = RSA_PKCS1_OAEP_PADDING;

RSA * createRSAWithFilename(char * filename,int public)
{
    FILE * fp = fopen(filename,"rb");
    if(fp == NULL)
    {
        printf("Unable to open file %s \n",filename);
        return NULL;    
    }
    RSA *rsa= RSA_new() ;
 
    if(public)
    {
        rsa = PEM_read_RSA_PUBKEY(fp, &rsa,NULL, NULL);
    }
    else
    {
        rsa = PEM_read_RSAPrivateKey(fp, &rsa,NULL, NULL);
    }
 
    return rsa;
}

// int private_decrypt(unsigned char * enc_data,int data_len,char * fname, unsigned char *decrypted)
// {
//     int  result = 0;
//     RSA *rsa = createRSAWithFilename(fname,1);
//     result = RSA_public_decrypt(data_len,enc_data,decrypted,rsa,padding);
// 	/* your code goes here */
//     return result;
// }
void handleErrors(void)
{
    ERR_print_errors_fp(stderr);
    abort();
}
int encryptAES(unsigned char *plaintext, int plaintext_len, unsigned char *key,
            unsigned char *iv, unsigned char *ciphertext)
{
    EVP_CIPHER_CTX *ctx;

    int len;

    int ciphertext_len;

    /* Create and initialise the context */
    if(!(ctx = EVP_CIPHER_CTX_new()))
        handleErrors();

    /*
     * Initialise the encryption operation. IMPORTANT - ensure you use a key
     * and IV size appropriate for your cipher
     * In this example we are using 256 bit AES (i.e. a 256 bit key). The
     * IV size for *most* modes is the same as the block size. For AES this
     * is 128 bits
     */
    if(1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv))
        handleErrors();

    /*
     * Provide the message to be encrypted, and obtain the encrypted output.
     * EVP_EncryptUpdate can be called multiple times if necessary
     */
    if(1 != EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len))
        handleErrors();
    ciphertext_len = len;

    /*
     * Finalise the encryption. Further ciphertext bytes may be written at
     * this stage.
     */
    if(1 != EVP_EncryptFinal_ex(ctx, ciphertext + len, &len))
        handleErrors();
    ciphertext_len += len;

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return ciphertext_len;
}
int decryptAES(unsigned char *ciphertext, int ciphertext_len, unsigned char *key,
            unsigned char *iv, unsigned char *plaintext)
{
    EVP_CIPHER_CTX *ctx;

    int len;

    int plaintext_len;

    /* Create and initialise the context */
    if(!(ctx = EVP_CIPHER_CTX_new()))
        handleErrors();

    /*
     * Initialise the decryption operation. IMPORTANT - ensure you use a key
     * and IV size appropriate for your cipher
     * In this example we are using 256 bit AES (i.e. a 256 bit key). The
     * IV size for *most* modes is the same as the block size. For AES this
     * is 128 bits
     */
    if(1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv))
        handleErrors();

    /*
     * Provide the message to be decrypted, and obtain the plaintext output.
     * EVP_DecryptUpdate can be called multiple times if necessary.
     */
    if(1 != EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len))
        handleErrors();
    plaintext_len = len;

    /*
     * Finalise the decryption. Further plaintext bytes may be written at
     * this stage.
     */
    if(1 != EVP_DecryptFinal_ex(ctx, plaintext + len, &len))
        handleErrors();
    plaintext_len += len;

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return plaintext_len;
}

void hybrid_encrypt(char * symmetric_key_file, char* message_file, char* public_key_file) {
    int  result = 0;
    RSA *rsa = createRSAWithFilename(public_key_file,1);
    unsigned char * encKey = malloc(sizeof(RSA_size(rsa))); 
    FILE *fp = fopen(symmetric_key_file, "r");
    fseek(fp, 0, SEEK_END); 
    int keyLen = ftell(fp); 
    rewind(fp);
    unsigned char* symmKey = malloc(sizeof(unsigned char)*keyLen + 1);
    fread(symmKey, keyLen, 1, fp);
    symmKey[keyLen]='\0';
    result = RSA_public_encrypt(keyLen,symmKey,encKey,rsa,padding);
	/* Init vector */
	unsigned char iv[AES_BLOCK_SIZE];
    for(int i = 0; i<AES_BLOCK_SIZE;i++)
        iv[i] = rand();
    fp = fopen(message_file, "r");
    puts("reading file");
    fseek(fp, 0, SEEK_END); 
    int msgLen = ftell(fp); 
    rewind(fp);
    unsigned char* aes_input = malloc(sizeof(unsigned char)*msgLen + 1);
    unsigned char* enc_out = malloc(sizeof(unsigned char)*msgLen*128 + 1);

    fread(aes_input, msgLen, 1, fp);
    int ciphertext_len = encryptAES (aes_input, strlen ((char *)aes_input), symmKey, iv,
                              enc_out);
    enc_out[ciphertext_len]='\0';
    FILE *C1 = fopen("C1","w");
    FILE *C2 = fopen("C2","w");
    FILE *IV = fopen("IV","w");
    fprintf(C1,"%s",encKey);
    fprintf(C2,"%s",enc_out);
    fprintf(IV,"%s",iv);
    fclose(C1);
    fclose(C2);
    fclose(IV);
    fclose(fp);
}

int main(int argc, char *argv[]) {

	if(argc < 4){
		printf("Usage : \narg1:message_file \narg2:private_key_file \narg3:public_key_file \narg4:symmetric_key\n");
    	exit(1);
    }
    puts("calling hybrid encyrpt");
    hybrid_encrypt(argv[4],argv[1], argv[3]);
    decrypt(argv[2], "decrypted_message.txt");
    
    printf("output : %d\n",compare_files("message.txt","decrypted_message.txt"));
	return 0;

}
