////////////////////////////////////////////////////////////////////////////////
// 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 "glog/logging.h"
#include "absl/strings/string_view.h"
#include "openssl/aes.h"
#include "openssl/cmac.h"
#include "openssl/evp.h"
#include "openssl/hmac.h"
#include "openssl/sha.h"
#include "util/endian/endian.h"

namespace widevine {
namespace crypto_util {

const char kEncryptionKeyLabel[] = "ENCRYPTION";
const int kEncryptionKeySizeBits = 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.
const char kPhonyGroupMasterKey[] = "fedcba9876543210";
const int kAes128KeySizeBits = 128;
const int kAes128KeySizeBytes = 16;

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) {
  HMAC_CTX ctx;
  HMAC_CTX_init(&ctx);
  HMAC_Init(&ctx, key.data(), key.size(), EVP_sha256());
  HMAC_Update(&ctx, reinterpret_cast<const unsigned char*>(message.data()),
              message.size());
  unsigned char digest[SHA256_DIGEST_LENGTH];
  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) {
  return CreateSignatureHmacSha256(key, message) == signature;
}

// Creates a SHA-1 HMAC signature for the given message.
std::string CreateSignatureHmacSha1(absl::string_view key,
                               absl::string_view message) {
  HMAC_CTX ctx;
  HMAC_CTX_init(&ctx);
  HMAC_Init(&ctx, key.data(), key.size(), EVP_sha1());
  HMAC_Update(&ctx, reinterpret_cast<const unsigned char*>(message.data()),
              message.size());
  unsigned char digest[SHA_DIGEST_LENGTH];
  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-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 an AES 128 key from the provided key and additional info.
std::string DeriveKey(absl::string_view key, absl::string_view label,
                 absl::string_view context, const uint32_t size_bits) {
  if (key.size() != kAes128KeySizeBytes) return "";

  // We only handle even multiples of 16 bytes (128 bits) right now.
  if ((size_bits % 128) || (size_bits > (128 * 255))) {
    return "";
  }

  std::string result;

  const EVP_CIPHER* cipher = EVP_aes_128_cbc();
  CMAC_CTX* cmac_ctx = CMAC_CTX_new();

  for (unsigned char counter = 1; counter <= (size_bits / 128); counter++) {
    if (CMAC_Init(cmac_ctx, key.data(), key.size(), cipher, 0)) {
      std::string message;
      message.append(1, counter);
      message.append(label.data(), label.size());
      message.append(1, '\0');
      message.append(context.data(), context.size());
      char size_string[4];
      BigEndian::Store32(&size_string, size_bits);
      message.append(&size_string[0], &size_string[0] + 4);
      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