// Copyright 2019 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
#include <assert.h>
#include <string.h>
#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <sstream>
#include <vector>

#include <openssl/cmac.h>
#include <openssl/err.h>
#include <openssl/evp.h>

// clang-format off

// This is a test keybox.  It will not be accepted by production systems.
static std::vector<uint8_t> test_keybox = {
  // sample keybox used for test vectors
  // deviceID = WidevineTestOnlyKeybox000
  0x57, 0x69, 0x64, 0x65, 0x76, 0x69, 0x6e, 0x65,
  0x54, 0x65, 0x73, 0x74, 0x4f, 0x6e, 0x6c, 0x79,
  0x4b, 0x65, 0x79, 0x62, 0x6f, 0x78, 0x30, 0x30,
  0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,

  // key
  0xe4, 0xff, 0x57, 0x4c, 0x32, 0x2e, 0xf5, 0x34,
  0x26, 0x21, 0x2c, 0xb3, 0xed, 0x37, 0xf3, 0x5e,

  // data (system ID 7912 = 1EE8).
  0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x1e, 0xe8,
  0xca, 0x1e, 0x71, 0x7c, 0xfb, 0xe8, 0xa3, 0x94,
  0x52, 0x0a, 0x6b, 0x71, 0x37, 0xd2, 0x69, 0xfa,
  0x5a, 0xc6, 0xb5, 0x4c, 0x6b, 0x46, 0x63, 0x9b,
  0xbe, 0x80, 0x3d, 0xbb, 0x4f, 0xf7, 0x4c, 0x5f,
  0x6f, 0x55, 0x0e, 0x3d, 0x3d, 0x9a, 0xcf, 0x81,
  0x12, 0x5d, 0x52, 0xe0, 0x47, 0x8c, 0xda, 0x0b,
  0xf4, 0x31, 0x41, 0x13, 0xd0, 0xd5, 0x2d, 0xa0,
  0x5b, 0x20, 0x9a, 0xed, 0x51, 0x5d, 0x13, 0xd6,

  // magic
  0x6b, 0x62, 0x6f, 0x78,

  // Crc
  0x39, 0xf2, 0x94, 0xa7,
};

const size_t DEVICE_KEY_OFFSET = 0x20;
const size_t CA_TOKEN_OFFSET = 0x30;
const size_t KEYBOX_VERSION_OFFSET = CA_TOKEN_OFFSET;
const size_t SYSTEM_ID_OFFSET = CA_TOKEN_OFFSET + sizeof(uint32_t);
const size_t CA_TOKEN_SIZE = 0x48;
const size_t CA_TOKEN_END = CA_TOKEN_OFFSET + CA_TOKEN_SIZE;

// sample renewal key for testing
static const std::vector<uint8_t> renewal_key = {
  0xfa, 0xfd, 0xc1, 0x1b, 0x55, 0x6f, 0xac, 0xd3,
  0x14, 0x62, 0x50, 0xa0, 0xf7, 0xa5, 0x1a, 0x0e
};

static const std::string label = "Keyboxv3";
static const std::vector<uint8_t> new_system_id = {0x00, 0x00, 0x23, 0x45};
static const std::vector<uint8_t> new_keybox_version = {0x00, 0x00, 0x00, 0x03};

// clang-format on

static std::string GetSSLError() {
  char error_buffer[128];
  ERR_error_string_n(ERR_get_error(), error_buffer, sizeof(error_buffer));
  return error_buffer;
}

// Derive the new device key from renewal key and context using
// NIST 800-108 key derivation with 128-bit AES-128-CMAC as the pseudorandom
// function in counter mode:
//
// New Device Key := PRF(renewal_key, 1 || label || 0x00 || context || L)
// PRF := AES-128-CMAC
// label := “Keyboxv3”
// L=(unsigned long)0x80: 0x00|0x00|0x00|0x80
static bool DeriveKey(const std::vector<uint8_t>& key,
                      const std::vector<uint8_t>& context,
                      std::vector<uint8_t>* out) {
  if (key.empty() || context.empty() || out == nullptr) {
    std::cerr << "DeriveKey(): Invalid inputs" << std::endl;
    return false;
  }

  std::vector<uint8_t> cmac_input;
  cmac_input.push_back(1);
  std::copy(label.begin(), label.end(), std::back_inserter(cmac_input));
  cmac_input.push_back(0x00);
  std::copy(context.begin(), context.end(), std::back_inserter(cmac_input));
  cmac_input.push_back(0x00);
  cmac_input.push_back(0x00);
  cmac_input.push_back(0x00);
  cmac_input.push_back(0x80);

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

  if (!cmac_ctx) {
    std::cerr << "DeriveKey(): Failed to create context: " << GetSSLError()
              << std::endl;
    return false;
  }

  if (!CMAC_Init(cmac_ctx, &key[0], key.size(), cipher, 0)) {
    std::cerr << "DeriveKey(): Failed to initialize: " << GetSSLError()
              << std::endl;
    CMAC_CTX_free(cmac_ctx);
    return false;
  }

  if (!CMAC_Update(cmac_ctx, &cmac_input[0], cmac_input.size())) {
    std::cerr << "DeriveKey(): Failed to update: " << GetSSLError()
              << std::endl;
    CMAC_CTX_free(cmac_ctx);
    return false;
  }

  size_t reslen;
  uint8_t res[128];
  if (!CMAC_Final(cmac_ctx, res, &reslen)) {
    std::cerr << "DeriveKey(): Failed to finalize: " << GetSSLError()
              << std::endl;
    CMAC_CTX_free(cmac_ctx);
    return false;
  }
  out->assign(res, res + reslen);
  CMAC_CTX_free(cmac_ctx);

  return true;
}

int main(int argc, char** argv) {
  if (argc != 2) {
    std::cerr << "usage: " << argv[0] << " <new_key_filename>" << std::endl;
    exit(0);
  }

  std::string filename = argv[1];
  std::ofstream new_key_file;
  new_key_file.open(filename, std::ios::binary);
  if (!new_key_file) {
    std::cerr << "unable to open " << filename << " for writing" << std::endl;
    exit(-1);
  }

  // patch the keybox with version 3 and system ID
  std::copy(new_keybox_version.begin(), new_keybox_version.end(),
            test_keybox.begin() + KEYBOX_VERSION_OFFSET);
  std::copy(new_system_id.begin(), new_system_id.end(),
            test_keybox.begin() + SYSTEM_ID_OFFSET);

  // context is Device Key || CA token
  std::vector<uint8_t> context(test_keybox.begin() + DEVICE_KEY_OFFSET,
                               test_keybox.begin() + CA_TOKEN_END);

  // derive the new device key
  std::vector<uint8_t> new_device_key;
  if (!DeriveKey(renewal_key, context, &new_device_key)) {
    std::cerr << "Failed to derive new renewal key" << std::endl;
    exit(-1);
  } else {
    std::ostream_iterator<uint8_t> output_iterator(new_key_file);
    std::copy(new_device_key.begin(), new_device_key.end(), output_iterator);
    new_key_file.close();
    std::cout << "New key written to " << filename << std::endl;
    ;
  }
}