/* Copyright 2005-2006, Voltage Security, all rights reserved.
 */
#include "vibecrypto.h"
#include "environment.h"
#include "base.h"
#include "libctx.h"
#include "keyobj.h"
#include "dh.h"
#include "dsa.h"
#include "mpint.h"
#include "random.h"
#include "fipsmodule.h"
#include "errorctx.h"
#include "surrender.h"

int DHGenerateKeyPair (
   VtKeyObject priKey,
   VtKeyObject pubKey,
   VtRandomObject random
   )
{
  int status, cmpResult;
  unsigned int pubValLen, qLen;
  VoltKeyObject *priObj = (VoltKeyObject *)priKey;
  VoltKeyObject *pubObj = (VoltKeyObject *)pubKey;
  VoltLibCtx *libCtx = (VoltLibCtx *)(priObj->voltObject.libraryCtx);
  VtRandomObject rand = (VtRandomObject)0;
  VtRandomObject randomToUse;
  VoltDHKeyGenCtx *dhGenCtx = (VoltDHKeyGenCtx *)(priObj->localGenerateCtx);
  VtParameterObject paramGen = (VtParameterObject)0;
  VtFips186PrngInfo randInfo;
  VoltSurrenderCtx *surrCtx = (VoltSurrenderCtx *)0;
  VtDHPubKeyInfo pubKeyInfo;
  VtDHPriKeyInfo priKeyInfo;
  unsigned char *pubVal = (unsigned char *)0;
  unsigned char *priVal = (unsigned char *)0;
  unsigned char xkey[VOLT_DH_XKEY_LEN];
  unsigned char xseed[VOLT_DH_XSEED_LEN];
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* If there's no random object, get one from the libCtx.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_NO_RANDOM_OBJECT;
    randomToUse = random;
    if (random == (VtRandomObject)0)
    {
      randomToUse = (VtRandomObject)VoltGetLibCtxInfo (
        (VtLibCtx)libCtx, VOLT_LIB_CTX_INFO_TYPE_RANDOM);

      if (randomToUse == (VtRandomObject)0)
        break;
    }

    /* Make sure the random object is valid.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_RANDOM_OBJ;
    if (VOLT_OBJECT_TYPE_NOT_EQUAL (randomToUse, VOLT_OBJECT_TYPE_RANDOM))
      break;

    /* If there's a surrender ctx, call the Surrender function.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_GET_OBJECT_SURR_CTX (surrCtx, priObj);
    /* If the surrenderCallback was not in the private key, is it in
     * the public key?
     */
    if (surrCtx == (VoltSurrenderCtx *)0)
    {
      VOLT_GET_OBJECT_SURR_CTX (surrCtx, pubObj);
    }
    VOLT_CALL_SURRENDER (surrCtx, VT_SURRENDER_FNCT_DH_KEY_GEN, 3, 1)

    /* Generate XKEY and XSEED.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGenerateRandomBytes (randomToUse, xkey, VOLT_DH_XKEY_LEN);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGenerateRandomBytes (randomToUse, xseed, VOLT_DH_XSEED_LEN);
    if (status != 0)
      break;

    /* Build a random object that will generate the random x in a FIPS
     * prescibed manner.
     */
    randInfo.variation = FIPS_186_PRNG_3_1_CERTIFY;
    randInfo.mpCtx = (VtMpIntCtx)(priObj->mpCtx);
    randInfo.primeQ.data = (unsigned char *)0;
    randInfo.primeQ.len = 0;
    randInfo.XKEY.data = xkey;
    randInfo.XKEY.len = VOLT_DH_XKEY_LEN;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateRandomObject (
      priObj->voltObject.libraryCtx, VtRandomImplFips186Prng,
      (Pointer)&randInfo, &rand);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSeedRandom (rand, xseed, VOLT_DH_XSEED_LEN);
    if (status != 0)
      break;

    /* Generate the random x. It will be > 1 and < q - 1.
     */
    qLen = 20;
    if (dhGenCtx->subprimeQ.len != 0)
      qLen = dhGenCtx->subprimeQ.len;

    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    priVal = (unsigned char *)Z2Malloc (qLen, VOLT_MEMORY_SENSITIVE);
    if (priVal == (unsigned char *)0)
      break;

    do
    {
      VOLT_SET_ERROR_TYPE (errorType, 0)
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VtGenerateRandomBytes (rand, priVal, qLen);
      if (status != 0)
        break;

      /* If there is no q, this x will do.
       */
      if (dhGenCtx->subprimeQ.data == (unsigned char *)0)
        break;

      /* If there is a q, make sure the generated value is less than
       * the q.
       */
      cmpResult = Z2Memcmp (priVal, dhGenCtx->subprimeQ.data, qLen);
      if (cmpResult < 0)
        break;

    } while (1);
    if (status != 0)
      break;

    VOLT_CALL_SURRENDER (surrCtx, VT_SURRENDER_FNCT_DH_KEY_GEN, 3, 2)

    /* Generate the public value from the private value.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VoltGeneratePubValAlloc (
      libCtx, dhGenCtx->mpCtx, &(dhGenCtx->primeP), &(dhGenCtx->baseG),
      priVal, qLen, &pubVal, &pubValLen);
    if (status != 0)
      break;

    /* Fill in the keyInfo struct with the data.
     */
    pubKeyInfo.primeP.data = dhGenCtx->primeP.data;
    pubKeyInfo.primeP.len = dhGenCtx->primeP.len;
    pubKeyInfo.subprimeQ.data = dhGenCtx->subprimeQ.data;
    pubKeyInfo.subprimeQ.len = dhGenCtx->subprimeQ.len;
    pubKeyInfo.baseG.data = dhGenCtx->baseG.data;
    pubKeyInfo.baseG.len = dhGenCtx->baseG.len;
    pubKeyInfo.pubValY.data = pubVal;
    pubKeyInfo.pubValY.len = pubValLen;

    /* Set the object.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (pubKey, VtKeyParamDHPublic, (Pointer)&pubKeyInfo);
    if (status != 0)
      break;

    /* Fill in the keyInfo struct with the data.
     */
    priKeyInfo.primeP.data = dhGenCtx->primeP.data;
    priKeyInfo.primeP.len = dhGenCtx->primeP.len;
    priKeyInfo.subprimeQ.data = dhGenCtx->subprimeQ.data;
    priKeyInfo.subprimeQ.len = dhGenCtx->subprimeQ.len;
    priKeyInfo.baseG.data = dhGenCtx->baseG.data;
    priKeyInfo.baseG.len = dhGenCtx->baseG.len;
    priKeyInfo.pubValY.data = pubVal;
    priKeyInfo.pubValY.len = pubValLen;
    priKeyInfo.priValX.data = priVal;
    priKeyInfo.priValX.len = 20;

    /* Set the object.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (
      priKey, VtKeyParamDHPrivate, (Pointer)&priKeyInfo);
    if (status != 0)
      break;

    status = VoltTestDHKeyPair (libCtx, surrCtx, pubKey, priKey, random);

    /* If we're building this for the FIPS shared library, we want to set
     * the FIPS error.
     */
#if VOLT_BUILD == VOLT_BUILD_FIPS_SHARED
    if (status == VT_ERROR_UNMATCHED_KEY_PAIR)
    {
      VoltSetFipsError (VT_ERROR_FIPS_DH_PAIR_GEN);
      status = VT_ERROR_FIPS_DH_PAIR_GEN;
    }
#endif  /* VOLT_BUILD == VOLT_BUILD_FIPS_SHARED */
    if (status != 0)
      break;

    VOLT_CALL_SURRENDER (surrCtx, VT_SURRENDER_FNCT_DH_KEY_GEN, 3, 3)

  } while (0);

  VtDestroyParameterObject (&paramGen);
  VtDestroyRandomObject (&rand);

  if (priVal != (unsigned char *)0)
    Z2Free (priVal);
  if (pubVal != (unsigned char *)0)
    Z2Free (pubVal);

  Z2Memset (xkey, 0, VOLT_DH_XKEY_LEN);
  Z2Memset (xseed, 0, VOLT_DH_XSEED_LEN);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, priKey, status, 0, errorType,
    (char *)0, "DHGenerateKeyPair", fnctLine, (char *)0)

  return (status);
}

int VoltTestDHKeyPair (
   VoltLibCtx *libCtx,
   VoltSurrenderCtx *surrCtx,
   VtKeyObject pubKey,
   VtKeyObject priKey,
   VtRandomObject random
   )
{
  int status;
  VtAlgorithmObject dhAgree = (VtAlgorithmObject)0;
  VoltKeyObject *obj = (VoltKeyObject *)pubKey;
  VtKeyObject newPub = (VtKeyObject)0;
  VtKeyObject shared = (VtKeyObject)0;
  VtItem *getShared;
  VtDHPubKeyInfo *getPubInfo;
  VtDHPubKeyInfo newInfo;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Get the public key data so we have prime, subprime, and base.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGetKeyParam (
      pubKey, VtKeyParamDHPublic, (Pointer *)&getPubInfo);
    if (status != 0)
      break;

    /* Build a new public key from the params, then use the base as the
     * public value.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateKeyObject (
      (VtLibCtx)libCtx, VtKeyImplMpCtx, (Pointer)(obj->mpCtx), &newPub);
    if (status != 0)
      break;

    newInfo.primeP.data = getPubInfo->primeP.data;
    newInfo.primeP.len = getPubInfo->primeP.len;
    newInfo.subprimeQ.data = getPubInfo->subprimeQ.data;
    newInfo.subprimeQ.len = getPubInfo->subprimeQ.len;
    newInfo.baseG.data = getPubInfo->baseG.data;
    newInfo.baseG.len = getPubInfo->baseG.len;
    newInfo.pubValY.data = getPubInfo->baseG.data;
    newInfo.pubValY.len = getPubInfo->baseG.len;
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (
      newPub, VtKeyParamDHPublic, (Pointer)&newInfo);
    if (status != 0)
      break;

    /* Now generate a shared secret using the new public key and the
     * existing private key.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateKeyObject (
      (VtLibCtx)libCtx, VtKeyImplMpCtx, (Pointer)(obj->mpCtx), &shared);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateAlgorithmObject (
      (VtLibCtx)libCtx, VtAlgorithmImplDHKeyAgree, (Pointer)0, &dhAgree);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGenerateSharedSecret (dhAgree, random, newPub, priKey, shared);
    if (status != 0)
      break;

    /* Get the shared secret data, it should be the same as the
     * original private key's public value.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGetKeyParam (
      shared, VtKeyParamDHSharedSecret, (Pointer *)&getShared);
    if (status != 0)
      break;

    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_UNMATCHED_KEY_PAIR;
    if (getShared->len != getPubInfo->pubValY.len)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if (Z2Memcmp (
      getShared->data, getPubInfo->pubValY.data, getShared->len) != 0)
      break;

    /* If we get this far, the test passes.
     */
    status = 0;

  } while (0);

  VtDestroyAlgorithmObject (&dhAgree);
  VtDestroyKeyObject (&newPub);
  VtDestroyKeyObject (&shared);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, priKey, status, 0, errorType,
    (char *)0, "VoltTestDHKeyPair", fnctLine, (char *)0)

  return (status);
}