// Copyright (c) 2013-2015 IONU Security, Inc. All rights reserved
//
// DRBG class for deterministic random bit generation and predictable RSA keys

#include <pthread.h>
#include <cstring>

#include <openssl/ec.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>

#include "eyeutils.h"
#include "eyelog.h"
#include "eyedrbg.h"

using namespace std;
using namespace sequencelogic;

pthread_mutex_t ionu_drbg_mutex = PTHREAD_MUTEX_INITIALIZER;

//#define JAVASCRIPT 1  // set to enable drbg information to match with libeye.js
//#define DEBUG_RSA 1
#ifdef DEBUG_RSA
int ionuRSA_check_key(const RSA *key);
#endif

EyeDRBG::EyeDRBG (const unsigned char* seed, size_t seedlen)
{
    _data = new unsigned char[IONU_DRBG_BUFFER_LEN];
    _ctr = (unsigned char) seedlen & 0x0ff;
    if (_data) {
        std::string mdhex = sequencelogic::DigestMessage (MESSAGE_DIGEST_SHA256, seed, seedlen);
        //cout << (char*)seed << " " << mdhex << endl;
        sequencelogic::HexToBinary (mdhex.c_str(), &_data[IONU_DRBG_BUFFER_LEN - 32]);
        FillBuffer();
    }
}

EyeDRBG::~EyeDRBG ()
{
    sequencelogic::MemSet (_data, 0, IONU_DRBG_BUFFER_LEN);
    delete[] _data;
}

void EyeDRBG::FillBuffer ()
{
    _pos = 0;
    size_t loc = IONU_DRBG_BUFFER_LEN - 32;
    unsigned char hashbuf[33];
    for (size_t i = 0; i < IONU_DRBG_BUFFER_LEN; i += 32) {
        memcpy (hashbuf, &_data[loc], 32);
        hashbuf[32] = _ctr++;
        std::string mdhex = sequencelogic::DigestMessage (MESSAGE_DIGEST_SHA256, hashbuf, 33);
#ifdef JAVASCRIPT
        //cout << mdhex << endl;
#endif
        sequencelogic::HexToBinary (mdhex.c_str(), &_data[i]);
        loc = i;
    }
}

bool EyeDRBG::RandBytes (unsigned char* bytes, size_t len)
{
    if (! _data)
        return false;

    // Set mutex lock while getting bytes to keep it all thread safe
    pthread_mutex_lock (&ionu_drbg_mutex);

    size_t ba = IONU_DRBG_BUFFER_LEN - _pos;

    // We have all the bytes needed
    if (ba >= len) {
        sequencelogic::MemCpy (bytes, &_data[_pos], len);
        _pos += len;
    }
    else {
        // Copy any remaining bytes
        if (ba > 0 ) {
            sequencelogic::MemCpy (bytes, &_data[_pos], ba);
            bytes += ba;
        }
        size_t fb = len - ba; // Calculate number of bytes still needed
        // Loop until we've generated all the required bytes
        while (fb > 0) {
            FillBuffer ();
            ba = IONU_DRBG_BUFFER_LEN;
            sequencelogic::MemCpy (bytes, &_data[0], ba > fb ? fb : ba);
            if (fb > ba) {
                bytes += ba;
                fb -= IONU_DRBG_BUFFER_LEN;
            }
            else {
                _pos = fb;
                fb = 0;
            }
        }
    }
    pthread_mutex_unlock (&ionu_drbg_mutex);
    return true;
}


// Create a deterministic random BIGNUM for RSA key generation
// Code adapted from OpenSSL and Secure Programming Cookbook
bool EyeDRBG::GenerateRSAPrime (BIGNUM *rsap)
{
    // Only works for IONU_RSA_KEY_BITS size keys, prime is IONU_RSA_KEY_BITS / 2
    size_t bytes = SL_RSA_KEY_BITS / 16;
    unsigned char* buff = new unsigned char[bytes];
    if (buff == NULL)
        return false;

    size_t count = 0;
    bool good = false;
    while (!good) {
       if (! RandBytes (buff, bytes)) { // Get bytes from our DRBG
           delete[] buff;
           return false;
       }
       // Set 2 most significant bits so p * q is exactly nbits, and least to ensure odd
       buff[0] |= 0xc0;
       buff[bytes - 1] |= 1;
       // Convert to BIGNUM and see if it is prime
       BN_bin2bn (buff, static_cast<int>(bytes), rsap);
       if (BN_is_prime_fasttest_ex (rsap, BN_prime_checks, NULL, 1, NULL) == 1)
          good = true;
       count++;
    }
#ifdef JAVASCRIPT
    cout << count << endl;
#endif
    // Clear out the memory containing the prime
    sequencelogic::MemSet (buff, 0, bytes);
    delete[] buff;
    return good;
}

// Create an RSA key of IONU_RSA_KEY_BITS size using DRBG random number generator
// Code adapted from OpenSSL and Secure Programming Cookbook
bool EyeDRBG::GenerateDRBGRSAKey (RSA* rsa)
{
    if (rsa == NULL)
        return false;

    BN_CTX *ctx = BN_CTX_new();
    if (ctx == NULL)
        return false;

    bool rc = false;
    BIGNUM* r0 = NULL;
    BIGNUM* r1 = NULL;
    BIGNUM* r2 = NULL;
    BIGNUM* tmp = NULL;

    // Get some temporary BIGNUMS for scratch pad calculations
    BN_CTX_start (ctx);
    r0 = BN_CTX_get (ctx);
    r1 = BN_CTX_get (ctx);
    r2 = BN_CTX_get (ctx);
    if (r2 == NULL) goto cleanup;

    // Allocate the RSA BIGNUMs
    rsa->n = BN_new();
    rsa->d = BN_new();
    rsa->e = BN_new();
    rsa->p = BN_new();
    rsa->q = BN_new();
    rsa->dmp1 = BN_new();
    rsa->dmq1 = BN_new();
    rsa->iqmp = BN_new();

    if (! (rsa->n && rsa->d && rsa->e && rsa->p && rsa->q && rsa->dmp1 && rsa->dmq1 && rsa->iqmp))
       goto cleanup;

    BN_set_word (rsa->e, SL_RSA_PUBLIC_EXP);

    /* generate p and q */
    for (;;) {
        if (! GenerateRSAPrime (rsa->p)) goto cleanup;
        if (! BN_sub (r2, rsa->p, BN_value_one())) goto cleanup;
        if (! BN_gcd (r1, r2, rsa->e, ctx)) goto cleanup;
        if (BN_is_one (r1)) break;
    }
    for (;;) {
        if (! GenerateRSAPrime (rsa->q)) goto cleanup;
        if (BN_cmp (rsa->p, rsa->q) == 0) continue; //goto cleanup;
        if (! BN_sub (r2, rsa->q, BN_value_one())) goto cleanup;
        if (! BN_gcd (r1, r2, rsa->e, ctx)) goto cleanup;
        if (BN_is_one (r1))
            break;
    }
    // If q is greater than p, swap them for performance
    if (BN_cmp (rsa->p, rsa->q) < 0) {
        tmp = rsa->p;
        rsa->p = rsa->q;
        rsa->q = tmp;
    }
    // Calculate n
    if (! BN_mul (rsa->n, rsa->p, rsa->q, ctx)) goto cleanup;

    // Calculate d, and leave p-1 in r1, q-1 in r2
    if (! BN_sub (r1, rsa->p, BN_value_one())) goto cleanup;     // p-1 */
    if (! BN_sub (r2, rsa->q, BN_value_one())) goto cleanup;     // q-1 */
    if (! BN_mul (r0, r1, r2, ctx)) goto cleanup;                // (p-1)(q-1) */
    if (! BN_mod_inverse(rsa->d, rsa->e, r0, ctx)) goto cleanup; // d */

    /* calculate d mod (p-1) */
    if (! BN_mod (rsa->dmp1, rsa->d, r1, ctx)) goto cleanup;

    /* calculate d mod (q-1) */
    if (!BN_mod(rsa->dmq1, rsa->d, r2, ctx)) goto cleanup;

    /* calculate inverse of q mod p */
    if (!BN_mod_inverse(rsa->iqmp, rsa->q, rsa->p, ctx)) goto cleanup;

    // Success?
#ifdef DEBUG_RSA
    if (sequencelogic::RSA_check_key (rsa) != 1)
       rc = false;
    else
#endif
       rc = true;
cleanup:
    if (! rc) {
        IONUDEBUG ("GenerateRSAPrime failed");
    }
    if (ctx != NULL) {
        BN_CTX_end(ctx);
        BN_CTX_free(ctx);
    }

    return rc;
}

#ifdef DEBUG_RSA
int sequencelogic::RSA_check_key(const RSA *key)
{
    BIGNUM *i, *j, *k, *l, *m;
    BN_CTX *ctx;
    int r;
    int ret=1;
    
    i = BN_new();
    j = BN_new();
    k = BN_new();
    l = BN_new();
    m = BN_new();
    ctx = BN_CTX_new();
    if (i == NULL || j == NULL || k == NULL || l == NULL ||
        m == NULL || ctx == NULL)
        {
        ret = -1;
        printf ("ERR_R_MALLOC_FAILURE\n");
        goto err;
        }
    
    /* p prime? */
    r = BN_is_prime_ex(key->p, BN_prime_checks, NULL, NULL);
    if (r != 1)
        {
        ret = r;
        if (r != 0)
            goto err;
        printf ("RSA_R_P_NOT_PRIME\n");
        }
    
    /* q prime? */
    r = BN_is_prime_ex(key->q, BN_prime_checks, NULL, NULL);
    if (r != 1)
        {
        ret = r;
        if (r != 0)
            goto err;
        printf ("RSA_R_Q_NOT_PRIME\n");
        }
    
    /* n = p*q? */
    r = BN_mul(i, key->p, key->q, ctx);
    if (!r) { ret = -1; goto err; }
    
    if (BN_cmp(i, key->n) != 0)
        {
        ret = 0;
        printf ("RSA_R_N_DOES_NOT_EQUAL_P_Q\n");
        }
    
    /* d*e = 1  mod lcm(p-1,q-1)? */

    r = BN_sub(i, key->p, BN_value_one());
    if (!r) { ret = -1; goto err; }
    r = BN_sub(j, key->q, BN_value_one());
    if (!r) { ret = -1; goto err; }

    /* now compute k = lcm(i,j) */
    r = BN_mul(l, i, j, ctx);
    if (!r) { ret = -1; goto err; }
    r = BN_gcd(m, i, j, ctx);
    if (!r) { ret = -1; goto err; }
    r = BN_div(k, NULL, l, m, ctx); /* remainder is 0 */
    if (!r) { ret = -1; goto err; }

    r = BN_mod_mul(i, key->d, key->e, k, ctx);
    if (!r) { ret = -1; goto err; }

    if (!BN_is_one(i))
        {
        ret = 0;
        printf ("RSA_R_D_E_NOT_CONGRUENT_TO_1\n");
        }
    
    if (key->dmp1 != NULL && key->dmq1 != NULL && key->iqmp != NULL)
        {
        /* dmp1 = d mod (p-1)? */
        r = BN_sub(i, key->p, BN_value_one());
        if (!r) { ret = -1; goto err; }

        r = BN_mod(j, key->d, i, ctx);
        if (!r) { ret = -1; goto err; }

        if (BN_cmp(j, key->dmp1) != 0)
            {
            ret = 0;
            printf ("RSA_R_DMP1_NOT_CONGRUENT_TO_D\n");
            }
    
        /* dmq1 = d mod (q-1)? */    
        r = BN_sub(i, key->q, BN_value_one());
        if (!r) { ret = -1; goto err; }
    
        r = BN_mod(j, key->d, i, ctx);
        if (!r) { ret = -1; goto err; }

        if (BN_cmp(j, key->dmq1) != 0)
            {
            ret = 0;
            printf ("RSA_R_DMQ1_NOT_CONGRUENT_TO_D\n");
            }
    
        /* iqmp = q^-1 mod p? */
        if(!BN_mod_inverse(i, key->q, key->p, ctx))
            {
            ret = -1;
            goto err;
            }

        if (BN_cmp(i, key->iqmp) != 0)
            {
            ret = 0;
            printf ("RSA_R_IQMP_NOT_INVERSE_OF_Q\n");
            }
        }

 err:
    if (i != NULL) BN_free(i);
    if (j != NULL) BN_free(j);
    if (k != NULL) BN_free(k);
    if (l != NULL) BN_free(l);
    if (m != NULL) BN_free(m);
    if (ctx != NULL) BN_CTX_free(ctx);
    return (ret);
}
#endif

typedef struct ec_extra_data_st {
        struct ec_extra_data_st *next;
        void *data;
        void *(*dup_func)(void *);
        void (*free_func)(void *);
        void (*clear_free_func)(void *);
} EC_EXTRA_DATA; /* used in EC_GROUP */

struct ec_key_st {
        int version;

        EC_GROUP *group;

        EC_POINT *pub_key;
        BIGNUM   *priv_key;

        unsigned int enc_flag;
        point_conversion_form_t conv_form;

        int     references;
        int     flags;

        EC_EXTRA_DATA *method_data;
} /* EC_KEY */;


bool EyeDRBG::GenerateDRBGECKey (struct ec_key_st  *eckey)
{   
   bool rc = true;
   BN_CTX   *ctx = NULL;
   BIGNUM   *priv_key = NULL, *order = NULL;
   EC_POINT *pub_key = NULL;

   if (!eckey || !eckey->group)
      return false;

   if ((order = BN_new()) == NULL) rc = false;
   if ((ctx = BN_CTX_new()) == NULL) rc = false;

   if (eckey->priv_key == NULL)
   {
      priv_key = BN_new();
   }
   else
      priv_key = eckey->priv_key;

   if (priv_key == NULL) {
      IONUDEBUG ("EyeDRBG::GenerateDRBGECKey(): failed to get private key");
      return false;
   }

   if (!EC_GROUP_get_order (eckey->group, order, ctx)) {
      IONUDEBUG ("EyeDRBG::GenerateDRBGECKey(): get_order failed");
      return false;
   }

   int n = BN_num_bits (order);
   size_t bytes = n / 8;
   unsigned char* buff = new unsigned char[bytes];
   if (buff == NULL)
        return false;

   do {
      if (! RandBytes (buff, bytes)) { // Get bytes from our DRBG
         delete[] buff;
         return false;
      }
      // Convert to BIGNUM and see if it is prime
      BN_bin2bn (buff, static_cast<int>(bytes), priv_key);
   }
   while (BN_is_zero (priv_key));
   delete[] buff;

   if (eckey->pub_key == NULL)
   {
      pub_key = EC_POINT_new (eckey->group);
      if (pub_key == NULL)
         rc = false;
   }
   else
      pub_key = eckey->pub_key;

   if (!EC_POINT_mul (eckey->group, pub_key, priv_key, NULL, NULL, ctx))
      rc = false;

   eckey->priv_key = priv_key;
   eckey->pub_key  = pub_key;

   if (order) BN_free (order);
   if (pub_key  != NULL && eckey->pub_key  == NULL) EC_POINT_free (pub_key);
   if (priv_key != NULL && eckey->priv_key == NULL) BN_free (priv_key);
   if (ctx != NULL) BN_CTX_free (ctx);

   return rc;
}