/*
 * Copyright 2019-present MongoDB, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "../mongocrypt-crypto-private.h"
#include "../mongocrypt-private.h"
#include <stdint.h>

#ifdef MONGOCRYPT_ENABLE_CRYPTO_CNG

#include <bcrypt.h>

static BCRYPT_ALG_HANDLE _algo_sha512_hmac = 0;
static BCRYPT_ALG_HANDLE _algo_sha256_hmac = 0;
static BCRYPT_ALG_HANDLE _algo_aes256_cbc = 0;
static BCRYPT_ALG_HANDLE _algo_aes256_ecb = 0;
static DWORD _aes256_key_blob_length;
static DWORD _aes256_block_length;

static BCRYPT_ALG_HANDLE _random;

#define STATUS_SUCCESS 0

bool _native_crypto_initialized = false;

void
_native_crypto_init (void)
{
   DWORD cbOutput;
   NTSTATUS nt_status;

   /* Note, there is no mechanism for libmongocrypt to close these providers,
    * If we ever add such a mechanism, call BCryptCloseAlgorithmProvider.
    */
   nt_status = BCryptOpenAlgorithmProvider (&_algo_sha512_hmac,
                                            BCRYPT_SHA512_ALGORITHM,
                                            MS_PRIMITIVE_PROVIDER,
                                            BCRYPT_ALG_HANDLE_HMAC_FLAG);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptOpenAlgorithmProvider (&_algo_sha256_hmac,
                                            BCRYPT_SHA256_ALGORITHM,
                                            MS_PRIMITIVE_PROVIDER,
                                            BCRYPT_ALG_HANDLE_HMAC_FLAG);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptOpenAlgorithmProvider (
      &_algo_aes256_cbc, BCRYPT_AES_ALGORITHM, MS_PRIMITIVE_PROVIDER, 0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptSetProperty (
      _algo_aes256_cbc,
      BCRYPT_CHAINING_MODE,
      (PUCHAR) (BCRYPT_CHAIN_MODE_CBC),
      (ULONG) (sizeof (wchar_t) * wcslen (BCRYPT_CHAIN_MODE_CBC)),
      0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptOpenAlgorithmProvider (
      &_algo_aes256_ecb, BCRYPT_AES_ALGORITHM, MS_PRIMITIVE_PROVIDER, 0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptSetProperty (
      _algo_aes256_ecb,
      BCRYPT_CHAINING_MODE,
      (PUCHAR) (BCRYPT_CHAIN_MODE_ECB),
      (ULONG) (sizeof (wchar_t) * wcslen (BCRYPT_CHAIN_MODE_ECB)),
      0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   cbOutput = sizeof (_aes256_key_blob_length);
   nt_status = BCryptGetProperty (_algo_aes256_cbc,
                                  BCRYPT_OBJECT_LENGTH,
                                  (PUCHAR) (&_aes256_key_blob_length),
                                  cbOutput,
                                  &cbOutput,
                                  0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   cbOutput = sizeof (_aes256_block_length);
   nt_status = BCryptGetProperty (_algo_aes256_cbc,
                                  BCRYPT_BLOCK_LENGTH,
                                  (PUCHAR) (&_aes256_block_length),
                                  cbOutput,
                                  &cbOutput,
                                  0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   nt_status = BCryptOpenAlgorithmProvider (
      &_random, BCRYPT_RNG_ALGORITHM, MS_PRIMITIVE_PROVIDER, 0);
   if (nt_status != STATUS_SUCCESS) {
      return;
   }

   _native_crypto_initialized = true;
}

typedef struct {
   unsigned char *key_object;
   uint32_t key_object_length;

   BCRYPT_KEY_HANDLE key_handle;

   unsigned char *iv;
   uint32_t iv_len;
} cng_encrypt_state;

static void
_crypto_state_destroy (cng_encrypt_state *state);

static cng_encrypt_state *
_crypto_state_init (const _mongocrypt_buffer_t *key,
                    const _mongocrypt_buffer_t *iv,
                    mongocrypt_status_t *status)
{
   cng_encrypt_state *state;
   uint32_t keyBlobLength;
   unsigned char *keyBlob;
   BCRYPT_KEY_DATA_BLOB_HEADER blobHeader;
   NTSTATUS nt_status;

   BSON_ASSERT_PARAM (key);
   BSON_ASSERT_PARAM (iv);

   keyBlob = NULL;

   state = bson_malloc0 (sizeof (*state));
   BSON_ASSERT (state);

   state->key_handle = INVALID_HANDLE_VALUE;

   /* Initialize key storage buffer */
   state->key_object = bson_malloc0 (_aes256_key_blob_length);
   BSON_ASSERT (state->key_object);

   state->key_object_length = _aes256_key_blob_length;

   /* Allocate temporary buffer for key import */
   BSON_ASSERT (sizeof (BCRYPT_KEY_DATA_BLOB_HEADER) + key->len <= UINT32_MAX);
   keyBlobLength = sizeof (BCRYPT_KEY_DATA_BLOB_HEADER) + key->len;
   keyBlob = bson_malloc0 (keyBlobLength);
   BSON_ASSERT (keyBlob);


   blobHeader.dwMagic = BCRYPT_KEY_DATA_BLOB_MAGIC;
   blobHeader.dwVersion = BCRYPT_KEY_DATA_BLOB_VERSION1;
   blobHeader.cbKeyData = key->len;

   memcpy (keyBlob, &blobHeader, sizeof (BCRYPT_KEY_DATA_BLOB_HEADER));

   memcpy (keyBlob + sizeof (BCRYPT_KEY_DATA_BLOB_HEADER), key->data, key->len);

   nt_status = BCryptImportKey (_algo_aes256_cbc,
                                NULL,
                                BCRYPT_KEY_DATA_BLOB,
                                &(state->key_handle),
                                state->key_object,
                                state->key_object_length,
                                keyBlob,
                                keyBlobLength,
                                0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("Import Key Failed: 0x%x", (int) nt_status);
      goto fail;
   }

   bson_free (keyBlob);

   state->iv = bson_malloc0 (iv->len);
   BSON_ASSERT (state->iv);

   state->iv_len = iv->len;
   memcpy (state->iv, iv->data, iv->len);

   return state;
fail:
   _crypto_state_destroy (state);
   bson_free (keyBlob);

   return NULL;
}


static void
_crypto_state_destroy (cng_encrypt_state *state)
{
   if (state) {
      /* Free the key handle before the key_object that contains it */
      if (state->key_handle != INVALID_HANDLE_VALUE) {
         BCryptDestroyKey (state->key_handle);
      }

      bson_free (state->key_object);
      bson_free (state->iv);
      bson_free (state);
   }
}


bool
_native_crypto_aes_256_cbc_encrypt (aes_256_args_t args)
{
   BSON_ASSERT (args.in);
   BSON_ASSERT (args.out);

   bool ret = false;
   mongocrypt_status_t *status = args.status;
   cng_encrypt_state *state = _crypto_state_init (args.key, args.iv, status);

   BSON_ASSERT (state);

   NTSTATUS nt_status;

   nt_status = BCryptEncrypt (state->key_handle,
                              (PUCHAR) (args.in->data),
                              args.in->len,
                              NULL,
                              state->iv,
                              state->iv_len,
                              args.out->data,
                              args.out->len,
                              args.bytes_written,
                              0);

   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error initializing cipher: 0x%x", (int) nt_status);
      goto done;
   }

   ret = true;
done:
   _crypto_state_destroy (state);
   return ret;
}


bool
_native_crypto_aes_256_cbc_decrypt (aes_256_args_t args)
{
   BSON_ASSERT (args.in);
   BSON_ASSERT (args.out);

   bool ret = false;
   mongocrypt_status_t *status = args.status;
   cng_encrypt_state *state = _crypto_state_init (args.key, args.iv, status);

   BSON_ASSERT (state);

   NTSTATUS nt_status;

   nt_status = BCryptDecrypt (state->key_handle,
                              (PUCHAR) (args.in->data),
                              args.in->len,
                              NULL,
                              state->iv,
                              state->iv_len,
                              args.out->data,
                              args.out->len,
                              args.bytes_written,
                              0);


   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error initializing cipher: 0x%x", (int) nt_status);
      goto done;
   }

   ret = true;
done:
   _crypto_state_destroy (state);
   return ret;
}

/* _hmac_with_algorithm computes an HMAC of @in with the algorithm specified by
 * @hAlgorithm.
 * @key is the input key.
 * @out is the output. @out must be allocated by the caller with
 * the expected length @expect_out_len for the output.
 * Returns false and sets @status on error. @status is required. */
static bool
_hmac_with_algorithm (BCRYPT_ALG_HANDLE hAlgorithm,
                      const _mongocrypt_buffer_t *key,
                      const _mongocrypt_buffer_t *in,
                      _mongocrypt_buffer_t *out,
                      uint32_t expect_out_len,
                      mongocrypt_status_t *status)
{
   bool ret = false;
   BCRYPT_HASH_HANDLE hHash;
   NTSTATUS nt_status;

   BSON_ASSERT_PARAM (key);
   BSON_ASSERT_PARAM (in);
   BSON_ASSERT_PARAM (out);

   if (out->len != expect_out_len) {
      CLIENT_ERR ("out does not contain " PRIu32 " bytes", expect_out_len);
      return false;
   }

   nt_status = BCryptCreateHash (
      hAlgorithm, &hHash, NULL, 0, (PUCHAR) key->data, (ULONG) key->len, 0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error initializing hmac: 0x%x", (int) nt_status);
      /* Only call BCryptDestroyHash if BCryptCreateHash succeeded. */
      return false;
   }

   nt_status = BCryptHashData (hHash, (PUCHAR) in->data, (ULONG) in->len, 0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error hashing data: 0x%x", (int) nt_status);
      goto done;
   }

   nt_status = BCryptFinishHash (hHash, out->data, out->len, 0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error finishing hmac: 0x%x", (int) nt_status);
      goto done;
   }

   ret = true;
done:
   (void) BCryptDestroyHash (hHash);
   return ret;
}

bool
_native_crypto_hmac_sha_512 (const _mongocrypt_buffer_t *key,
                             const _mongocrypt_buffer_t *in,
                             _mongocrypt_buffer_t *out,
                             mongocrypt_status_t *status)
{
   return _hmac_with_algorithm (
      _algo_sha512_hmac, key, in, out, MONGOCRYPT_HMAC_SHA512_LEN, status);
}


bool
_native_crypto_random (_mongocrypt_buffer_t *out,
                       uint32_t count,
                       mongocrypt_status_t *status)
{
   BSON_ASSERT_PARAM (out);

   NTSTATUS nt_status = BCryptGenRandom (_random, out->data, count, 0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("BCryptGenRandom Failed: 0x%x", (int) nt_status);
      return false;
   }

   return true;
}

typedef struct {
   BCRYPT_KEY_HANDLE key_handle;

   unsigned char *input_block;
   uint32_t input_block_len;

   unsigned char *output_block;
   uint32_t output_block_len;
   uint32_t output_block_ptr;
} cng_ctr_encrypt_state;

static bool
_cng_ctr_crypto_generate (cng_ctr_encrypt_state *state,
                          mongocrypt_status_t *status)
{
   BSON_ASSERT (state);

   uint32_t bytesEncrypted = 0;
   NTSTATUS nt_status = BCryptEncrypt (state->key_handle,
                                       state->input_block,
                                       state->input_block_len,
                                       NULL,
                                       NULL,
                                       0,
                                       state->output_block,
                                       state->output_block_len,
                                       &bytesEncrypted,
                                       0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("error encrypting: 0x%x", (int) nt_status);
      return false;
   }
   BSON_ASSERT (bytesEncrypted);
   state->output_block_ptr = 0;
   return true;
}

static void
_cng_ctr_crypto_advance (cng_ctr_encrypt_state *state)
{
   BSON_ASSERT_PARAM (state);

   /* Assert rather than return false/NULL since this function's type is void */
   BSON_ASSERT (sizeof (BCRYPT_KEY_DATA_BLOB_HEADER) <=
                UINT32_MAX - state->input_block_len);

   uint32_t carry = 1;
   for (int i = (int) state->input_block_len - 1; i >= 0 && carry != 0; --i) {
      uint32_t bpp = (uint32_t) (state->input_block[i]) + carry;
      carry = bpp >> 8;
      state->input_block[i] = bpp & 0xFF;
   }
}

static bool
_cng_ctr_crypto_next (cng_ctr_encrypt_state *state,
                      mongocrypt_status_t *status,
                      unsigned char *mask)
{
   BSON_ASSERT_PARAM (state);
   BSON_ASSERT_PARAM (mask);

   if (state->output_block_ptr >= state->output_block_len) {
      _cng_ctr_crypto_advance (state);
      if (!_cng_ctr_crypto_generate (state, status)) {
         return false;
      };
   }

   *mask = state->output_block[state->output_block_ptr];
   ++state->output_block_ptr;

   return true;
}

static void
_cng_ctr_crypto_state_destroy (cng_ctr_encrypt_state *state);

static cng_ctr_encrypt_state *
_cng_ctr_crypto_state_init (const _mongocrypt_buffer_t *key,
                            const _mongocrypt_buffer_t *iv,
                            mongocrypt_status_t *status)
{
   cng_ctr_encrypt_state *state;
   uint32_t keyBlobLength;
   unsigned char *keyBlob;
   BCRYPT_KEY_DATA_BLOB_HEADER blobHeader;
   NTSTATUS nt_status;

   BSON_ASSERT_PARAM (key);
   BSON_ASSERT_PARAM (iv);

   keyBlob = NULL;

   state = bson_malloc0 (sizeof (*state));
   BSON_ASSERT (state);

   if (UINT32_MAX - key->len < sizeof (BCRYPT_KEY_DATA_BLOB_HEADER)) {
      CLIENT_ERR ("key is too long");
      goto fail;
   }

   state->key_handle = INVALID_HANDLE_VALUE;

   /* Initialize input storage buffer */
   state->input_block = bson_malloc0 (_aes256_block_length);
   BSON_ASSERT (state->input_block);
   state->input_block_len = _aes256_block_length;
   BSON_ASSERT (iv->len == _aes256_block_length);
   memcpy (state->input_block, iv->data, iv->len);

   /* Initialize output storage buffer */
   state->output_block = bson_malloc0 (_aes256_block_length);
   BSON_ASSERT (state->output_block);
   state->output_block_len = _aes256_block_length;
   state->output_block_ptr = 0;

   /* Allocate temporary buffer for key import */
   keyBlobLength = sizeof (BCRYPT_KEY_DATA_BLOB_HEADER) + key->len;
   keyBlob = bson_malloc0 (keyBlobLength);
   BSON_ASSERT (keyBlob);

   blobHeader.dwMagic = BCRYPT_KEY_DATA_BLOB_MAGIC;
   blobHeader.dwVersion = BCRYPT_KEY_DATA_BLOB_VERSION1;
   blobHeader.cbKeyData = key->len;

   memcpy (keyBlob, &blobHeader, sizeof (BCRYPT_KEY_DATA_BLOB_HEADER));

   memcpy (keyBlob + sizeof (BCRYPT_KEY_DATA_BLOB_HEADER), key->data, key->len);

   nt_status = BCryptImportKey (_algo_aes256_ecb,
                                NULL,
                                BCRYPT_KEY_DATA_BLOB,
                                &(state->key_handle),
                                NULL,
                                0,
                                keyBlob,
                                keyBlobLength,
                                0);
   if (nt_status != STATUS_SUCCESS) {
      CLIENT_ERR ("Import Key Failed: 0x%x", (int) nt_status);
      goto fail;
   }

   bson_free (keyBlob);

   if (!_cng_ctr_crypto_generate (state, status)) {
      goto fail;
   }

   return state;
fail:
   _cng_ctr_crypto_state_destroy (state);
   bson_free (keyBlob);

   return NULL;
}

static void
_cng_ctr_crypto_state_destroy (cng_ctr_encrypt_state *state)
{
   if (state) {
      if (state->key_handle != INVALID_HANDLE_VALUE) {
         BCryptDestroyKey (state->key_handle);
      }

      bson_free (state->input_block);
      bson_free (state->output_block);
      bson_free (state);
   }
}

bool
_native_crypto_aes_256_ctr_encrypt (aes_256_args_t args)
{
   bool ret = false;
   cng_ctr_encrypt_state *state = NULL;
   mongocrypt_status_t *status = args.status;

   BSON_ASSERT (args.in && args.in->data);
   BSON_ASSERT (args.out && args.out->data);
   if (args.out->len < args.in->len) {
      CLIENT_ERR ("Output buffer is too small");
      goto fail;
   }

   state = _cng_ctr_crypto_state_init (args.key, args.iv, status);
   if (!state) {
      goto fail;
   }

   for (uint32_t i = 0; i < args.in->len; ++i) {
      unsigned char mask;
      if (!_cng_ctr_crypto_next (state, status, &mask)) {
         goto fail;
      }
      args.out->data[i] = args.in->data[i] ^ mask;
   }

   if (args.bytes_written) {
      *args.bytes_written = args.in->len;
   }

   ret = true;

fail:
   _cng_ctr_crypto_state_destroy (state);
   return ret;
}

bool
_native_crypto_aes_256_ctr_decrypt (aes_256_args_t args)
{
   bool ret = false;
   cng_ctr_encrypt_state *state = NULL;
   mongocrypt_status_t *status = args.status;

   BSON_ASSERT (args.in && args.in->data);
   BSON_ASSERT (args.out && args.out->data);
   if (args.out->len < args.in->len) {
      CLIENT_ERR ("Output buffer is too small");
      goto fail;
   }

   state = _cng_ctr_crypto_state_init (args.key, args.iv, status);
   if (!state) {
      goto fail;
   }

   for (uint32_t i = 0; i < args.in->len; ++i) {
      unsigned char mask;
      if (!_cng_ctr_crypto_next (state, status, &mask)) {
         goto fail;
      }
      args.out->data[i] = args.in->data[i] ^ mask;
   }

   if (args.bytes_written) {
      *args.bytes_written = args.in->len;
   }

   ret = true;

fail:
   _cng_ctr_crypto_state_destroy (state);
   return ret;
}

bool
_native_crypto_hmac_sha_256 (const _mongocrypt_buffer_t *key,
                             const _mongocrypt_buffer_t *in,
                             _mongocrypt_buffer_t *out,
                             mongocrypt_status_t *status)
{
   return _hmac_with_algorithm (
      _algo_sha256_hmac, key, in, out, MONGOCRYPT_HMAC_SHA256_LEN, status);
}

#endif /* MONGOCRYPT_ENABLE_CRYPTO_CNG */