// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
#include <arpa/inet.h>
#include <ctype.h>
#include <errno.h>
#include <inttypes.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include "OEMCryptoCENC.h"
#include "string_conversions.h"

namespace {
// Size of a valid keybox.
constexpr size_t kKeyboxSize = 128;
// Size of a valid keybox token section.
constexpr size_t kKeyDataSize = 72;
// Offset of the system ID in key data.
constexpr size_t kSystemIdOffset = 4;
// Offset of the keybox version in key data.
constexpr size_t kVersionOffset = 0;

// == Utils ==

bool IsRegularFile(const std::string& path) {
  struct stat st;
  if (stat(path.c_str(), &st) != 0) {
    if (errno == ENOENT || errno == ENOTDIR) {
      std::cerr << "File does not exist: path = " << path << std::endl;
      return false;
    }
    std::cerr << "Failed to call stat: path = " << path;
    std::cerr << ", errno = " << errno << std::endl;
    return false;
  }
  if (!S_ISREG(st.st_mode)) {
    std::cerr << "Not a regular file: path = " << path << ", mode = ";
    std::cerr << std::setfill('0') << std::setw(7) << std::oct << st.st_mode;
    std::cerr << std::endl;
    return false;
  }
  return true;
}

void PrintDeviceId(const uint8_t* device_id_u8, size_t device_id_length) {
  std::string device_id(reinterpret_cast<const char*>(device_id_u8),
                        device_id_length);
  // Trim null bytes.
  while (!device_id.empty() && device_id.back() == '\0') device_id.pop_back();
  // Check if empty.
  if (device_id.empty()) {
    std::cerr << "Device ID was all null bytes: ";
    std::cerr << "length = " << device_id_length << std::endl;
    std::cout << "device_id = <empty>" << std::endl;
    return;
  }
  // Check if printable.
  if (!std::all_of(device_id.begin(), device_id.end(), ::isprint)) {
    device_id = wvutil::b2a_hex(device_id);
  }
  std::cout << "device_id = " << device_id << std::endl;
}

void PrintSystemId(const uint8_t* key_data) {
  // Assumes that |key_data| length as already been verified.
  const uint32_t* system_id_ptr =
      reinterpret_cast<const uint32_t*>(&key_data[kSystemIdOffset]);
  const uint32_t system_id = ntohl(*system_id_ptr);
  std::cout << "system_id = " << system_id << std::endl;
  std::cout << "hex(system_id) = 0x";
  std::cout << std::setfill('0') << std::setw(8) << std::hex << system_id;
  std::cout << std::endl;
}

void PrintKeyboxVersion(const uint8_t* key_data) {
  // Assumes that |key_data| length as already been verified.
  const uint32_t* version_ptr =
      reinterpret_cast<const uint32_t*>(&key_data[kVersionOffset]);
  const uint32_t version = ntohl(*version_ptr);
  std::cout << "version = " << version << std::endl;
}

// == Primary ==

bool RetrieveKeybox(const std::string& path, std::vector<uint8_t>* keybox) {
  using StreamIter = std::istreambuf_iterator<char>;
  using PosType = std::iostream::pos_type;
  std::ifstream fin(path, std::ios::in | std::ios::binary);
  if (!fin) {
    std::cerr << "Failed to open input file: path = " << path << std::endl;
    return false;
  }

  // Verify size.
  fin.seekg(0, std::ios::end);
  if (!fin) {
    std::cerr << "Failed to seek to end of file: path = " << path << std::endl;
    return false;
  }
  const PosType keybox_size_pt = fin.tellg();
  if (keybox_size_pt == PosType(-1)) {
    std::cerr << "Failed to obtain the file size: path = " << path << std::endl;
    return false;
  }
  const size_t keybox_size = static_cast<size_t>(keybox_size_pt);
  fin.seekg(0, std::ios::beg);
  if (!fin) {
    std::cerr << "Failed to seek to beginning of file: path = ";
    std::cerr << path << std::endl;
    return false;
  }
  if (keybox_size != kKeyboxSize) {
    std::cerr << "Unexpected keybox size: ";
    std::cerr << "size = " << keybox_size;
    std::cerr << ", expected = " << kKeyboxSize;
    std::cerr << ", path = " << path << std::endl;
    return false;
  }

  // Read keybox data.
  keybox->clear();
  keybox->reserve(kKeyboxSize);
  keybox->assign(StreamIter(fin), StreamIter());

  if (keybox->size() != kKeyboxSize) {
    std::cerr << "Failed to read entire keybox: ";
    std::cerr << "read = " << keybox->size();
    std::cerr << ", expected = " << kKeyboxSize;
    std::cerr << ", path = " << path << std::endl;
    keybox->clear();
    return false;
  }
  return true;
}

bool InstallKeyboxAndPrintInfo(const std::vector<uint8_t>& keybox) {
  std::cout << "Install keybox: " << wvutil::b2a_hex(keybox) << std::endl;

  // Step 1:  Initialize.
  OEMCryptoResult result = OEMCrypto_Initialize();
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to initialize: result = " << result << std::endl;
    return false;
  }
  std::cout << "OEMCrypto initialized" << std::endl;

  // Step 2: Install keybox.
  const OEMCrypto_ProvisioningMethod method = OEMCrypto_GetProvisioningMethod();
  if (method != OEMCrypto_Keybox) {
    std::cerr << "OEMCrypto is not keybox type: method = ";
    std::cerr << method << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  result = OEMCrypto_InstallKeyboxOrOEMCert(keybox.data(), keybox.size());
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to install keybox: result = " << result << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  std::cout << "OEMCrypto keybox installed" << std::endl;

  // Step 3: Verify device ID.
  uint8_t buffer[128] = {};
  memset(buffer, 0, sizeof(buffer));
  size_t buffer_length = sizeof(buffer);
  result = OEMCrypto_GetDeviceID(buffer, &buffer_length);
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to get device ID: result = " << result << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  PrintDeviceId(buffer, buffer_length);

  // Step 4: Verify system ID and keybox version.
  memset(buffer, 0, sizeof(buffer));
  buffer_length = sizeof(buffer);

  result = OEMCrypto_GetKeyData(buffer, &buffer_length);
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to get key data: result = " << result << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  if (buffer_length != kKeyDataSize) {
    std::cerr << "Unexpected key data size: ";
    std::cerr << "size = " << buffer_length;
    std::cerr << ", expected = " << kKeyDataSize << std::endl;
    OEMCrypto_Terminate();
    return false;
  }
  PrintSystemId(buffer);
  PrintKeyboxVersion(buffer);

  // Step 5: Cleanup.
  result = OEMCrypto_Terminate();
  if (result != OEMCrypto_SUCCESS) {
    std::cerr << "Failed to terminate: result = " << result << std::endl;
    return false;
  }
  std::cout << "OEMCrypto terminated" << std::endl;
  return true;
}
}  // namespace

int main(int argc, char** argv) {
  if (argc != 2) {
    std::cerr << "Usage: " << argv[0] << " <filename>" << std::endl;
    return 1;
  }
  const std::string keybox_path = argv[1];
  if (!IsRegularFile(keybox_path)) {
    std::cerr << "Bad keybox path: " << keybox_path << std::endl;
    return 1;
  }
  std::vector<uint8_t> keybox;
  if (!RetrieveKeybox(keybox_path, &keybox)) {
    std::cerr << "Failed to retrieve keybox" << std::endl;
    return 1;
  }
  if (!InstallKeyboxAndPrintInfo(keybox)) {
    std::cerr << "Failed to install keybox" << std::endl;
    return 1;
  }
  return 0;
}