#include <stdio.h>
#include <string.h>
#include "fips.h"
#include "../utils.h"
#include "../error.h"

/*
 * SSL Format
 * RSA PUBLIC KEY -> PKCS#1 format
 * PUBLIC KEY -> PEM Format
 */

static int DEBUG = 0;

int fips(const char *pkey, struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const int type, const int is_pubkey, const int to_stdout){
    int res;

    if(type == TYPE_RSA){
        /* If it's a pub key, we are going to analyse it */
        if (is_pubkey == 1)
            res = fips_pubkey_rsa(st_audit_fips, st_keyinfo, pkey, to_stdout);
        else 
            res = fips_privkey_rsa(st_audit_fips, st_keyinfo, pkey, to_stdout);
    }
    else if (type == TYPE_X509){
        if (openssl_version() == 1)
            res = fips_x509_v1(st_audit_fips, st_keyinfo, pkey, to_stdout);
        else
            res = fips_x509_v3(st_audit_fips, st_keyinfo, pkey, to_stdout);
    }
    return res;
}
/*
 * This function load public RSA key and make an audit on it
 */
static int fips_pubkey_rsa(struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const char *pkey, const int to_stdout) {
    struct rsa *rsa;
    size_t keysize;
    int res;

    /*
     * Now, we check if the public certificate is compliant with FIPS
     * The program check the length of the key, the exponent
     */
    if (openssl_version() == 1)
        res = loadkeys_rsa_v1(&rsa, pkey, &st_keyinfo->format);
    else /* OpenSSL version 3 */
        res = loadkeys_rsa_v3(&rsa, pkey, &st_keyinfo->format);

    if (res > 0 || rsa == NULL){
        printf("Failed to read the public key\n");
        clean_rsa_st(rsa);
        return res;
    }

    audit_rsa_keys(rsa, st_audit_fips, st_keyinfo, pkey);

    // Clean
    clean_rsa_st(rsa);
    return 0;
}
/*
 * This function audit RSA private key
 */
static int fips_privkey_rsa(struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const char *pkey, const int to_stdout) {
    int res;
    struct rsa *rsa = NULL;

    if (openssl_version() == 1)
        res = load_priv_rsa_keys_v1(&rsa, pkey);
    else /* OpenSSL version 3 */
        res = load_priv_rsa_keys_v3(&rsa, pkey);

    if (res > 0){
        printf("Failed to read the private key\n");
        clean_rsa_st(rsa);
        return res;
    }

    audit_rsa_keys(rsa, st_audit_fips, st_keyinfo, pkey);

    // Clean
    clean_rsa_st(rsa);
    return 0;
}

/*
 * If OpenSSL v3, we need tu use new OpenSSL functions for reading keys
 */
static int loadkeys_rsa_v3(struct rsa **rsa, const char *pkey, int *format){
    *rsa = (struct rsa*)malloc(sizeof(struct rsa*));
    
    if (*rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }
    
    memset(*rsa, 0, sizeof(struct rsa*));

    (*rsa)->bio = BIO_new(BIO_s_file());
    if (BIO_read_filename((*rsa)->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }
    
    #if OPENSSL_VERSION_NUMBER > 0x03000000f
    (*rsa)->evp = PEM_read_bio_PUBKEY_ex((*rsa)->bio, NULL, NULL, NULL, NULL, NULL);
    if ((*rsa)->evp == NULL){
        return FIPS_ERR_LOAD_KEY;
    }

    //printf("Keysize: %d\n", EVP_PKEY_bits((*rsa)->evp));
    (*rsa)->rsa = EVP_PKEY_get1_RSA((*rsa)->evp);
    if ((*rsa)->rsa == NULL){
        return FIPS_ERR_LOAD_RSA_KEY;
    }
    #endif

    // Get the format of the key
    *format = 0;
    return 0;
}
static int loadkeys_rsa_v1(struct rsa **rsa, const char *pkey, int *format){
    *rsa = (struct rsa*)malloc(sizeof(struct rsa*));

    if (*rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }
    memset(*rsa, 0, sizeof(struct rsa*));

    (*rsa)->bio = BIO_new(BIO_s_file());
    if(BIO_read_filename((*rsa)->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }
    
    /*
     * Works with PEM_read_RSAPublicKey, but when we try to read the file
     * we cannot. This function "block" the access to the file
     */
    /*rsa->rsa = PEM_read_RSAPublicKey(f, NULL, NULL, NULL); */

    // Deprecated in OpenSSL v3
    /*
     * RSAPublicKey read publickey at the PEM format
     * RSA_PUBKEY read publickey at the PKCS1 format
     */
    (*rsa)->rsa = PEM_read_bio_RSAPublicKey((*rsa)->bio, NULL, NULL, NULL);
    if ((*rsa)->rsa == NULL){
        //print_error();

        /*
         * We need to reset or reseek the BIO, otherwise, we cannot read it
         * https://docs.openssl.org/3.0/man3/BIO_ctrl/#synopsis
         */
        //BIO_reset((*rsa)->bio); /* Works too */
        BIO_seek((*rsa)->bio, 0);

        (*rsa)->rsa = PEM_read_bio_RSA_PUBKEY((*rsa)->bio, NULL, NULL,NULL);

        if ((*rsa)->rsa == NULL){
            if (DEBUG)
                printf("Cannot read the SPKI format of the public key\n");
            return FIPS_ERR_LOAD_KEY;
        }

        *format = RSA_FORMAT_SPKI;
    }
    else
        *format = RSA_FORMAT_PKCS1;

    return 0;
}

/*
 * This function load RSA Private key for OpenSSL v1
 */
static int load_priv_rsa_keys_v1(struct rsa **rsa, const char *pkey){
    *rsa = (struct rsa*)malloc(sizeof(struct rsa*));

    if (*rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }

    (*rsa)->bio = BIO_new(BIO_s_file());
    if ((*rsa)->bio == NULL){
        if (DEBUG)
            printf("Failed to create new BIO\n");
        return FIPS_ERR_NEW_BIO;
    }
    if(BIO_read_filename((*rsa)->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }

    (*rsa)->rsa = PEM_read_bio_RSAPrivateKey((*rsa)->bio, NULL, NULL, NULL);
    if ((*rsa)->rsa == NULL){
        if (DEBUG)
            printf("Failed to read BIO RSAPrivateKey\n");
        return FIPS_ERR_LOAD_RSA_PRIV_KEY;
    }

    return 0;
}
/*
 * This function load RSA Private key for OpenSSL v3
 */
static int load_priv_rsa_keys_v3(struct rsa **rsa, const char *pkey){
    *rsa = (struct rsa*)malloc(sizeof(struct rsa*));

    if (*rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }

    (*rsa)->bio = BIO_new(BIO_s_file());
    if ((*rsa)->bio == NULL){
        if (DEBUG)
            printf("Failed to create new BIO\n");
        return FIPS_ERR_NEW_BIO;
    }
    if(BIO_read_filename((*rsa)->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }

    #if OPENSSL_VERSION_NUMBER > 0x03000000f
    (*rsa)->evp = PEM_read_bio_PrivateKey_ex((*rsa)->bio, NULL, NULL, NULL, NULL, NULL);

    if ((*rsa)->evp == NULL){
        if (DEBUG)
            printf("Failed to read BIO PrivateKey\n");
        return FIPS_ERR_READ_BIO;
    }
    (*rsa)->rsa = EVP_PKEY_get1_RSA((*rsa)->evp);
    if ((*rsa)->rsa == NULL)
        return FIPS_ERR_LOAD_RSA_KEY;
    #endif

    return 0;
}

/*
 * This function audit the RSA keys, both public and private
 * For the audit, the function check the exponent (modulus) and the key size
 */
static void audit_rsa_keys(struct rsa *rsa, struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const char *pkey){
    int res;

    st_keyinfo->keysize = RSA_size(rsa->rsa);

    st_keyinfo->algo = ALGO_RSA;

    // The return value is a const, shouldn't be freed
    const BIGNUM *e = RSA_get0_e(rsa->rsa);
    char *exponent = BN_bn2dec(e);
    //free(exponent);
    OPENSSL_free(exponent);

    /* Exponent has been set up, we can check it */
    res = check_exponent(e, st_audit_fips->audit_exponent.result, &st_keyinfo->exponent);
    
    /*
     * Audit the key size. For a better security, the key size is at least 2048 bits
     */
    if (st_keyinfo->keysize * 8 < 2048){
        sprintf(st_audit_fips->audit_keysize.result, "The key size is lower than 2048. The key should be at least 2048 bits.");
        st_audit_fips->audit_keysize.audit = FALSE;
    }
    else{
        sprintf(st_audit_fips->audit_keysize.result, "The key size is upper or equal than 2048. The audit is passed with success.");
        st_audit_fips->audit_keysize.audit = TRUE;
    }
}

/*
 * In this function, we are going to check the exponent
 * For testing if the exponent is odd or even, we apply a modulo 2 on the exponent
 * If the result is 1, means the key has a remainder and the key is odd, if the result is 0, the exponent is even. 
 * Regarding to the FIPS 186-5, the exponent must be odd.
 * The function check also the size of the exponent. 
 * When the key has been generated with OpenSSL, by default the exponent is 65537.
 * The exponent e size must be 2 ** 16 < e < 2 ** 256
 */
static int check_exponent(const BIGNUM *e, char *buf, unsigned long *exponent){
    BIGNUM *rem = BN_new(), *a = BN_new(), *m = BN_new();
    char nExponent[4];
    BN_CTX *ctx;
    char *r;
    int error = 0;

    sprintf(nExponent, "%d", 2);
    BN_dec2bn(&m, nExponent); 
    //BN_mod(rem, e, m); 
    ctx = BN_CTX_new();
    BN_div(NULL, rem, e, m, ctx);
    //r = malloc(4);
    r = BN_bn2dec(rem);
    //printf("%s\n", r);

    /*
     * According to the FIPS 186-5, the exponent size must be 2 ** 16 < e < 2 ** 256
     * The exponent must be odd too
     */
    char *exp = BN_bn2dec(e);
    *exponent = (char2dec(exp[0]) * 10000) + (char2dec(exp[1]) * 1000) + (char2dec(exp[2])* 100) + (char2dec(exp[3]) * 10) + (char2dec(exp[4]));
    // printf("Exp: %lu\n", exp);

    if (strcmp(r, "0") == 0){
        strncpy(buf, "The exponent is even, should be odd", BUF_SIZE_AUDIT);
        error += 1;
    }

    // Check the exponent size
    double minSize = pow(2, 16);
    double maxSize = pow(2, 256);
    if (*exponent < minSize || *exponent > maxSize){
        strncpy(buf, "The exponent size is not correct. The minimum size is 2 ** 16 and maximum size 2 ** 256.", BUF_SIZE_AUDIT);
        error += 1;
    }

    // If no error
    if (error == 0)
        strncpy(buf, "The exponent is correct, the FIPS compliance is respected.", BUF_SIZE_AUDIT);

    // Cleaning
    free(r);
    BN_free(rem);
    BN_free(a);
    BN_free(m);
    BN_CTX_free(ctx);
    OPENSSL_free(exp);
    return 0;
}

/*
 * This function load X509 certificate for OpenSSL v1
 */
static int fips_x509_v1(struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const char *pkey, const int to_stdout){
    struct rsa *rsa = (struct rsa*)malloc(sizeof(struct rsa*));

    if (rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }
    memset(rsa, 0, sizeof(struct rsa*));

    rsa->bio = BIO_new(BIO_s_file());
    if (BIO_read_filename(rsa->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }

    X509 *x = PEM_read_bio_X509(rsa->bio, NULL, 0, NULL);
    if (x == NULL){
        printf("Failed to read the X509 certificate\n");
        clean_rsa_st(rsa);
        return FIPS_ERR_LOAD_X509;
    }

    EVP_PKEY *evp = X509_get_pubkey(x);
    if (evp == NULL){
        printf("Failed to get public certificate\n");
        X509_free(x);
        return FIPS_ERR_LOAD_RSA_KEY;
    }
    rsa->rsa = EVP_PKEY_get1_RSA(evp);
    if (rsa->rsa == NULL){
        X509_free(x);
        EVP_PKEY_free(evp);
        return FIPS_ERR_LOAD_RSA_KEY;
    }

    // We have the RSA key, we can audit it
    audit_rsa_keys(rsa, st_audit_fips, st_keyinfo, pkey);

    X509_free(x);
    EVP_PKEY_free(evp);
    clean_rsa_st(rsa);
    return 0;
}
/*
 * This function load X509 certificate for OpenSSL v3
 */
static int fips_x509_v3(struct audit_fips *st_audit_fips, struct keyinfo *st_keyinfo, const char *pkey, const int to_stdout){
    struct rsa *rsa = (struct rsa*)malloc(sizeof(struct rsa*));

    if (rsa == NULL){
        if (DEBUG)
            printf("Cannot malloc the structure\n");
        return COMMON_ERR_MALLOC;
    }
    memset(rsa, 0, sizeof(struct rsa*));

    rsa->bio = BIO_new(BIO_s_file());
    if (BIO_read_filename(rsa->bio, pkey) == 0){
        printf("Failed to read BIO\n");
        return FIPS_ERR_READ_BIO;
    }

    #if OPENSSL_VERSION_NUMBER > 0x03000000f
    rsa->evp = NULL; // Otherwise, I have a seg fault during clean rsa st

    X509 *x = PEM_read_bio_X509(rsa->bio, NULL, 0, NULL);
    if (x == NULL){
        printf("Failed to read the X509 certificate\n");
        clean_rsa_st(rsa);
        return FIPS_ERR_LOAD_X509;
    }

    rsa->evp = X509_get_pubkey(x);
    if (rsa->evp == NULL){
        printf("Failed to get public certificate\n");
        X509_free(x);
        return FIPS_ERR_LOAD_RSA_KEY;
    }
    rsa->rsa = EVP_PKEY_get1_RSA(rsa->evp);
    if (rsa->rsa == NULL){
        X509_free(x);
        return FIPS_ERR_LOAD_RSA_KEY;
    }
    #endif

    /*
     * TODO: Need to identify public-key cryptosystem: RSA or EC
     */

    // We have the RSA key, we can audit it
    audit_rsa_keys(rsa, st_audit_fips, st_keyinfo, pkey);

    #if OPENSSL_VERSION_NUMBER > 0x03000000f
    X509_free(x);
    #endif
    clean_rsa_st(rsa);
    return 0;
}

/*
 * This function clean the RSA structure
 */
static void clean_rsa_st(struct rsa *rsa){
    if (DEBUG)
        printf("\nCleaning RSA\n");

    if (rsa == NULL)
        return;

    if (rsa->bio != NULL)
        BIO_free(rsa->bio);

    if (openssl_version() == 3){
        #if OPENSSL_VERSION_NUMBER > 0x03000000f
        if (rsa->evp != NULL)
            EVP_PKEY_free(rsa->evp);
        #endif
    }

    if (rsa->rsa != NULL)
        RSA_free(rsa->rsa);

    free(rsa);
    rsa = NULL;
}

/*
 * Return 1 if the version is upper than 1 and less than 3
 * Return 3 for the version v3
 */
static int openssl_version(){
    unsigned long version = OPENSSL_VERSION_NUMBER;
    
    if (DEBUG)
        printf("OpenSSL Version: %lx\n", version);
    if (version <= 0x03000000f)
        return 1;
    if (version >= 0x03000000f)
        return 3; 
}
/*
 * In case we have an error with OpenSSL librairy, we can print the error message
 */
static void print_error(){
    unsigned long err = ERR_get_error();
    char b[256];
    ERR_error_string(err, b);
    printf("%s\n", b);
}