#include "SSLCert.hpp"

namespace httpsserver {

SSLCert::SSLCert(unsigned char * certData, uint16_t certLength, unsigned char * pkData, uint16_t pkLength):
  _certLength(certLength),
  _certData(certData),
  _pkLength(pkLength),
  _pkData(pkData) {

}

SSLCert::~SSLCert() {
  // TODO Auto-generated destructor stub
}


uint16_t SSLCert::getCertLength() {
  return _certLength;
}

uint16_t SSLCert::getPKLength() {
  return _pkLength;
}

unsigned char * SSLCert::getCertData() {
  return _certData;
}

unsigned char * SSLCert::getPKData() {
  return _pkData;
}

void SSLCert::setPK(unsigned char * pkData, uint16_t length) {
  _pkData = pkData;
  _pkLength = length;
}

void SSLCert::setCert(unsigned char * certData, uint16_t length) {
  _certData = certData;
  _certLength = length;
}

void SSLCert::clear() {
  for(uint16_t i = 0; i < _certLength; i++) _certData[i]=0;
  delete _certData;
  _certLength = 0;

  for(uint16_t i = 0; i < _pkLength; i++) _pkData[i] = 0;
  delete _pkData;
  _pkLength = 0;
}

#ifndef HTTPS_DISABLE_SELFSIGNING

/**
 * Function to create the key for a self-signed certificate.
 * 
 * Writes private key as DER in certCtx
 * 
 * Based on programs/pkey/gen_key.c
 */
static int gen_key(SSLCert &certCtx, SSLKeySize keySize) {

  // Initialize the entropy source
  mbedtls_entropy_context entropy;
  mbedtls_entropy_init( &entropy );

  // Initialize the RNG
  mbedtls_ctr_drbg_context ctr_drbg;
  mbedtls_ctr_drbg_init( &ctr_drbg );
  int rngRes = mbedtls_ctr_drbg_seed(
    &ctr_drbg, mbedtls_entropy_func, &entropy,
    NULL, 0
  );
  if (rngRes != 0) {
    mbedtls_entropy_free( &entropy );
    return HTTPS_SERVER_ERROR_KEYGEN_RNG;
  }

  // Initialize the private key
  mbedtls_pk_context key;
  mbedtls_pk_init( &key );
  int resPkSetup = mbedtls_pk_setup( &key, mbedtls_pk_info_from_type( MBEDTLS_PK_RSA ) );
  if ( resPkSetup != 0) {
    mbedtls_ctr_drbg_free( &ctr_drbg );
    mbedtls_entropy_free( &entropy );
    return HTTPS_SERVER_ERROR_KEYGEN_SETUP_PK;
  }

  // Actual key generation 
  int resPkGen = mbedtls_rsa_gen_key(
    mbedtls_pk_rsa( key ),
    mbedtls_ctr_drbg_random,
    &ctr_drbg,
    keySize,
    65537
  );
  if ( resPkGen != 0) {
    mbedtls_pk_free( &key );
    mbedtls_ctr_drbg_free( &ctr_drbg );
    mbedtls_entropy_free( &entropy );
    return HTTPS_SERVER_ERROR_KEYGEN_GEN_PK;
  }

  // Free the entropy source and the RNG as they are no longer needed
  mbedtls_ctr_drbg_free( &ctr_drbg );
  mbedtls_entropy_free( &entropy );

  // Allocate the space on the heap, as stack size is quite limited
  unsigned char * output_buf = new unsigned char[4096];
  if (output_buf == NULL) {
    mbedtls_pk_free( &key );
    return HTTPS_SERVER_ERROR_KEY_OUT_OF_MEM;
  }
  memset(output_buf, 0, 4096);

  // Write the key to the temporary buffer and determine its length
  int resPkWrite = mbedtls_pk_write_key_der( &key, output_buf, 4096 );
  if (resPkWrite < 0) {
    delete[] output_buf;
    mbedtls_pk_free( &key );
    return HTTPS_SERVER_ERROR_KEY_WRITE_PK;
  }
  size_t pkLength = resPkWrite;
  unsigned char *pkOffset = output_buf + sizeof(unsigned char) * 4096 - pkLength;

  // Copy the key into a new, fitting space on the heap
  unsigned char * output_pk = new unsigned char[pkLength];
  if (output_pk == NULL) {
    delete[] output_buf;
    mbedtls_pk_free( &key );
    return HTTPS_SERVER_ERROR_KEY_WRITE_PK;
  }
  memcpy(output_pk, pkOffset, pkLength);

  // Clean up the temporary buffer and clear the key context
  delete[] output_buf;
  mbedtls_pk_free( &key );

  // Set the private key in the context
  certCtx.setPK(output_pk, pkLength);

  return 0;
}

/**
 * Function to generate the X509 certificate data for a private key
 * 
 * Writes certificate in certCtx
 *
 * Based on programs/x509/cert_write.c
 */
static int cert_write(SSLCert &certCtx, std::string dn, std::string validityFrom, std::string validityTo) {
  int funcRes = 0;
  int stepRes = 0;

  mbedtls_entropy_context entropy;
  mbedtls_ctr_drbg_context ctr_drbg;
  mbedtls_pk_context key;
  mbedtls_x509write_cert crt;
  mbedtls_mpi serial;
  unsigned char * primary_buffer;
  unsigned char *certOffset;
  unsigned char * output_buffer;
  size_t certLength;

  // Make a C-friendly version of the distinguished name
  char dn_cstr[dn.length()+1];
  strcpy(dn_cstr, dn.c_str());

  // Initialize the entropy source
  mbedtls_entropy_init( &entropy );

  // Initialize the RNG
  mbedtls_ctr_drbg_init( &ctr_drbg );
  stepRes = mbedtls_ctr_drbg_seed( &ctr_drbg, mbedtls_entropy_func, &entropy, NULL, 0 );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_RNG;
    goto error_after_entropy;
  }

  mbedtls_pk_init( &key );
  stepRes = mbedtls_pk_parse_key( &key, certCtx.getPKData(), certCtx.getPKLength(), NULL, 0 );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_READKEY;
    goto error_after_key;
  }

  // Start configuring the certificate
  mbedtls_x509write_crt_init( &crt );
  // Set version and hash algorithm
  mbedtls_x509write_crt_set_version( &crt, MBEDTLS_X509_CRT_VERSION_3 );
  mbedtls_x509write_crt_set_md_alg( &crt, MBEDTLS_MD_SHA256 );

  // Set the keys (same key as we self-sign)
  mbedtls_x509write_crt_set_subject_key( &crt, &key );
  mbedtls_x509write_crt_set_issuer_key( &crt, &key );

  // Set issuer and subject (same, as we self-sign)
  stepRes = mbedtls_x509write_crt_set_subject_name( &crt, dn_cstr );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_NAME;
    goto error_after_cert;
  }
  stepRes = mbedtls_x509write_crt_set_issuer_name( &crt, dn_cstr );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_NAME;
    goto error_after_cert;
  }

  // Set the validity of the certificate. At the moment, it's fixed from 2019 to end of 2029.
  stepRes = mbedtls_x509write_crt_set_validity( &crt, validityFrom.c_str(), validityTo.c_str());
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_VALIDITY;
    goto error_after_cert;
  }

  // Make this a CA certificate
  stepRes = mbedtls_x509write_crt_set_basic_constraints( &crt, 1, 0 );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_VALIDITY;
    goto error_after_cert;
  }

  // generate random serial number
  mbedtls_mpi_init( &serial );
  stepRes = mbedtls_mpi_fill_random( &serial, 10, mbedtls_ctr_drbg_random, &ctr_drbg );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_SERIAL;
    goto error_after_cert_serial;
  }
  stepRes = mbedtls_x509write_crt_set_serial( &crt, &serial );
  if (stepRes != 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_SERIAL;
    goto error_after_cert_serial;
  }

  // Create buffer to write the certificate
  primary_buffer = new unsigned char[4096];
  if (primary_buffer == NULL) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_OUT_OF_MEM;
    goto error_after_cert_serial;
  }

  // Write the actual certificate
  stepRes = mbedtls_x509write_crt_der(&crt, primary_buffer, 4096, mbedtls_ctr_drbg_random, &ctr_drbg );
  if (stepRes < 0) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_WRITE;
    goto error_after_primary_buffer;
  }

  // Create a matching buffer
  certLength = stepRes;
  certOffset = primary_buffer + sizeof(unsigned char) * 4096 - certLength;

  // Copy the cert into a new, fitting space on the heap
  output_buffer = new unsigned char[certLength];
  if (output_buffer == NULL) {
    funcRes = HTTPS_SERVER_ERROR_CERTGEN_OUT_OF_MEM;
    goto error_after_primary_buffer;
  }
  memcpy(output_buffer, certOffset, certLength);

  // Configure the cert in the context
  certCtx.setCert(output_buffer, certLength);

  // Run through the cleanup process
error_after_primary_buffer:
  delete[] primary_buffer;

error_after_cert_serial:
  mbedtls_mpi_free( &serial );

error_after_cert:
  mbedtls_x509write_crt_free( &crt );

error_after_key:
  mbedtls_pk_free(&key);

error_after_entropy:
  mbedtls_ctr_drbg_free( &ctr_drbg );
  mbedtls_entropy_free( &entropy );
  return funcRes;
}

int createSelfSignedCert(SSLCert &certCtx, SSLKeySize keySize, std::string dn, std::string validFrom, std::string validUntil) {

  // Add the private key
  int keyRes = gen_key(certCtx, keySize);
  if (keyRes != 0) {
    // Key-generation failed, return the failure code
    return keyRes;
  }

  // Add the self-signed certificate
  int certRes = cert_write(certCtx, dn, validFrom, validUntil);
  if (certRes != 0) {
    // Cert writing failed, reset the pk and return failure code
    certCtx.setPK(NULL, 0);
    return certRes;
  }

  // If all went well, return 0
  return 0;
}

#endif // !HTTPS_DISABLE_SELFSIGNING

} /* namespace httpsserver */