// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "crypto/encryptor.h"

#include <vector>

#include "crypto/symmetric_key.h"

namespace crypto {

namespace {

// On success, returns the block size (in bytes) for the algorithm that |key|
// is for. On failure, returns 0.
DWORD GetCipherBlockSize(HCRYPTKEY key) {
  DWORD block_size_in_bits = 0;
  DWORD param_size = sizeof(block_size_in_bits);
  BOOL ok = CryptGetKeyParam(key, KP_BLOCKLEN,
                             reinterpret_cast<BYTE*>(&block_size_in_bits),
                             &param_size, 0);
  if (!ok)
    return 0;

  return block_size_in_bits / 8;
}

}  // namespace

Encryptor::Encryptor()
    : key_(NULL),
      mode_(CBC),
      block_size_(0) {
}

Encryptor::~Encryptor() {
}

bool Encryptor::Init(SymmetricKey* key, Mode mode, const std::string& iv) {
  DCHECK(key);
  DCHECK_EQ(CBC, mode) << "Unsupported mode of operation";

  // In CryptoAPI, the IV, padding mode, and feedback register (for a chaining
  // mode) are properties of a key, so we have to create a copy of the key for
  // the Encryptor.  See the Remarks section of the CryptEncrypt MSDN page.
  BOOL ok = CryptDuplicateKey(key->key(), NULL, 0, capi_key_.receive());
  if (!ok)
    return false;

  // CRYPT_MODE_CBC is the default for Microsoft Base Cryptographic Provider,
  // but we set it anyway to be safe.
  DWORD cipher_mode = CRYPT_MODE_CBC;
  ok = CryptSetKeyParam(capi_key_.get(), KP_MODE,
                        reinterpret_cast<BYTE*>(&cipher_mode), 0);
  if (!ok)
    return false;

  block_size_ = GetCipherBlockSize(capi_key_.get());
  if (block_size_ == 0)
    return false;

  if (iv.size() != block_size_)
    return false;

  ok = CryptSetKeyParam(capi_key_.get(), KP_IV,
                        reinterpret_cast<const BYTE*>(iv.data()), 0);
  if (!ok)
    return false;

  DWORD padding_method = PKCS5_PADDING;
  ok = CryptSetKeyParam(capi_key_.get(), KP_PADDING,
                        reinterpret_cast<BYTE*>(&padding_method), 0);
  if (!ok)
    return false;

  return true;
}

bool Encryptor::Encrypt(const std::string& plaintext, std::string* ciphertext) {
  DWORD data_len = plaintext.size();
  DWORD total_len = data_len + block_size_;

  // CryptoAPI encrypts/decrypts in place.
  std::vector<BYTE> tmp(total_len);
  memcpy(&tmp[0], plaintext.data(), data_len);

  BOOL ok = CryptEncrypt(capi_key_.get(), NULL, TRUE, 0, &tmp[0],
                         &data_len, total_len);
  if (!ok)
    return false;

  ciphertext->assign(reinterpret_cast<char*>(&tmp[0]), data_len);
  return true;
}

bool Encryptor::Decrypt(const std::string& ciphertext, std::string* plaintext) {
  DWORD data_len = ciphertext.size();
  if (data_len == 0)
    return false;

  std::vector<BYTE> tmp(data_len);
  memcpy(&tmp[0], ciphertext.data(), data_len);

  BOOL ok = CryptDecrypt(capi_key_.get(), NULL, TRUE, 0, &tmp[0], &data_len);
  if (!ok)
    return false;

  DCHECK_GT(tmp.size(), data_len);

  plaintext->assign(reinterpret_cast<char*>(&tmp[0]), data_len);
  return true;
}

}  // namespace crypto