/* Copyright 2005-2006, Voltage Security, all rights reserved.
 */

#include "vibecrypto.h"
#include "environment.h"
#include "base.h"
#include "libctx.h"
#include "algobj.h"
#include "keyobj.h"
#include "keyagree.h"
#include "dh.h"
#include "errorctx.h"

/* Gets the shared secret data out of a key object.
 *
 * @param object The object from which the data is to be extracted.
 * @param getInfo The address where the function will deposit the
 * pointer to the info.
 * @return an int, 0 if the function completed successfully or a
 * non-zero error code.
 */
static int VOLT_CALLING_CONV VoltKeyGetDHSharedSecret VOLT_PROTO_LIST ((
   VtKeyObject object,
   Pointer *getInfo
));

int VtAlgorithmImplDHKeyAgree (
   VtAlgorithmObject *object,
   Pointer info,
   unsigned int flag
   )
{
  int status;
  unsigned int bufferSize;
  VoltAlgorithmObject *obj = (VoltAlgorithmObject *)(*object);
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  unsigned char *buffer = (unsigned char *)0;
  VoltKeyAgreeClassCtx *ctx;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Check the flag, it should be VOLT_ALG_SET_TYPE_FLAG.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_TYPE;
    if (flag != VOLT_ALG_SET_TYPE_FLAG)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_ASSOCIATED_INFO;
    if (info != (Pointer)0)
      break;

    /* Allocate space for the Key Agree Class Ctx.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    bufferSize = sizeof (VoltKeyAgreeClassCtx);
    buffer = (unsigned char *)Z3Malloc (bufferSize);
    if (buffer == (unsigned char *)0)
      break;
    Z2Memset (buffer, 0, bufferSize);

    ctx = (VoltKeyAgreeClassCtx *)buffer;
    ctx->algorithmImpl = VtAlgorithmImplDHKeyAgree;
    ctx->algorithmImplInfo = (Pointer)0;
    ctx->GenerateSharedSecret = VoltDHGenerateSharedSecret;

    obj->algClass = VOLT_CLASS_KEY_AGREE;
    obj->classCtx = (Pointer)ctx;
    obj->ClassCtxDestroy = VoltSimpleCtxDestroy;

    status = 0;

  } while (0);

  if (status == 0)
    return (0);

  /* If there was an error, destroy what we created but did not return.
   */
  if (buffer != (unsigned char *)0)
    Z2Free (buffer);

  obj->state = VOLT_STATE_ERROR;

  VOLT_LOG_ERROR_INFO (
    0, *object, status, 0, errorType,
    (char *)0, "VtAlgorithmImplDHKeyAgree", fnctLine, (char *)0)

  return (status);
}

int VoltDHGenerateSharedSecret (
   VoltAlgorithmObject *obj,
   VtKeyObject otherPartyPubKey,
   VtKeyObject myPriKey,
   VtKeyObject sharedSecret,
   VtRandomObject random
   )
{
  int status, cmpResult;
  unsigned int theSecretLen, sign;
  VoltKeyObject *pubKey = (VoltKeyObject *)otherPartyPubKey;
  VoltKeyObject *priKey = (VoltKeyObject *)myPriKey;
  VoltKeyObject *secret = (VoltKeyObject *)sharedSecret;
  VoltMpIntCtx *mpCtx = secret->mpCtx;
  VoltLibCtx *libCtx = (VoltLibCtx *)(secret->voltObject.libraryCtx);
  VtKeyObject tempPub = (VtKeyObject)0;
  VtKeyObject tempPri = (VtKeyObject)0;
  VoltMpInt *result = (VoltMpInt *)0;
  unsigned char *theSecret = (unsigned char *)0;
  VoltDHPublicKey *pubKeyData;
  VtDHPubKeyInfo *pubKeyInfo;
  VoltDHPrivateKey *priKeyData;
  VtDHPriKeyInfo *priKeyInfo;
  VtItem secretItem;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Make sure the keys are DH.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_KEY_OBJ;
    if ((pubKey->keyType & VOLT_KEY_TYPE_MASK_ASYM_ALG) != VOLT_KEY_ALG_DH)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if ((pubKey->keyType & VOLT_KEY_TYPE_PUBLIC) == 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if (pubKey->mpCtx == (VoltMpIntCtx *)0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if ((priKey->keyType & VOLT_KEY_TYPE_MASK_ASYM_ALG) != VOLT_KEY_ALG_DH)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if ((priKey->keyType & VOLT_KEY_TYPE_PRIVATE) == 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if (priKey->mpCtx == (VoltMpIntCtx *)0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if (secret->mpCtx == (VoltMpIntCtx *)0)
      break;

    /* Get the key data in a form we want and make sure the mpCtx in
     * pub and pri matches.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGetKeyParam (
      otherPartyPubKey, VtKeyParamDHPublic, (Pointer *)&pubKeyInfo);
    if (status == VT_ERROR_GET_INFO_UNAVAILABLE)
      status = VT_ERROR_INVALID_KEY_OBJ;
    if (status != 0)
      break;

    /* Create a key object we know will possess the key data the way
     * we want it.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateKeyObject (
      secret->voltObject.libraryCtx, VtKeyImplMpCtx,
      (Pointer)(secret->mpCtx), &tempPub);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (
      tempPub, VtKeyParamDHPublic, (Pointer)pubKeyInfo);
    if (status != 0)
      break;

    pubKey = (VoltKeyObject *)tempPub;
    pubKeyData = (VoltDHPublicKey *)(pubKey->keyData);

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGetKeyParam (
      myPriKey, VtKeyParamDHPrivate, (Pointer *)&priKeyInfo);
    if (status == VT_ERROR_GET_INFO_UNAVAILABLE)
      status = VT_ERROR_INVALID_KEY_OBJ;
    if (status != 0)
      break;

    /* Create a key object we know will possess the key data the way
     * we want it.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateKeyObject (
      secret->voltObject.libraryCtx, VtKeyImplMpCtx,
      (Pointer)(secret->mpCtx), &tempPri);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (
      tempPri, VtKeyParamDHPrivate, (Pointer)priKeyInfo);
    if (status != 0)
      break;

    priKey = (VoltKeyObject *)tempPri;
    priKeyData = (VoltDHPrivateKey *)(priKey->keyData);

    /* Make sure the pub key's base and prime are the same as the pri
     * key's base and prime.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->Compare (
      pubKeyData->primeP, priKeyData->primeP, &cmpResult);
    if (status != 0)
      break;

    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_INPUT;
    if (cmpResult != 0)
      break;

    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->Compare (
      pubKeyData->baseG, priKeyData->baseG, &cmpResult);
    if (status != 0)
      break;

    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_INPUT;
    if (cmpResult != 0)
      break;

    /* To generate the shared secret, find
     * (otherPub ^ myPri) mod prime. We'll need a result mpInt.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &result);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->ModExp (
      pubKeyData->pubValY, priKeyData->priValX, pubKeyData->primeP,
      result);
    if (status != 0)
      break;

    /* Get the result as a byte array. First, how long.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->MpIntToOctetString (
      result, &sign, (unsigned char *)0, 0, &theSecretLen);
    if (status == 0)
      status = VT_ERROR_GENERAL;
    if (status != VT_ERROR_BUFFER_TOO_SMALL)
      break;

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

    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->MpIntToOctetString (
      result, &sign, theSecret, theSecretLen, &theSecretLen);
    if (status != 0)
      break;

    /* Set the secret key object with the shared secret.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    secretItem.data = theSecret;
    secretItem.len = theSecretLen;
    status = VtSetKeyParam (
      sharedSecret, VtKeyParamDHSharedSecret, (Pointer)&secretItem);

  } while (0);

  if (theSecret != (unsigned char *)0)
    Z2Free (theSecret);

  if (result != (VoltMpInt *)0)
    mpCtx->DestroyMpInt (&result);

  VtDestroyKeyObject (&tempPub);
  VtDestroyKeyObject (&tempPri);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, sharedSecret, status, 0, errorType,
    (char *)0, "VtSharedSecretGenDH", fnctLine, (char *)0)

  return (status);
}

int VtKeyParamDHSharedSecret (
   VtKeyObject object,
   Pointer info,
   unsigned int flag
   )
{
  int status;
  VoltKeyObject *obj = (VoltKeyObject *)object;
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  VtItem *keyInfo = (VtItem *)info;
  VtItem *localCopy = (VtItem *)0;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    if (flag == VOLT_KEY_GET_TYPE_FLAG)
    {
      VOLT_SET_ERROR_TYPE (errorType, 0)
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VoltKeyGetDHSharedSecret (object, (Pointer *)info);
      break;
    }

    /* Check the flag, it should be VOLT_KEY_SET_TYPE_FLAG.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_TYPE;
    if (flag != VOLT_KEY_SET_TYPE_FLAG)
      break;

    /* The associated info should be a pointer to a VtItem.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_ASSOCIATED_INFO;
    if (info == (Pointer)0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    if ( (keyInfo->data == (unsigned char *)0) || (keyInfo->len == 0) )
      break;

    /* Allocate space for a VtItem and the data.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    localCopy = (VtItem *)Z2Malloc (
      sizeof (VtItem) + keyInfo->len, VOLT_MEMORY_SENSITIVE);
    if (localCopy == (VtItem *)0)
      break;

    localCopy->data = ((unsigned char *)localCopy) + sizeof (VtItem);
    Z2Memcpy (localCopy->data, keyInfo->data, keyInfo->len);
    localCopy->len = keyInfo->len;

    obj->voltObject.CloneObject = VoltCloneDHSharedSecret;
    obj->keyData = (Pointer)localCopy;
    obj->KeyDataDestroy = VoltSimpleCtxDestroy;
    obj->keyType =
      VOLT_KEY_ALG_DH | VOLT_KEY_TYPE_AGREE | VOLT_KEY_TYPE_SHARED_SECRET |
      VOLT_KEY_TYPE_DATA;

    status = 0;

  } while (0);

  /* If everything worked, return 0.
   */
  if (status == 0)
    return (0);

  /* If something went wrong, indicate that this object is not usable.
   */
  obj->keyType = 0;

  if (localCopy != (VtItem *)0)
    Z2Free (localCopy);

  VOLT_LOG_ERROR_INFO (
    0, obj, status, 0, errorType,
    (char *)0, "VtKeyParamDHSharedSecret", fnctLine, (char *)0)

  return (status);
}

static int VoltKeyGetDHSharedSecret (
   VtKeyObject object,
   Pointer *getInfo
   )
{
  int status;
  VoltKeyObject *obj = (VoltKeyObject *)object;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* Is there data?
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_GET_INFO_UNAVAILABLE;
    if (obj->keyData == (Pointer)0)
      break;

    /* Is the algorithm DH?
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    if ((obj->keyType & VOLT_KEY_TYPE_MASK_ASYM_ALG) != VOLT_KEY_ALG_DH)
      break;

    /* Check the keyType in the object, if must contain
     * VOLT_KEY_TYPE_SHARED_SECRET.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    if ((obj->keyType & VOLT_KEY_TYPE_SHARED_SECRET) == 0)
      break;

    /* Is the key in data form?
     */
    if ((obj->keyType & VOLT_KEY_TYPE_MASK_DATA) != VOLT_KEY_TYPE_DATA)
    {
      /* The data is not available, does the object have a GetData
       * function?
       */
      VOLT_SET_FNCT_LINE (fnctLine)
      *getInfo = (Pointer)0;
      if (obj->GetKeyData == (VGetKeyData)0)
        break;

      /* Call the Get function.
       */
      VOLT_SET_ERROR_TYPE (errorType, 0)
      VOLT_SET_FNCT_LINE (fnctLine)
      status = obj->GetKeyData ((VtKeyObject)obj, getInfo);
      break;
    }

    *getInfo = obj->keyData;
    status = 0;

  } while (0);

  VOLT_LOG_ERROR_INFO_COMPARE (
    status, 0, object, status, 0, errorType,
    (char *)0, "VoltKeyGetDHSharedSecret", fnctLine, (char *)0)

  return (status);
}

int VoltCloneDHSharedSecret (
   Pointer sourceObject,
   Pointer *destObject
   )
{
  int status;
  VoltKeyObject *obj = (VoltKeyObject *)sourceObject;
  VtKeyObject newKey = (VtKeyObject)0;
  VtItem *getInfo;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* We know the source is an object, is it a key object?
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_KEY_OBJ;
    if (VOLT_OBJECT_TYPE_NOT_EQUAL (sourceObject, VOLT_OBJECT_TYPE_KEY))
      break;

    /* Is the object a DH shared secret?
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_KEY_OBJ;
    if ((obj->keyType & VOLT_KEY_TYPE_MASK_ASYM_ALG) != VOLT_KEY_ALG_DH)
      break;
    if ((obj->keyType & VOLT_KEY_TYPE_SHARED_SECRET) == 0)
      break;

    /* Get the DH key data out of the source object.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtGetKeyParam (
      (VtKeyObject)obj, VtKeyParamDHSharedSecret, (Pointer *)&getInfo);
    if (status != 0)
      break;

    /* Create and set a key object with this data.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateKeyObject (
      obj->voltObject.libraryCtx, VtKeyImplMpCtx, (Pointer)(obj->mpCtx),
      &newKey);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtSetKeyParam (
      newKey, VtKeyParamDHSharedSecret, (Pointer)getInfo);
    if (status != 0)
      break;

    /* If all that worked, we have our clone.
     */
    *destObject = (Pointer)newKey;

  } while (0);

  /* If successful, return 0.
   */
  if (status == 0)
    return (0);

  /* If there was an error, destroy anything we created.
   */
  VtDestroyKeyObject (&newKey);

  VOLT_LOG_ERROR_INFO (
    0, obj, status, 0, 0,
    (char *)0, "VoltCloneDHSharedSecret", fnctLine, (char *)0)

  return (status);
}