#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <aes.h>
#include <time.h>

#include "hwinfo.h"
#include "util.h"

int main(int argc, char** argv) {
    if (argc != 4) {
        return 1;
    }

    char* input = argv[1];
    char* output = argv[2];
    char* readme = argv[3];

    hw_info_t hwinfo = {0};
    if(get_hw_info(&hwinfo) != 0){
        return 1;
    }

    char* env_vendor_id = getenv("ENCRYPTOR_VENDOR_ID");
    if (env_vendor_id != NULL) {
        hwinfo.vendor_id = env_vendor_id;
    } else {
        fprintf(stderr, "Warning: No vendor_id specified, keeping original '%s'\n", hwinfo.vendor_id);
    }

    char* env_model_name = getenv("ENCRYPTOR_MODEL_NAME");
    if (env_model_name != NULL) {
        hwinfo.model_name = env_model_name;
    } else {
        fprintf(stderr, "Warning: No model_name specified, keeping original '%s'\n", hwinfo.model_name);
    }

    char* env_os_version = getenv("ENCRYPTOR_OS_VERSION");
    if (env_os_version != NULL) {
        hwinfo.os_version = env_os_version;
    } else {
        fprintf(stderr, "Warning: No os_version specified, keeping original '%s'\n", hwinfo.os_version);
    }

    char* env_mem_total = getenv("ENCRYPTOR_MEM_TOTAL");
    if (env_mem_total != NULL) {
        hwinfo.mem_total = env_mem_total;
    } else {
        fprintf(stderr, "Warning: No mem_total specified, keeping original '%s'\n", hwinfo.mem_total);
    }

    uint8_t* key = hwinfo_to_key(hwinfo);

    if(key == NULL) {
        return 1;
    }

    fprintf(stderr, "HWINFO: %s\n", hw_info_flatten(hwinfo));

    fprintf(stderr, "Key: ");
    for (int i = 0; i < 16; i++)
        fprintf(stderr, "%02X", key[i]);
    fprintf(stderr, "\n");

    struct AES_ctx ctx;
    uint8_t nonce[16] = {0};

    time_t ts = time(NULL);
    srand(ts);

    for (size_t i = 0; i < sizeof(nonce); i+= 4) {
        uint32_t random = rand();
        nonce[i + 0] = random & 0xFF;
        nonce[i + 1] = (random >>  8) & 0xFF;
        nonce[i + 2] = (random >> 16) & 0xFF;
        nonce[i + 3] = (random >> 24) & 0xFF;
    }

    fprintf(stderr, "Nonce: ");
    for (size_t i = 0; i < sizeof(nonce); i++) {
        fprintf(stderr, "%02X", nonce[i]);
    }
    fprintf(stderr, "\n");

    FILE* license = fopen(input, "rb");
    long data_start = ftell(license);
    fseek(license, 0, SEEK_END);
    int license_size = ftell(license) - data_start;
    fseek(license, data_start, SEEK_SET);
    uint8_t* buf = malloc(license_size);
    fread(buf, 1, license_size, license);


    AES_init_ctx(&ctx, key);
    AES_init_ctx_iv(&ctx, key, nonce);
    AES_CTR_xcrypt_buffer(&ctx, buf, license_size);

    FILE* encrypted_license = fopen(output, "wb+");
    fwrite(&nonce, 1, sizeof(nonce), encrypted_license);
    fwrite(buf, 1, license_size, encrypted_license);

    fclose(license);
    fclose(encrypted_license);

    fprintf(stderr, "Wrote encrypted license file to %s\n", output);

    FILE* f_readme = fopen(readme, "w+");
    fprintf(f_readme, "Oh, and one more thing you should know...\n");
    fprintf(f_readme, "This doesn't just run on any hardware, I don't know how they did it\n");
    fprintf(f_readme, "but it only works on my %s processor\n", hwinfo.model_name);
    fprintf(f_readme, "running linux (version %s) with %s kb of RAM\n", hwinfo.os_version, hwinfo.mem_total);
    fprintf(f_readme, "You'll find a way, slime. I'm sorry that I can't give you any more details\n");
    fprintf(f_readme, "this is the last message I could get out before they... you know... ⛓️‍👮‍♂️\n");
    fprintf(f_readme, "\n");
    fprintf(f_readme, "- ASNT aka. Apfelsaftnaturtrüb\n");

    return 0;
}