#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <openssl/bn.h>
#include <string.h>
#include "dh.h"


// gcc -Wall -o dh dh.c -lcrypto && ./dh


static int generatePrimeNumber(unsigned long long int p){
    for (int i = 2; i < p; i++){
        if ((unsigned long long int)p % i == 0)
            return 1;
    }
    return 0;
}
static BIGNUM *generatePrivateKey(){
    BIGNUM *bn = NULL;
    char n[CHAR_SIZE];
    unsigned long long int r = rand() % (PRIVATEKEY_MAX - PRIVATEKEY_MIN) + PRIVATEKEY_MIN;
    sprintf(n, "%lld", r);
    BN_dec2bn(&bn, n);
    return bn;
}
static void dumpBN(const char *s, BIGNUM *bn){
    char *n = BN_bn2dec(bn);
    printf("%s: %s\n", s, n);
    OPENSSL_free(n);
}
static int combineKey(BIGNUM **pubKey, BIGNUM *e, BIGNUM *x, BIGNUM *bnP, BN_CTX *ctx){
    if (ctx == NULL)
        ctx = BN_CTX_new();

    if (*pubKey == NULL)
        *pubKey = BN_new();

    BIGNUM *exp = BN_new();
    int res = BN_exp(exp, e, x, ctx);

    //dumpBN("publicKey", exp);
    if (res == 0){
        BN_free(exp);
        return  0;
    }
    res = BN_mod(*pubKey, exp, bnP, ctx);

    BN_free(exp);
    return 1;
}
static BIGNUM *generateP(){
    BIGNUM *bnP = NULL;
    unsigned long long int p = rand() % (500 - 200) + 500;
    char nP[CHAR_SIZE];
    int res = 0;

    // First, we need to find a prime number of p
    while (!generatePrimeNumber(p++));

    sprintf(nP, "%lld", p);
    res = BN_dec2bn(&bnP, nP);
    if (res == 0)
        return NULL;
    return bnP;
}
static int charToDec(const char *s){
    int dec = 0;
    size_t l = strlen(s);
    int pos = 1;
    for (int i = 0; i < l; i++){
        dec += (s[i] - '0') * pos;
        pos *= 10;
    }
    return dec;
}
static BIGNUM *generateG(const char *nP){
    BIGNUM *bnG = NULL;
    char nG[CHAR_SIZE];
    int res = 0;
    int p = charToDec(nP); // Convert char* to int
    unsigned long long int g = rand() % (p - 1);
    sprintf(nG, "%lld", g);
    res = BN_dec2bn(&bnG, nG);
    if (res == 0)
        return NULL;
    return bnG;
}
static void cleanDH(struct bn_dh *dh){
    if (dh->bnPrivateKey != NULL)
        BN_free(dh->bnPrivateKey);
    if (dh->bnPublicKey != NULL)
        BN_free(dh->bnPublicKey);
    if (dh->bnPublicKeyPeer != NULL)
        BN_free(dh->bnPublicKeyPeer);
    if (dh->bnPublicEphemeralKey != NULL)
        BN_free(dh->bnPublicEphemeralKey);
    if (dh->ctx != NULL)
        BN_CTX_free(dh->ctx);
}
int main(void){
    char *nP;
    struct bn_dh dhA = {NULL};
    struct bn_dh dhB = {NULL};
    BIGNUM *bnP, *bnG = NULL;

    // Init rand()
    srand(time(NULL));

    // First, we need to generate p, a prime number
    bnP = generateP();
    if (bnP == NULL)
        exit(-1);

    // Generate g, g must be between 1 and p - 1
    nP = BN_bn2dec(bnP);
    bnG = generateG(nP); 

    if (bnG == NULL){
        cleanDH(&dhA);
        cleanDH(&dhB);
        BN_free(bnP);
        exit(-1);
    }
    OPENSSL_free(nP);

    dumpBN("p", bnP);
    dumpBN("g", bnG);

    // We generate the private key for both A and B
    dhA.bnPrivateKey = generatePrivateKey();

    if (dhA.bnPrivateKey == NULL){
        cleanDH(&dhA);
        cleanDH(&dhB);
        BN_free(bnP);
        BN_free(bnG);
        exit(-1);
    }

    // Generate private keyB
    dhB.bnPrivateKey = generatePrivateKey();
    if (dhB.bnPrivateKey == NULL){
        cleanDH(&dhA);
        cleanDH(&dhB);
        BN_free(bnP);
        BN_free(bnG);
        exit(-1);
    }

    dumpBN("Private keyA", dhA.bnPrivateKey);
    dumpBN("Private keyB", dhB.bnPrivateKey);

    // We are going to combine keys
    combineKey(&dhA.bnPublicKey, bnG, dhA.bnPrivateKey, bnP, dhA.ctx);
    dumpBN("Public keyA", dhA.bnPublicKey);

    combineKey(&dhB.bnPublicKey, bnG, dhB.bnPrivateKey, bnP, dhB.ctx);
    dumpBN("Public keyB", dhB.bnPublicKey);

    // Send public key to the other
    dhA.bnPublicKeyPeer = BN_new();
    dhB.bnPublicKeyPeer = BN_new();

    if(BN_copy(dhA.bnPublicKeyPeer, dhB.bnPublicKey) == NULL) {
        printf("Failed to copy the public key\n");
        cleanDH(&dhA);
        cleanDH(&dhB);
        BN_free(bnP);
        BN_free(bnG);
    }
    if(BN_copy(dhB.bnPublicKeyPeer, dhA.bnPublicKey) == NULL){
        printf("Failed to copy the public key\n");
        cleanDH(&dhA);
        cleanDH(&dhB);
        BN_free(bnP);
        BN_free(bnG);
    }
      
    // For the ephemeral shared key: public keyB ** private keyA % p
    combineKey(&dhA.bnPublicEphemeralKey, dhA.bnPublicKeyPeer, dhA.bnPrivateKey, bnP, dhA.ctx);
    dumpBN("Public ephemeral keyA", dhA.bnPublicEphemeralKey);
    combineKey(&dhB.bnPublicEphemeralKey, dhB.bnPublicKeyPeer, dhB.bnPrivateKey, bnP, dhB.ctx);
    dumpBN("Public ephemeral keyB", dhB.bnPublicEphemeralKey);
    
    cleanDH(&dhA);
    cleanDH(&dhB);
    BN_free(bnP);
    BN_free(bnG);
    return 0;
}