#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "util.h"

#include <openssl/sha.h>
#include <openssl/ec.h>
#include <openssl/bn.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/core_names.h>

#include "embedded_device.h"

#define BUFSIZE 16384

typedef struct {
    uint8_t *firmware_path;
    uint8_t *public_key_path;
} embedded_device_args;

char *embedded_device_syntax(void) {
    return
        "embedded            mock up implementation of a embedded device\n"
        "    -ppub  <path>   file path of the PEM encoded public key of the proxy\n"
        "    -firm  <path>   path of to firmware binary\n";
}

static EVP_PKEY *read_public_key(char *public_key_file_path, EVP_PKEY **key) {
    if(public_key_file_path == NULL) {
        fprintf(stderr, "public_key_file_path is a null pointer!\n");
        return NULL;
    }

    FILE *fd = fopen(public_key_file_path, "rb");
    if(fd == NULL) {
        fprintf(stderr, "failed to open public key file!\n");
        return NULL;
    }

    *key = PEM_read_PUBKEY(fd, key, NULL, NULL);
    fclose(fd);

    return *key;
}

static void hash_firmware(uint8_t *firmware_path, EVP_MD_CTX **ctx) {
    if(firmware_path == NULL) {
        fprintf(stderr, "firmware_path is a null pointer!\n");
        return;
    }

    FILE *fd = fopen(firmware_path, "rb");
    if(fd == NULL) {
        fprintf(stderr, "failed to open firmware!\n");
        goto exit;
    }

    size_t size;
    uint8_t buf[BUFSIZE];
    while((size = fread(buf, 1, BUFSIZE, fd)) != 0) {
        EVP_DigestVerifyUpdate(*ctx, buf, size);
    }
    
    exit: fclose(fd);
}

int handle_embedded_device(int argc, char **argv) {
    embedded_device_args args = {
        .firmware_path = NULL,
        .public_key_path = NULL
    };

    for(int i = 0; i < argc; i += 2) {
        if((strcmp(argv[i], "-ppub") == 0) && (argc - i >= 2)) {
            args.public_key_path = argv[i+1];
        } else if((strcmp(argv[i], "-firm") == 0) && (argc - i >= 2)) {
            args.firmware_path = argv[i+1];
        } else {
            syntax_exit();
        }
    }

    if((args.firmware_path == NULL) || (args.public_key_path == NULL)) {
        fprintf(stderr, "failed to parse arguments");
        exit(EXIT_FAILURE);
    }

    EVP_PKEY *key = NULL;
    if(read_public_key(args.public_key_path, &key) == NULL) {
        fprintf(stderr, "failed to import public key");
        goto clean;
    }

    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (EVP_DigestVerifyInit(ctx, NULL, EVP_sha256(), NULL, key) != 1) {
        fprintf(stderr, "failed to initialize context\n");
        goto clean;
    }

    uint8_t signature[BUFSIZE] = {0};
    size_t signature_size = read(0, signature, BUFSIZE);
    if(signature_size < 70) {
        fprintf(stderr, "failed to read firmware signature\n");
        goto clean;
    }

    hash_firmware(args.firmware_path, &ctx);
    if (EVP_DigestVerifyFinal(ctx, signature, signature_size) != 1) {
        fprintf(stderr, "failed to verify firmware signature\n");
    }else {
        printf("successfully verified firmware signature\n");
    }

    clean: ;
    if(key != NULL)
        EVP_PKEY_free(key);
    if(ctx != NULL)
        EVP_MD_CTX_free(ctx);

    return 0;
}
