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

#include "vibecrypto.h"
#include "environment.h"
#include "base.h"
#include "libctx.h"
#include "paramobj.h"
#include "dsa.h"
#include "mpint.h"
#include "prime.h"
#include "surrender.h"
#include "errorctx.h"

int DSAGenerateParameters (
   VtParameterObject paramObj,
   VtRandomObject random
   )
{
  int status;
  unsigned int seedLen, counter, callNumber;
  VoltParameterObject *obj = (VoltParameterObject *)paramObj;
  VoltDsaParamGenCtx *dsaParamGenCtx =
    (VoltDsaParamGenCtx *)(obj->localGenerateCtx);
  VtRandomObject randomToUse;
  VoltMpIntCtx *mpCtx = obj->mpCtx;
  VoltLibCtx *libCtx = (VoltLibCtx *)(obj->voltObject.libraryCtx);
  VoltMpInt *primeP = (VoltMpInt *)0;
  VoltMpInt *subprimeQ = (VoltMpInt *)0;
  VoltMpInt *baseG = (VoltMpInt *)0;
  VoltMpInt *hVal = (VoltMpInt *)0;
  VoltSurrenderCtx *surrCtx = (VoltSurrenderCtx *)0;
  unsigned char SEED[20];
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  do
  {
    /* If the object already contains parameters, error.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_PARAM_OBJ;
    if (obj->paramData != (Pointer)0)
      break;

    /* If there's a surrender ctx, call the Surrender function.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_GET_OBJECT_SURR_CTX (surrCtx, obj);
    VOLT_CALL_SURRENDER (surrCtx, VT_SURRENDER_FNCT_DSA_PARAM_GEN, 0, 1)

    /* 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;

    /* Create the MPInt's.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &primeP);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &subprimeQ);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &baseG);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &hVal);
    if (status != 0)
      break;

    /* General purpose PQG gen to share with DH.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    callNumber = 1;
    status = VoltGeneratePQGFips186 (
      mpCtx, surrCtx, VT_SURRENDER_FNCT_DSA_PARAM_GEN, &callNumber,
      dsaParamGenCtx->primeSizeBits, 160, randomToUse,
      primeP, subprimeQ, baseG, hVal, SEED, &seedLen, &counter);
    if (status != 0)
      break;

    /* If that succeeded, set the object with the values.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = AddDSAParametersMpInt (
      obj, primeP, subprimeQ, baseG, hVal, SEED, seedLen, counter);
    if (status != 0)
      break;

    /* If there's a surrender ctx, call it for the last time.
     */
    VOLT_CALL_SURRENDER (surrCtx, VT_SURRENDER_FNCT_DSA_PARAM_GEN, 0, 0)

  } while (0);

  mpCtx->DestroyMpInt (&hVal);
  Z2Memset (SEED, 0, sizeof (SEED));

  /* If successful, we're done.
   */
  if (status == 0)
    return (0);

  /* If error, clean up some things.
   */
  if (mpCtx != (VoltMpIntCtx *)0)
  {
    mpCtx->DestroyMpInt (&primeP);
    mpCtx->DestroyMpInt (&subprimeQ);
    mpCtx->DestroyMpInt (&baseG);
  }

  VOLT_LOG_ERROR (
    (VtLibCtx)libCtx, status, errorType, fnctLine,
    "DSAGenerateParameters", (char *)0)

  return (status);
}

int VoltGeneratePQGFips186 (
   VoltMpIntCtx *mpCtx,
   VoltSurrenderCtx *surrCtx,
   unsigned int surrFlag,
   unsigned int *callNumber,
   unsigned int primeSizeBits,
   unsigned int subprimeSizeBits,
   VtRandomObject random,
   VoltMpInt *primeP,
   VoltMpInt *subprimeQ,
   VoltMpInt *baseG,
   VoltMpInt *hVal,
   unsigned char *SEED,
   unsigned int *seedLen,
   unsigned int *count
   )
{
  int status;
  unsigned int bufSize, counter, offset, kVal, nVal, totalLen, digestLen;
  unsigned int bufOffset, bitMask, bitSet, isPrime, msbit, callNum;
  unsigned char *buffer = (unsigned char *)0;
  unsigned char *seedCopy, *digest;
  VoltLibCtx *libCtx = (VoltLibCtx *)(mpCtx->voltObject.libraryCtx);
  VtAlgorithmObject sha1 = (VtAlgorithmObject)0;
  VoltMpInt *remainder = (VoltMpInt *)0;
  VoltMpInt *expo = (VoltMpInt *)0;
  VOLT_DECLARE_ERROR_TYPE (errorType)
  VOLT_DECLARE_FNCT_LINE (fnctLine)

  callNum = 1;
  if (callNumber != (unsigned int *)0)
    callNum = *callNumber;

  do
  {
    /* For now, the toolkit supports only 1024-bit DSA prime p and
     * 160-bit subprime q.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_INVALID_PARAM_LENGTH;
    if ( (primeSizeBits != 1024) || (subprimeSizeBits != 160) )
      break;

    /* We'll need this later on.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VtCreateAlgorithmObject (
      (VtLibCtx)libCtx, VtAlgorithmImplSHA1, (Pointer)0, &sha1);
    if (status != 0)
      break;

    /* We'll also need some masking values to make sure the appropriate
     * high order bits are set or not set depending on the prime size.
     */
    bitSet = primeSizeBits % 8;
    if (bitSet == 0)
      bitSet = 8;
    bitMask = 0xff >> (8 - bitSet);
    /* Make sure the msbit is set. Create a value to OR with the first
     * byte that will set the msbit.
     */
    bitSet = 1 << (bitSet - 1);

    /* FIPS specifies L - 1 = n*160 + b, where L is the bit length of
     * the primeP.
     */
    nVal = primeSizeBits / subprimeSizeBits;

    /* For this implementation, the seedLen is the number of bytes
     * needed to hold subprimeSizeBits.
     */
    *seedLen = (subprimeSizeBits + 7) / 8;

    /* Create the mpInt's we'll need.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &remainder);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->CreateMpInt ((Pointer)mpCtx, &expo);
    if (status != 0)
      break;

    /* Use baseG as a temp to subtract 1 later on.
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->IntToMpInt (0, 1, baseG);
    if (status != 0)
      break;

    /* After we find the subprime, we'll need a buffer to hold the
     * primeP starting point. While we're at it, allocate space for
     * a copy of SEED that we'll be able to manipulate, and a digest
     * buffer.
     */
    VOLT_SET_ERROR_TYPE (errorType, VT_ERROR_TYPE_PRIMARY)
    VOLT_SET_FNCT_LINE (fnctLine)
    status = VT_ERROR_MEMORY;
    bufSize = (primeSizeBits + 7) / 8;
    buffer = (unsigned char *)Z2Malloc (
      bufSize + *seedLen + 20, VOLT_MEMORY_SENSITIVE);
    if (buffer == (unsigned char *)0)
      break;

    seedCopy = buffer + bufSize;
    digest = seedCopy + *seedLen;

    /* This is the beginning of the total loop: find subprimeQ and
     * primeP.
     */
    VOLT_SET_ERROR_TYPE (errorType, 0)
    do
    {
      callNum++;
      VOLT_CALL_SURRENDER (surrCtx, surrFlag, 0, callNum)

      /* Steps 1 - 5: Generate the subprime. (See the comments in
       * VoltGeneratePrimeFips for more on steps 1 - 5.
       */

      VOLT_SET_FNCT_LINE (fnctLine)
      status = VoltGeneratePrimeFips (
        subprimeSizeBits, random, SEED, seedLen, subprimeQ);
      if (status != 0)
        break;

      /* Step 6: counter = 0, offset = 2
       */
      counter = 0;
      offset = 2;

      /* This is the beginning of the inner loop: find primeP based on
       * the given subprimeQ.
       */
      do
      {
        callNum++;
        VOLT_CALL_SURRENDER (surrCtx, surrFlag, 0, callNum)

        /* Step 7: Build the primeP starting point by digesting SEED along
         * with the offset and a count as many times as necessary to
         * generate the appropriate number of bytes.
         */
        Z2Memcpy (seedCopy, SEED, *seedLen);
        totalLen = bufSize;
        kVal = 0;
        bufOffset = bufSize;

        VoltAddValueToBuffer (seedCopy, *seedLen, offset - 1); 
        do
        {
          /* Compute SEED + offset + kVal.
           */
          VoltAddValueToBuffer (seedCopy, *seedLen, 1); 
          kVal++;

          /* Digest this value.
           */
          VOLT_SET_FNCT_LINE (fnctLine)
          status = VtDigestInit (sha1);
          if (status != 0)
            break;

          VOLT_SET_FNCT_LINE (fnctLine)
          status = VtDigestFinal (
            sha1, seedCopy, *seedLen, digest, 20, &digestLen);
          if (status != 0)
            break;

          if (digestLen > totalLen)
            digestLen = totalLen;

          /* "Prepend" the current digest. This is the part of step 8
           * that builds W.
           */
          bufOffset -= digestLen;
          totalLen  -= digestLen;
          Z2Memcpy (buffer + bufOffset, digest + (20 - digestLen), digestLen);

        } while (totalLen > 0);
        if (status != 0)
          break;

        /* Finish step 8, compute X.
         */

        /* The FIPS technique actually specifies that W is a value of
         * length primeSizeBits - 1, so we need to clear the most
         * significant bit of the buffer to get the appropriate W. But
         * then we're going to add that bit in to get X, which is
         * W + (2 ^ (primeSizeBits - 1)). So at this point, just make
         * sure the most significant bit is set.
         */
        buffer[0] &= (unsigned char)bitMask;
        buffer[0] |= (unsigned char)bitSet;

        /* Step 9: c = X mod 2q, p = X - (c - 1)
         */
        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->OctetStringToMpInt (0, buffer, bufSize, primeP);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->Add (subprimeQ, subprimeQ, expo);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->ModReduce (primeP, expo, remainder);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->Subtract (remainder, baseG, remainder);
        if (status != 0)
          break;

        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->Subtract (primeP, remainder, primeP);
        if (status != 0)
          break;

        /* Step 10: is primeP < 2 ^ (primeSizeBits - 1).
         * In other words, is the most significant bit still set?
         */
        VOLT_SET_FNCT_LINE (fnctLine)
        status = mpCtx->GetBit (primeP, primeSizeBits - 1, &msbit);
        if (status == VT_ERROR_MP_INT_RANGE)
          status = 0;
        if (status != 0)
          break;

        /* Continuing Step 10, if the bit is set, move on to step 11.
         */
        if (msbit != 0)
        {
          /* Step 11: Is primeP actually prime?
           */
          VOLT_SET_FNCT_LINE (fnctLine)
          status = VoltRabinMillerTest (
            primeP, primeSizeBits, VOLT_RABIN_MILLER_COUNT_186_2,
            random, &isPrime);
          if (status != 0)
            break;

          /* Step 12: If the value is prime, move on to step 15.
           */
          if (isPrime != 0)
            break;
        }

        /* Step 13: The candidate "failed", either the msbit was not set
         * after the subtraction or it was not prime. Update counter and
         * offset.
         */
        counter++;
        offset += (nVal + 1);

        /* Step 14: If counter >= 4096, find a new subprimeQ.
         */
        if (counter >= 4096)
          break;

        /* Step 14: If counter < 4096, try another primeP.
         */

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

      /* At this point, we exited the inner loop because either we
       * found a prime or the counter hit 4096. If isPrime is not 0, we
       * found a prime.
       */
      if (isPrime != 0)
        break;

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

    /* If we reach this point, we have a subprimeQ and a primeP. The
     * SEED buffer contains the seed we used and counter is set to the
     * value we'll need to return at the address given by the count arg.
     */
    *count = counter;

    /* Now that we have p and q, compute g. It is
     *    g = h ^ ((p-1)/q)
     * where h is a random number.
     */

    /* Find (p-1) / q
     */
    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->MpIntToMpInt (primeP, expo);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->SetBit (expo, 0, 0);
    if (status != 0)
      break;

    VOLT_SET_FNCT_LINE (fnctLine)
    status = mpCtx->Divide (expo, subprimeQ, expo, remainder);
    if (status != 0)
      break;

    /* Generate a random h < p.
     */
    do
    {
      VOLT_SET_FNCT_LINE (fnctLine)
      status = VtGenerateRandomBytes (random, buffer, bufSize);
      if (status != 0)
        break;
      buffer[0] = 2; 

      VOLT_SET_FNCT_LINE (fnctLine)
      status = mpCtx->OctetStringToMpInt (0, buffer, /**/1/*/bufSize/*/, hVal); /* FIX */
      if (status != 0)
        break;

      /* compute g
       */
      VOLT_SET_FNCT_LINE (fnctLine)
      status = mpCtx->ModExp (hVal, expo, primeP, baseG);
      if (status != 0)
        break;

      /* Make sure baseG is not a small number.
       */
      VOLT_SET_FNCT_LINE (fnctLine)
      status = mpCtx->MpIntToInt (baseG, 1, &msbit);
      /* If status is 0, the g computed fit into an unsigned int,
       * that's small, try again with new random material.
       */
      if (status == 0)
        continue;

      /* If status is VT_ERROR_MP_INT_RANGE, the number is big enough, we're
       * done. Make sure status is 0.
       * If status was some other error, pass it along.
       */
      if (status == VT_ERROR_MP_INT_RANGE)
        status = 0;

      break;
    } while (1);

  } while (0);

  if (callNumber != (unsigned int *)0)
    *callNumber = callNum;
  if (buffer != (unsigned char *)0)
    Z2Free (buffer);

  mpCtx->DestroyMpInt (&remainder);
  mpCtx->DestroyMpInt (&expo);
  VtDestroyAlgorithmObject (&sha1);

  VOLT_LOG_ERROR_COMPARE (
    status, (VtLibCtx)libCtx, status, errorType, fnctLine,
    "VoltGeneratePQGFips186", (char *)0)

  return (status);
}