////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google LLC.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////

// Implementation of Common crypto utilities used by Widevine services.

#include "common/crypto_util.h"

#include <cstdint>

#include "glog/logging.h"
#include "absl/strings/escaping.h"
#include "absl/strings/string_view.h"
#include "openssl/aes.h"
#include "openssl/cmac.h"
#include "openssl/digest.h"
#include "openssl/evp.h"
#include "openssl/hmac.h"
#include "openssl/sha.h"

namespace widevine {
namespace crypto_util {

const char kWrappingKeyLabel[] = "ENCRYPTION";
const int kWrappingKeySizeBits = 128;
const char kSigningKeyLabel[] = "AUTHENTICATION";
const int kSigningKeySizeBits = 256;
const size_t kSigningKeySizeBytes = 32;
const char kIvMasterKey[] = "1234567890123456";
const char kIvLabel[] = "IV_ENCRYPTION";
const int kIvSizeBits = 128;
const char kKeyIdMasterKey[] = "0123456789abcdef";
const char kKeyIdLabel[] = "KEY_ID_ENCRYPTION";
const int kKeyIdSizeBits = 128;
const char kGroupKeyLabel[] = "GROUP_ENCRYPTION";
// TODO(user): This is a temporary key for development. Replace this with
// a real group master key in keystore.
// TODO(user): figure out why VerifySignatureHmacSha256 can not crypto_mcmcpy
// like VerifySignatureHmacSha1.
// TODO(user): Revert logging signature in VerifySignatureHmacSha256.
// function.
const char kPhonyGroupMasterKey[] = "fedcba9876543210";
const int kAes128KeySizeBits = 128;
const int kAes128KeySizeBytes = 16;
const int kAes256KeySizeBytes = 32;
const char kKeyboxV3Label[] = "Keyboxv3";

const int kAesBlockSizeBits = AES_BLOCK_SIZE * 8;
const int kAesMaxDerivedBlocks = 255;

const uint32_t kCENCSchemeID = 0x63656E63;  // 'cenc' (AES-CTR): 0x63656E63
const uint32_t kCBC1SchemeID = 0x63626331;  // 'cbc1' (AES-CBC): 0x63626331
const uint32_t kCENSSchemeID =
    0x63656E73;  // 'cens' (AES-CTR subsample): 0x63656E73
const uint32_t kCBCSSchemeID =
    0x63626373;  // 'cbcs' (AES-CBC subsample): 0x63626373

// Creates a SHA-256 HMAC signature for the given message.
std::string CreateSignatureHmacSha256(absl::string_view key,
                                      absl::string_view message) {
  unsigned char digest[SHA256_DIGEST_LENGTH];
  return CreateSignatureHmac(EVP_sha256(), digest, key, message);
}

// Creates a SHA-384 HMAC signature for the given message.
std::string CreateSignatureHmacSha384(absl::string_view key,
                                      absl::string_view message) {
  unsigned char digest[SHA384_DIGEST_LENGTH];
  return CreateSignatureHmac(EVP_sha384(), digest, key, message);
}

std::string CreateSignatureHmac(const EVP_MD* hash_algorithm,
                                unsigned char* digest, absl::string_view key,
                                absl::string_view message) {
  if (hash_algorithm == nullptr || digest == nullptr) {
    return "";
  }
  HMAC_CTX ctx;
  HMAC_CTX_init(&ctx);
  HMAC_Init(&ctx, key.data(), key.size(), hash_algorithm);
  HMAC_Update(&ctx, reinterpret_cast<const unsigned char*>(message.data()),
              message.size());
  unsigned int digest_len;
  HMAC_Final(&ctx, digest, &digest_len);
  HMAC_CTX_cleanup(&ctx);
  std::string s(reinterpret_cast<char*>(digest), digest_len);
  return s;
}

// Compares the SHA-256 HMAC against the provided signature.
bool VerifySignatureHmacSha256(absl::string_view key,
                               absl::string_view signature,
                               absl::string_view message) {
  std::string actual_signature = CreateSignatureHmacSha256(key, message);
  bool result = actual_signature == signature;
  if (!result) {
    VLOG(1) << "Invalid signature. actual signature: " << actual_signature;
  }
  return result;
}

// Compares the SHA-384 HMAC against the provided signature.
bool VerifySignatureHmacSha384(absl::string_view key,
                               absl::string_view signature,
                               absl::string_view message) {
  return CreateSignatureHmacSha384(key, message) == signature;
}

// Creates a SHA-1 HMAC signature for the given message.
std::string CreateSignatureHmacSha1(absl::string_view key,
                                    absl::string_view message) {
  unsigned char digest[SHA_DIGEST_LENGTH];
  return CreateSignatureHmac(EVP_sha1(), digest, key, message);
}

// Compares the SHA-1 HMAC against the provided signature.
bool VerifySignatureHmacSha1(absl::string_view key, absl::string_view signature,
                             absl::string_view message) {
  return CreateSignatureHmacSha1(key, message) == signature;
}

// Derives a key from the provided AES 128 or 256 key and additional info.
std::string DeriveKey(absl::string_view key, absl::string_view label,
                      absl::string_view context, const uint32_t size_bits) {
  // We only handle multiples of AES blocks (16 bytes) with a maximum of 255
  // blocks.
  const uint32_t output_block_count = size_bits / kAesBlockSizeBits;
  if (size_bits % kAesBlockSizeBits ||
      output_block_count > kAesMaxDerivedBlocks) {
    return "";
  }

  const EVP_CIPHER* cipher = nullptr;
  const size_t key_size_bytes = key.size();

  switch (key_size_bytes) {
    case kAes128KeySizeBytes:
      cipher = EVP_aes_128_cbc();
      break;
    case kAes256KeySizeBytes:
      cipher = EVP_aes_256_cbc();
      break;
    default:
      return "";
  }

  std::string result;
  CMAC_CTX* cmac_ctx = CMAC_CTX_new();

  for (unsigned char counter = 0; counter < output_block_count; counter++) {
    if (CMAC_Init(cmac_ctx, key.data(), key_size_bytes, cipher, nullptr)) {
      std::string message;
      message.append(1, counter + 1);
      message.append(label.data(), label.size());
      message.append(1, '\0');
      message.append(context.data(), context.size());
      message.append(1, (size_bits >> 24) & 0xFF);
      message.append(1, (size_bits >> 16) & 0xFF);
      message.append(1, (size_bits >> 8) & 0xFF);
      message.append(1, size_bits & 0xFF);
      if (CMAC_Update(cmac_ctx,
                      reinterpret_cast<const uint8_t*>(message.data()),
                      message.size())) {
        size_t reslen;
        unsigned char res[AES_BLOCK_SIZE];
        if (CMAC_Final(cmac_ctx, res, &reslen)) {
          result.append(reinterpret_cast<const char*>(res), reslen);
        }
        DCHECK(reslen == AES_BLOCK_SIZE);
      }
    }
  }

  CMAC_CTX_free(cmac_ctx);
  return result;
}

// Derives an IV from the provided info.
std::string DeriveIv(absl::string_view context) {
  return DeriveKey(kIvMasterKey, kIvLabel, context, kIvSizeBits);
}

// Derives a key ID from the provided info.
std::string DeriveKeyId(absl::string_view context) {
  return DeriveKey(kKeyIdMasterKey, kKeyIdLabel, context, kKeyIdSizeBits);
}

std::string DeriveGroupSessionKey(absl::string_view context,
                                  const uint32_t size_bits) {
  return DeriveKey(kPhonyGroupMasterKey, kGroupKeyLabel, context, size_bits);
}

std::string DeriveSigningKey(absl::string_view key, absl::string_view context,
                             const uint32_t size_bits) {
  return DeriveKey(key, kSigningKeyLabel, context, size_bits);
}

bool FourCCEncryptionSchemeIDFromString(const std::string& requested,
                                        uint32_t* four_cc_code) {
  if (requested.size() != 4 || four_cc_code == nullptr) return false;

  uint32_t result = 0;
  for (auto i = 0; i < 4; ++i) {
    result <<= 8;
    result |= requested[i];
  }

  switch (result) {
    case kCENCSchemeID:
    case kCBC1SchemeID:
    case kCENSSchemeID:
    case kCBCSSchemeID:
      *four_cc_code = result;
      return true;
    default:
      return false;
  }
}

}  // namespace crypto_util
}  // namespace widevine