#include <stdio.h>
#include<string.h>
#include <openssl/bn.h>

#define NBITS 512 

/*
Compute RSA public and private keys, then encrypt message using public key, then decrypt it using private key
Author: Tim Pierson, Dartmouth CS55, Winter 2021
based heavily on Kevin Du https://raw.githubusercontent.com/kevin-w-du/BookCode/master/Public_Key_Cryptography/rsa.c

compile: gcc rsa.c -lcrypto
(must have big number installed use: sudo apt-get install libssl-dev )

run: a.out
*/

void printBN(char *msg, BIGNUM * a)
{
   char * number_str = BN_bn2hex(a);
   printf("%s %s\n", msg, number_str);
   OPENSSL_free(number_str);
}


/*
Convert a string to its hexidecimal representation
based on  https://stackoverflow.com/questions/46210513/how-to-convert-a-string-to-hex-and-vice-versa-in-c
Input: msg = string to convert to hex
Output: hex representation.
*/
char *convertStringtoHex(char *msg) {
  const int len = strlen(msg);
  char *hex = malloc((len+1)*sizeof(char)*2);

  for (int i = 0, j = 0; i < len; ++i, j += 2)
      sprintf(hex + j, "%02x", msg[i] & 0xff);

  //printf("'%s' in hex is %s.\n", msg, hex);
  return hex;
}

/*
Convert a hexidecimal to a string representation
based on  https://stackoverflow.com/questions/46210513/how-to-convert-a-string-to-hex-and-vice-versa-in-c
Input: hex = hex to convert to string
Output: string representation.
*/

char *convertHexToString(char *hex) {
  const int len = strlen(hex);
  char *string = malloc((len/2+1)*sizeof(char));
    for (int i = 0, j = 0; j < len; ++i, j += 2) {
      int val[1];
      sscanf(hex + j, "%2x", val);
      string[i] = val[0];
      string[i + 1] = '\0';
    }
  return string;
}

int main () { 
  char *msgToSend = "This is a secret.";
  

  BN_CTX *ctx = BN_CTX_new();

  BIGNUM *p, *q, *n, *phi, *e, *d, *m, *c, *res;
  BIGNUM *new_m, *p_minus_one, *q_minus_one;
  p = BN_new(); q = BN_new(); n = BN_new(); e = BN_new();
  d = BN_new(); m = BN_new(); c = BN_new(); 
  res = BN_new(); phi = BN_new(); new_m = BN_new(); 
  p_minus_one = BN_new(); q_minus_one = BN_new(); 

  printf("Creating public and private keys\n");
  // Set the public key exponent e
  BN_dec2bn(&e, "65537");
  printBN("Public key (hex):",e);

  // Check whether e and (*@$\phi(n)$@*) are relatively prime. 
  do { 
    // Generate random p and q.
    BN_generate_prime_ex(p, NBITS, 1, NULL, NULL, NULL); //random prime p
    BN_generate_prime_ex(q, NBITS, 1, NULL, NULL, NULL); //random prime q
    BN_sub(p_minus_one, p, BN_value_one());     // Compute p-1
    BN_sub(q_minus_one, q, BN_value_one());     // Compute q-1
    BN_mul(n, p, q, ctx);                       // Compute n=pq
    //make sure e is relatively prime to phi
    BN_mul(phi, p_minus_one, q_minus_one, ctx); // Compute (*@$\phi(n)$@*)
    BN_gcd(res, phi, e, ctx);
  } while (!BN_is_one(res)); // try again if not relatively prime
  

  // Compute the private key exponent d, s.t. ed mod phi(n) = 1
  BN_mod_inverse(d, e, phi, ctx);                         
  printBN("Private key (hex):", d);

  // Encryption: calculate m^e mod n
  printf("\nEncrypting message: %s\n",msgToSend);
  char *msgHex = convertStringtoHex(msgToSend); //convert msgToSend to hex
  BN_hex2bn(&m, msgHex);    
  BN_mod_exp(c, m, e, n, ctx); //encrypt with public key e
  printBN("Encryption result:", c);

  // Decryption: calculate c^d mod n
  BN_mod_exp(new_m, c, d, n, ctx); //decrypt with private key d
  printBN("Decryption result:", new_m);
  char *string = BN_bn2hex(new_m);
  char *recovered_message = convertHexToString(string);
  printf("Decrypted message: %s\n",recovered_message);

  // Clear the sensitive data from the memory             
  BN_clear_free(p); BN_clear_free(q); BN_clear_free(d); 
  BN_clear_free(phi); BN_clear_free(m); BN_clear_free(new_m); 
  BN_clear_free(c); BN_clear_free(res); 
  BN_clear_free(p_minus_one); BN_clear_free(q_minus_one); 

  return 0;
}