#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/aes.h>

#define CBC 0
#define ECB 1

int mode;
int type;
unsigned char vector[AES_BLOCK_SIZE];
int bytes_to_write = AES_BLOCK_SIZE;
unsigned char in[AES_BLOCK_SIZE], out[AES_BLOCK_SIZE];

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

    //file_in file_out [-encrypt/-decrypt] key [-ecb/-cbc] [vector]

    FILE* file_in = fopen(argv[1],"rb");
    if (file_in == NULL) {
        printf("Otwieranie pliku wejsciowego nie powiodlo sie.\n");
        return -1;
    }

    FILE* file_out = fopen(argv[2],"wb");
    if (file_out == NULL) {
        printf("Otwieranie pliku wyjsciowego nie powiodlo sie.\n");
        return -1;
    }

    if ( strcmp(argv[3],"-encrypt") == 0) {
        mode = AES_ENCRYPT;
    } else if (strcmp(argv[3],"-decrypt") == 0) {
        mode = AES_DECRYPT;
    } else {
        return -1;
    }


    if ( strcmp(argv[5],"-ecb") == 0) {
        type = ECB;
    } else if (strcmp(argv[5],"-cbc") == 0) {
        type = CBC;
        FILE * vector_file = fopen(argv[6], "rb");
            if (vector_file == NULL) {
                return -1;
            }
            fread(vector, sizeof(vector), 1, vector_file);
         fclose(vector_file);
    } else {
        return -1;
    }

    // AES
    AES_KEY key;
    if (mode == AES_ENCRYPT){
        AES_set_encrypt_key(argv[4],128,&key);
    } else {
        AES_set_decrypt_key(argv[4],128,&key);
    }


    int i;
    int padding;
    int withoutOne = 0;
    unsigned char in[AES_BLOCK_SIZE], out[AES_BLOCK_SIZE];
    while(1) {

        size_t read_bytes = fread(in, 1, AES_BLOCK_SIZE, file_in);

        if(read_bytes == 0){
            break;
        }

        if(withoutOne && mode == AES_DECRYPT){
            fwrite(out, 1, AES_BLOCK_SIZE, file_out);
        }

        withoutOne = 1;

        //PKCS7
        if(read_bytes < AES_BLOCK_SIZE){
            if(mode == AES_ENCRYPT){
                padding = AES_BLOCK_SIZE - read_bytes;
                printf("Padding-Encrypt: %d\n\n",padding);
                for(i=read_bytes;i<AES_BLOCK_SIZE; i++){
                    in[i] = padding;
                }
            }
        }else{
            padding = 0;
        }

        //szyfrowanie
        if (type == CBC){
            AES_cbc_encrypt(in, out, sizeof(in), &key, vector, mode);
        }

        if (type == ECB){
            AES_ecb_encrypt(in, out, &key, mode);
        }

        if(mode == AES_ENCRYPT){
            fwrite(out, 1, AES_BLOCK_SIZE, file_out);
        }
    }

    // jesli nie ma paddingu
    if(mode == AES_ENCRYPT && padding == 0){
        memset(in, AES_BLOCK_SIZE, AES_BLOCK_SIZE);

        if (type == CBC){
            AES_cbc_encrypt(in, out, sizeof(in), &key, vector, mode);
        }

        if (type == ECB){
            AES_ecb_encrypt(in, out, &key, mode);
        }
        fwrite(out, 1, AES_BLOCK_SIZE, file_out);
    }

    if(mode == AES_DECRYPT){
        int omit = AES_BLOCK_SIZE - out[AES_BLOCK_SIZE -1];
        fwrite(out, 1, omit, file_out);
    }

    fclose(file_in);
    fclose(file_out);
    return 0;
}