// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#include <stdint.h>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "base/logging.h"
#include "base/strings/string_number_conversions.h"
#include "cdm/protos/license_protocol.pb.h"
#include "crypto_utils/aes_cbc_decryptor.h"
#include "crypto_utils/aes_ctr_encryptor.h"
#include "crypto_utils/crypto_util.h"
#include "crypto_utils/rsa_key.h"
#include "impl/reference/memory_util.h"
#include "impl/reference/string_view_util.h"
#include "third_party/boringssl/src/include/openssl/aes.h"
#include "third_party/boringssl/src/include/openssl/cmac.h"
#include "third_party/boringssl/src/include/openssl/err.h"
#include "third_party/boringssl/src/include/openssl/evp.h"
#include "third_party/boringssl/src/include/openssl/hmac.h"
#include "third_party/boringssl/src/include/openssl/rsa.h"
#include "third_party/boringssl/src/include/openssl/sha.h"

namespace {
using KeyContainer = video_widevine::License_KeyContainer;
using RsaPrivateKey = widevine::RsaPrivateKey;
using AesCbcDecryptor = widevine::AesCbcDecryptor;
using AesCtrDecryptor = widevine::AesCtrEncryptor;
using SecurityLevel = video_widevine::License_KeyContainer_SecurityLevel;

struct ContentKey {
  // Minimum security level that this key requires in order to be used.
  SecurityLevel level;

  // Key used to decrypt content.
  std::vector<uint8_t> key;
};
}  // namespace

// The white-box type can't be in the namespace as it is defined in the header.
struct WB_License_Whitebox {
  // CDM key, used for license requests.
  std::unique_ptr<RsaPrivateKey> key;

  // Keys used for license renewal.
  std::string server_signing_key;
  std::string client_signing_key;

  // We use two data structures to track keys. We track all the key ids seen
  // when loading the license and the content keys. This allows us to tell the
  // difference between a key missing and a key being misused.
  std::set<std::string> all_key_ids;
  std::map<std::string, ContentKey> content_keys;
};

namespace {
// Secret string value. For simplicity the pattern will just be 0xAAAAAAAA which
// is xor'd to the buffer, and removed by xor'ing it a second time.
const uint8_t kSecretStringPattern[] = {
    0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA,
    0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA,
};
const size_t kSecretStringPatternSize = sizeof(kSecretStringPattern);
static_assert(kSecretStringPatternSize == AES_BLOCK_SIZE,
              "Secret string must be AES_BLOCK_SIZE.");

std::string AsString(const uint8_t* buffer, size_t buffer_size) {
  return std::string(reinterpret_cast<const char*>(buffer), buffer_size);
}

void ApplyPattern(const uint8_t* input_buffer,
                  size_t input_buffer_size,
                  const uint8_t* pattern,
                  size_t pattern_size,
                  uint8_t* output_buffer) {
  DCHECK(input_buffer);
  DCHECK(pattern);
  DCHECK(output_buffer);
  DCHECK_GT(pattern_size, 0u);
  DCHECK_EQ(input_buffer_size % pattern_size, 0u);

  for (size_t i = 0; i < input_buffer_size; ++i) {
    output_buffer[i] = input_buffer[i] ^ pattern[i % pattern_size];
  }
}

WB_Result FindKey(const WB_License_Whitebox* whitebox,
                  const uint8_t* id,
                  size_t id_size,
                  ContentKey* key) {
  DCHECK(whitebox);
  DCHECK(id);
  DCHECK_GT(id_size, 0u);
  DCHECK(key);

  const std::string key_id = AsString(id, id_size);

  if (whitebox->all_key_ids.empty()) {
    DVLOG(1) << "Invalid state: no keys have been loaded.";
    return WB_RESULT_INVALID_STATE;
  }

  if (whitebox->all_key_ids.find(key_id) == whitebox->all_key_ids.end()) {
    DVLOG(1) << "No such key: " << base::HexEncode(id, id_size) << ".";
    return WB_RESULT_NO_SUCH_KEY;
  }

  const auto found = whitebox->content_keys.find(key_id);

  if (found == whitebox->content_keys.end()) {
    DVLOG(1) << "Wrong key type: " << base::HexEncode(id, id_size) << ".";
    return WB_RESULT_WRONG_KEY_TYPE;
  }

  *key = found->second;
  return WB_RESULT_OK;
}

WB_Result DecryptBuffer(WB_CipherMode mode,
                        const std::vector<uint8_t>& key,
                        const uint8_t* input_data,
                        size_t input_data_size,
                        const uint8_t* iv,
                        size_t iv_size,
                        uint8_t* output_data,
                        size_t* output_data_size) {
  if (!input_data || !iv || !output_data || !output_data_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (input_data_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // Data must be "block aligned" for CBC. CTR does not have this requirement.
  if (mode == WB_CIPHER_MODE_CBC && input_data_size % AES_BLOCK_SIZE != 0) {
    DVLOG(1) << "Invalid parameter: bad block size.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // We only support 16 byte IVs. CENC allows 8 byte, but it is expected that
  // the IV will be padded before entering the whitebox.
  if (iv_size != 16) {
    DVLOG(1) << "Invalid parameter: invalid iv size.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // There is no padding, so the output will be the same length as the input.
  if (*output_data_size < input_data_size) {
    DVLOG(1) << "Buffer too small: needs " << input_data_size << " but got "
             << *output_data_size << ".";
    *output_data_size = input_data_size;
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  // By this point, we have verified everything that need to be verified.
  // Decryption should just work.
  if (mode == WB_CIPHER_MODE_CBC) {
    AesCbcDecryptor decryptor;
    CHECK(decryptor.SetKey(key.data(), key.size()));

    *output_data_size = input_data_size;
    CHECK(decryptor.Decrypt(iv, iv_size, input_data, input_data_size,
                            output_data));
  } else if (mode == WB_CIPHER_MODE_CTR) {
    AesCtrDecryptor decryptor;
    CHECK(decryptor.SetKey(key.data(), key.size()));

    // Encrypt and Decrypt for CBC use the same interface.
    *output_data_size = input_data_size;
    CHECK(decryptor.Encrypt(iv, iv_size, input_data, input_data_size,
                            output_data));
  } else {
    DVLOG(1) << "Invalid parameter: invalid cipher mode.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  return WB_RESULT_OK;
}
}  // namespace

WB_Result WB_License_Create(const uint8_t* whitebox_init_data,
                            size_t whitebox_init_data_size,
                            WB_License_Whitebox** whitebox) {
  if (!whitebox_init_data || !whitebox) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (whitebox_init_data_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // |whitebox_init_data| is simply the bytes of a PKCS #8 PrivateKeyInfo block
  // of a RSA 2048 bit key.
  std::unique_ptr<RsaPrivateKey> key(RsaPrivateKey::Create(
      AsString(whitebox_init_data, whitebox_init_data_size)));
  if (!key) {
    DVLOG(1) << "Invalid parameter: invalid init data.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // Should always be non-null on modern compilers
  // (https://isocpp.org/wiki/faq/freestore-mgmt).
  *whitebox = new WB_License_Whitebox();
  (*whitebox)->key.swap(key);

  return WB_RESULT_OK;
}

void WB_License_Delete(WB_License_Whitebox* whitebox) {
  // Safe to delete nullptr (https://isocpp.org/wiki/faq/freestore-mgmt).
  delete whitebox;
}

WB_Result WB_License_SignLicenseRequest(const WB_License_Whitebox* whitebox,
                                        const uint8_t* license_request,
                                        size_t license_request_size,
                                        uint8_t* signature,
                                        size_t* signature_size) {
  if (!whitebox || !license_request || !signature || !signature_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (license_request_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  std::string result;
  DCHECK(whitebox->key->GenerateSignature(
      AsString(license_request, license_request_size), &result));

  if (!widevine::MemCopy(result.data(), result.size(), signature,
                         *signature_size)) {
    DVLOG(1) << "Buffer too small: signature needs " << result.size() << ".";
    *signature_size = result.size();
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  return WB_RESULT_OK;
}

WB_Result WB_License_ProcessLicenseResponse(WB_License_Whitebox* whitebox,
                                            const uint8_t* message,
                                            size_t message_size,
                                            const uint8_t* signature,
                                            size_t signature_size,
                                            const uint8_t* session_key,
                                            size_t session_key_size,
                                            const uint8_t* license_request,
                                            size_t license_request_size) {
  if (!whitebox || !message || !signature || !session_key || !license_request) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (message_size == 0 || license_request_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  std::string decrypted_session_key;
  if (!whitebox->key->Decrypt(AsString(session_key, session_key_size),
                              &decrypted_session_key)) {
    DVLOG(1) << "Invalid parameter: invalid session key.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  std::string signing_key_material = widevine::crypto_util::DeriveKey(
      decrypted_session_key, widevine::crypto_util::kSigningKeyLabel,
      widevine::AsStringView(license_request, license_request_size),
      widevine::crypto_util::kSigningKeySizeBits * 2);

  if (signing_key_material.size() <
      widevine::crypto_util::kSigningKeySizeBytes * 2) {
    DVLOG(1) << "Invalid parameter: invalid session key size.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (!widevine::crypto_util::VerifySignatureHmacSha256(
          signing_key_material,
          widevine::AsStringView(signature, signature_size),
          widevine::AsStringView(message, message_size))) {
    DVLOG(1) << "Failed to verify signed message.";
    return WB_RESULT_INVALID_SIGNATURE;
  }

  std::string decryption_key = widevine::crypto_util::DeriveKey(
      decrypted_session_key, widevine::crypto_util::kWrappingKeyLabel,
      widevine::AsStringView(license_request, license_request_size),
      widevine::crypto_util::kWrappingKeySizeBits);
  if (decryption_key.empty()) {
    DVLOG(1) << "Failed to decrypt decryption key";
    return WB_RESULT_INVALID_SIGNATURE;
  }

  AesCbcDecryptor decryptor;
  CHECK(
      decryptor.SetKey(reinterpret_cast<const uint8_t*>(decryption_key.data()),
                       decryption_key.size()));

  video_widevine::License license;
  if (!license.ParseFromArray(message, message_size)) {
    DVLOG(1) << "Invalid parameter: Invalid license.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  CHECK(!license.key().empty());
  for (const auto& key : license.key()) {
    const std::string wrapped_key = key.key();
    std::vector<uint8_t> unwrapped_key(wrapped_key.size());

    if (!decryptor.Decrypt(reinterpret_cast<const uint8_t*>(key.iv().data()),
                           key.iv().size(),
                           reinterpret_cast<const uint8_t*>(wrapped_key.data()),
                           wrapped_key.size(), unwrapped_key.data())) {
      // The input has to be a specific length, so if it is not, it means that
      // something is wrong with the license.
      DVLOG(1) << "Invalid parameter: Invalid license.";
      return WB_RESULT_INVALID_PARAMETER;
    }

    whitebox->all_key_ids.insert(key.id());

    switch (key.type()) {
      case KeyContainer::SIGNING: {
        const size_t kSigningKeySizeBytes =
            widevine::crypto_util::kSigningKeySizeBytes;
        if (unwrapped_key.size() < kSigningKeySizeBytes * 2) {
          DVLOG(1) << "Invalid parameter: Invalid signing key.";
          return WB_RESULT_INVALID_PARAMETER;
        }

        const std::string signing_key =
            AsString(unwrapped_key.data(), unwrapped_key.size());
        whitebox->server_signing_key =
            signing_key.substr(0, kSigningKeySizeBytes);
        whitebox->client_signing_key =
            signing_key.substr(kSigningKeySizeBytes, kSigningKeySizeBytes);
        break;
      }

      case KeyContainer::CONTENT: {
        constexpr size_t kContentKeySizeBytes = 16;
        if (unwrapped_key.size() < kContentKeySizeBytes) {
          DVLOG(1) << "Invalid parameter: Invalid content key.";
          return WB_RESULT_INVALID_PARAMETER;
        }

        unwrapped_key.resize(kContentKeySizeBytes);
        whitebox->content_keys[key.id()] = {key.level(),
                                            std::move(unwrapped_key)};
        break;
      }

      default:
        // We purposefully ignore the other cases. In the case that any other
        // key appeared in the license, the key id would have been tracked in
        // |all_key_ids|.
        break;
    }
  }

  return WB_RESULT_OK;
}

WB_Result WB_License_SignRenewalRequest(const WB_License_Whitebox* whitebox,
                                        const uint8_t* message,
                                        size_t message_size,
                                        uint8_t* signature,
                                        size_t* signature_size) {
  if (!whitebox || !message || !signature || !signature_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (message_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // License request must have been processed at least once before calling this
  // function.
  if (whitebox->client_signing_key.empty()) {
    DVLOG(1) << __func__ << " called before license received.";
    return WB_RESULT_INVALID_STATE;
  }

  const std::string computed_signature =
      widevine::crypto_util::CreateSignatureHmacSha256(
          whitebox->client_signing_key,
          widevine::AsStringView(message, message_size));

  if (!widevine::MemCopy(computed_signature.data(), computed_signature.size(),
                         signature, *signature_size)) {
    DVLOG(1) << "Buffer too small: signature needs "
             << computed_signature.size() << ".";
    *signature_size = computed_signature.size();
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  return WB_RESULT_OK;
}

WB_Result WB_License_VerifyRenewalResponse(const WB_License_Whitebox* whitebox,
                                           const uint8_t* message,
                                           size_t message_size,
                                           const uint8_t* signature,
                                           size_t signature_size) {
  if (!whitebox || !message || !signature) {
    DVLOG(1) << "Invalid parameter: null pointer";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (message_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // License request must have been processed at least once before calling this
  // function.
  if (whitebox->server_signing_key.empty()) {
    DVLOG(1) << __func__ << " called before license received.";
    return WB_RESULT_INVALID_STATE;
  }

  const std::string computed_signature =
      widevine::crypto_util::CreateSignatureHmacSha256(
          whitebox->server_signing_key,
          widevine::AsStringView(message, message_size));

  if (signature_size < computed_signature.size()) {
    DVLOG(1) << "Invalid parameters: invalid signature "
                "size.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (computed_signature != widevine::AsStringView(signature, signature_size)) {
    DVLOG(1) << "Data verification error: signatures do not match.";
    return WB_RESULT_DATA_VERIFICATION_ERROR;
  }

  return WB_RESULT_OK;
}

WB_Result WB_License_GetSecretString(const WB_License_Whitebox* whitebox,
                                     WB_CipherMode mode,
                                     const uint8_t* key_id,
                                     size_t key_id_size,
                                     uint8_t* secret_string,
                                     size_t* secret_string_size) {
  if (!whitebox || !key_id || !secret_string || !secret_string_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (mode != WB_CIPHER_MODE_CTR && mode != WB_CIPHER_MODE_CBC) {
    DVLOG(1) << "Invalid parameter: invalid cipher mode.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (key_id_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // The secret string can differ between keys, so we need to make sure that
  // the key id is actually a content key.
  ContentKey content_key;
  const WB_Result key_lookup_result =
      FindKey(whitebox, key_id, key_id_size, &content_key);
  if (key_lookup_result != WB_RESULT_OK) {
    return key_lookup_result;
  }

  if (!widevine::MemCopy(kSecretStringPattern, kSecretStringPatternSize,
                         secret_string, *secret_string_size)) {
    DVLOG(1) << "Buffer too small: needs " << kSecretStringPatternSize << ".";
    *secret_string_size = kSecretStringPatternSize;
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  return WB_RESULT_OK;
}

WB_Result WB_License_Decrypt(const WB_License_Whitebox* whitebox,
                             WB_CipherMode mode,
                             const uint8_t* key_id,
                             size_t key_id_size,
                             const uint8_t* input_data,
                             size_t input_data_size,
                             const uint8_t* iv,
                             size_t iv_size,
                             uint8_t* output_data,
                             size_t* output_data_size) {
  if (!whitebox || !key_id) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (key_id_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  ContentKey content_key;
  const WB_Result key_lookup_result =
      FindKey(whitebox, key_id, key_id_size, &content_key);
  if (key_lookup_result != WB_RESULT_OK) {
    return key_lookup_result;
  }

  // Check if decryption is allowed by the policy.
  switch (content_key.level) {
    case video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_CRYPTO:
      break;
    case video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_DECODE:
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_CRYPTO:
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_DECODE:
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_ALL:
      return WB_RESULT_INSUFFICIENT_SECURITY_LEVEL;
  }

  // DecryptBuffer() will validate the remaining decryption parameters.
  return DecryptBuffer(mode, content_key.key, input_data, input_data_size, iv,
                       iv_size, output_data, output_data_size);
}

WB_Result WB_License_MaskedDecrypt(const WB_License_Whitebox* whitebox,
                                   WB_CipherMode mode,
                                   const uint8_t* key_id,
                                   size_t key_id_size,
                                   const uint8_t* input_data,
                                   size_t input_data_size,
                                   const uint8_t* iv,
                                   size_t iv_size,
                                   uint8_t* masked_output_data,
                                   size_t* masked_output_data_size) {
  if (!whitebox || !key_id) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (key_id_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  ContentKey content_key;
  const WB_Result key_lookup_result =
      FindKey(whitebox, key_id, key_id_size, &content_key);
  if (key_lookup_result != WB_RESULT_OK) {
    return key_lookup_result;
  }

  // Check if decryption is allowed by the policy.
  switch (content_key.level) {
    case video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_CRYPTO:
    case video_widevine::License_KeyContainer_SecurityLevel_SW_SECURE_DECODE:
      break;
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_CRYPTO:
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_DECODE:
    case video_widevine::License_KeyContainer_SecurityLevel_HW_SECURE_ALL:
      return WB_RESULT_INSUFFICIENT_SECURITY_LEVEL;
  }

  // DecryptBuffer() will validate all the parameters, so just make sure it is
  // safe to resize this and let the normal validation process handle anything
  // wrong with the output size.
  std::vector<uint8_t> output;
  if (masked_output_data_size) {
    output.resize(*masked_output_data_size);
  }

  // DecryptBuffer() will validate the remaining decryption parameters and set
  // |masked_output_data_size|.
  const WB_Result result =
      DecryptBuffer(mode, content_key.key, input_data, input_data_size, iv,
                    iv_size, output.data(), masked_output_data_size);

  if (result != WB_RESULT_OK) {
    return result;
  }

  // Trivial implementation that simply takes the decrypted output and XORs it
  // with a fixed pattern. |output|'s size is based on |masked_output_data| so
  // we shouldn't need to worry about overflow.
  ApplyPattern(output.data(), output.size(), kSecretStringPattern,
               kSecretStringPatternSize, masked_output_data);

  return WB_RESULT_OK;
}

void WB_License_Unmask(const uint8_t* secret_string,
                       size_t secret_string_size,
                       uint8_t* buffer,
                       size_t buffer_size) {
  // No return code, so only check if parameters are valid.
  DCHECK(secret_string);
  DCHECK(secret_string_size);
  DCHECK(buffer);
  DCHECK(buffer_size);

  ApplyPattern(buffer, buffer_size, secret_string, secret_string_size, buffer);
}