#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "aes.h"

#define BUFFER_SIZE 0x4000

unsigned char inbuf[BUFFER_SIZE];
unsigned char outbuf[BUFFER_SIZE];

unsigned char iv[16] = "abcdefghijklmnop";

int main(int argc, char **argv) {
    FILE *infile, *outfile;
    size_t count;
    char *filename;
    struct aes_context context;
    int encrypt;
    
    if (argc != 4) {
        fprintf(stderr, "wrong number of arguments\n");
        exit(EXIT_FAILURE);
    }
    
    if (strncmp(argv[1], "-d", 2) == 0) {
        encrypt = 0;
    } else if (strncmp(argv[1], "-e", 2) == 0) {
        encrypt = 1;
    }
    
    infile = fopen(argv[3], "rb");
    if (infile == NULL) {
        fprintf(stderr, "error opening input file\n");
        return EXIT_FAILURE;
    }

    filename = malloc(strlen(argv[3]) + 5);
    strcpy(filename, argv[3]);
    if (encrypt) {
        strcat(filename, ".aes");
    } else {
        strcat(filename, ".dec");
    }
    
    outfile = fopen(filename, "wb");
    if (filename == NULL) {
        fprintf(stderr, "error opening output file\n");
        return EXIT_FAILURE;
    }
    
    memcpy(context.data, iv, 16);
    strncpy(context.key, argv[2], 32);
    
    aes128_expandkey(&context);

    do {
        count = fread(inbuf, 1, BUFFER_SIZE, infile);
        count = (count + 15) / 16;
        if (encrypt) {
            aes_cbc_encrypt(&context, inbuf, outbuf, count);
        } else {
            aes_cbc_decrypt(&context, inbuf, outbuf, count, iv);
            if (count != 0)
                memcpy(iv, inbuf + (count - 1) * 16, 16);
        }
        count *= 16;
        if (fwrite(outbuf, 1, count, outfile) != count) {
            fprintf(stderr, "error writing file\n");
            goto finish;
        }
    } while (count != 0);

finish:
    fclose(infile);
    fclose(outfile);
    
    free(filename);
}