#include<string.h>
#include<stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <openssl/rsa.h>
#include <openssl/bn.h>
#include <openssl/sha.h>
#include <openssl/pem.h>
#include <openssl/evp.h>
#include <openssl/objects.h>

void WrongUsage(){
    printf("Usage: docsign -sv\n");
    printf("-s sign document\n");
    printf("-v validate signature\n");
    exit(0);
}
void generateRSAkeys(){
    printf("Generating keys\n");
    RSA *rsa=RSA_new();
    BIGNUM *bn=BN_new();
    BN_set_word(bn, RSA_F4);
    RSA_generate_key_ex(rsa, 2048, bn, NULL);
    BN_free(bn);

    FILE* private=fopen("private.pem", "w");
    PEM_write_RSAPrivateKey(private, rsa, NULL, NULL, 0, 0, NULL);
    fclose(private);

    FILE* public=fopen("public.pem", "w");
    PEM_write_RSA_PUBKEY(public, rsa);
    fclose(public);

    RSA_free(rsa);
}
void printHEX(char* arr, int sz){
    for(int i=0; i<sz; i++)
        printf("%02x", arr[i]);
    printf("\n");
}
int main(int argc, char** argv){
    OpenSSL_add_all_digests();
    if(argc!=3)
        WrongUsage();
    if(access( argv[2], F_OK ) == -1) {
        printf("File doesn't exist!!!\n");
        WrongUsage();
    }
    if(!strcmp("-s", argv[1])){
        printf("Starting Signing Proccess\n");
        printf("-------------------------\n");
        printf("Check if keys exist\n");
        if(access("private.pem", F_OK ) == -1 || access("public.pem", F_OK )==-1){
            printf("Keys Don't exist");
            generateRSAkeys();
        }
        printf("Reading private key\n");
        FILE* private=fopen("private.pem", "r");
        RSA *rsa = NULL;
        PEM_read_RSAPrivateKey(private, &rsa, NULL, NULL);
        fclose(private);
        EVP_PKEY *key=EVP_PKEY_new();
        EVP_PKEY_assign_RSA(key, rsa);
        printf("Calculating Digest\n");
        FILE* fin=fopen(argv[2], "rb");
        FILE* fout=fopen(strcat(argv[2],".signed"), "wb");
        EVP_MD_CTX *ctx = EVP_MD_CTX_create();
        const EVP_MD *md = EVP_get_digestbyname("SHA256");
        EVP_SignInit_ex(ctx, md, NULL);
        unsigned char *buffer = malloc(32768);
        int bytesRead;
        while((bytesRead = fread(buffer, 1, 32768, fin))){
            EVP_SignUpdate(ctx, buffer, bytesRead);
            fwrite(buffer, 1, bytesRead, fout);
        }
        printf("Signing the digest\n");
        unsigned char signature[256];
        int len;
        EVP_SignFinal(ctx, signature, &len, key);
        printf("Writing signature\n");
        fwrite(signature, 1, len, fout);
        printHEX(signature, 256);
        fclose(fin);
        fclose(fout);
    }else if(!strcmp("-v", argv[1])){
        printf("Starting Validating Proccess\n");
        printf("----------------------------\n");
        FILE *fin = fopen(argv[2], "rb");
        fseek (fin, 0, SEEK_END);
        int size=ftell(fin);
        printf("%d\n", size);
        char* dataplus=malloc(size);
        fseek (fin, 0, SEEK_SET);
        fread(dataplus, 1, size, fin);
        char* data=malloc(size-256);
        memcpy(data, dataplus, size-256);
        char* signature=malloc(256);
        memcpy(signature, dataplus+size-256, 256);
        size-=256;
        if(access("public.pem", F_OK ) == -1){
            printf("Public Key doesn't exist");
            exit(0);
        }
        RSA *rsa = NULL;
        FILE* public=fopen("public.pem", "r");
        PEM_read_RSA_PUBKEY(public, &rsa, NULL, NULL);
        fclose(public);
        EVP_PKEY *key=EVP_PKEY_new();
        EVP_PKEY_assign_RSA(key, rsa);
        EVP_MD_CTX *ctx = EVP_MD_CTX_create();
        const EVP_MD *md = EVP_get_digestbyname("SHA256");
        EVP_VerifyInit_ex(ctx, md, NULL);
        EVP_VerifyUpdate(ctx, data, size);
        int res=EVP_VerifyFinal(ctx, signature, 256, key);
        if(res==1)
            printf("Digital signature is valid.\n");
        else if(!res)
            printf("Digital signature is not valid.\n");
    }else{
        WrongUsage();
    }
}