//
// Copyright 2022 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//

#define LOG_TAG "wv_factory_extraction_tool"

#include <cppbor.h>
#include <cppbor_parse.h>
#include <sys/random.h>

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <iterator>
#include <utility>
#include <vector>

#include "BccParser.h"
#include "WidevineProvisioner.h"
#include "log.h"
#include "properties.h"

constexpr size_t kChallengeSize = 64;

// The Google root key for the Endpoint Encryption Key chain, encoded as
// COSE_Sign1
inline constexpr uint8_t kCoseEncodedRootCert[] = {
    0x84, 0x43, 0xa1, 0x01, 0x27, 0xa0, 0x58, 0x2a, 0xa4, 0x01, 0x01, 0x03,
    0x27, 0x20, 0x06, 0x21, 0x58, 0x20, 0x99, 0xb9, 0xee, 0xdd, 0x5e, 0xe4,
    0x52, 0xf6, 0x85, 0xc6, 0x4c, 0x62, 0xdc, 0x3e, 0x61, 0xab, 0x57, 0x48,
    0x7d, 0x75, 0x37, 0x29, 0xad, 0x76, 0x80, 0x32, 0xd2, 0xb3, 0xcb, 0x63,
    0x58, 0xd9, 0x58, 0x40, 0x1e, 0x22, 0x08, 0x4b, 0xa4, 0xb7, 0xa4, 0xc8,
    0xd7, 0x4e, 0x03, 0x0e, 0xfe, 0xb8, 0xaf, 0x14, 0x4c, 0xa7, 0x3b, 0x6f,
    0xa5, 0xcd, 0xdc, 0xda, 0x79, 0xc6, 0x2b, 0x64, 0xfe, 0x99, 0x39, 0xaf,
    0x76, 0xe7, 0x80, 0xfa, 0x66, 0x00, 0x85, 0x0d, 0x07, 0x98, 0x2a, 0xac,
    0x91, 0x5c, 0xa7, 0x25, 0x14, 0x49, 0x06, 0x34, 0x75, 0xca, 0x8a, 0x27,
    0x7a, 0xd9, 0xe3, 0x5a, 0x49, 0xeb, 0x02, 0x03};

// The Google Endpoint Encryption Key certificate, encoded as COSE_Sign1
inline constexpr uint8_t kCoseEncodedGeekCert[] = {
    0x84, 0x43, 0xa1, 0x01, 0x27, 0xa0, 0x58, 0x4e, 0xa5, 0x01, 0x01, 0x02,
    0x58, 0x20, 0xd0, 0xae, 0xc1, 0x15, 0xca, 0x2a, 0xcf, 0x73, 0xae, 0x6b,
    0xcc, 0xcb, 0xd1, 0x96, 0x1d, 0x65, 0xe8, 0xb1, 0xdd, 0xd7, 0x4a, 0x1a,
    0x37, 0xb9, 0x43, 0x3a, 0x97, 0xd5, 0x99, 0xdf, 0x98, 0x08, 0x03, 0x38,
    0x18, 0x20, 0x04, 0x21, 0x58, 0x20, 0xbe, 0x85, 0xe7, 0x46, 0xc4, 0xa3,
    0x42, 0x5a, 0x40, 0xd9, 0x36, 0x3a, 0xa6, 0x15, 0xd0, 0x2c, 0x58, 0x7e,
    0x3d, 0xdc, 0x33, 0x02, 0x32, 0xd2, 0xfc, 0x5e, 0x1e, 0x87, 0x25, 0x5f,
    0x72, 0x60, 0x58, 0x40, 0x9b, 0xcf, 0x90, 0xe2, 0x2e, 0x4b, 0xab, 0xd1,
    0x18, 0xb1, 0x0e, 0x8e, 0x5d, 0x20, 0x27, 0x4b, 0x84, 0x58, 0xfe, 0xfc,
    0x32, 0x90, 0x7e, 0x72, 0x05, 0x83, 0xbc, 0xd7, 0x82, 0xbe, 0xfa, 0x64,
    0x78, 0x2d, 0x54, 0x10, 0x4b, 0xc0, 0x31, 0xbf, 0x6b, 0xe8, 0x1e, 0x35,
    0xe2, 0xf0, 0x2d, 0xce, 0x6c, 0x2f, 0x4f, 0xf2, 0xf5, 0x4f, 0xa5, 0xd4,
    0x83, 0xad, 0x96, 0xa2, 0xf1, 0x87, 0x58, 0x04};

std::vector<uint8_t> generateChallenge() {
  std::vector<uint8_t> challenge(kChallengeSize);

  ssize_t bytesRemaining = static_cast<ssize_t>(challenge.size());
  uint8_t* writePtr = challenge.data();
  while (bytesRemaining > 0) {
    int bytesRead = getrandom(writePtr, bytesRemaining, /*flags=*/0);
    if (bytesRead < 0) {
      if (errno == EINTR) {
        continue;
      } else {
        std::cerr << errno << ": " << strerror(errno) << std::endl;
        exit(-1);
      }
    }
    bytesRemaining -= bytesRead;
    writePtr += bytesRead;
  }

  return challenge;
}

std::vector<uint8_t> getEekChain() {
  cppbor::Array chain;
  chain.add(cppbor::EncodedItem(std::vector<uint8_t>(
      std::begin(kCoseEncodedRootCert), std::end(kCoseEncodedRootCert))));
  chain.add(cppbor::EncodedItem(std::vector<uint8_t>(
      std::begin(kCoseEncodedGeekCert), std::end(kCoseEncodedGeekCert))));
  return chain.encode();
}

cppbor::Array composeCertificateRequest(
    const std::vector<uint8_t>& protectedData,
    const std::vector<uint8_t>& verifiedDeviceInfo,
    const std::vector<uint8_t>& challenge) {
  cppbor::Array macedKeysToSign =
      cppbor::Array()
          .add(std::vector<uint8_t>(0))   // empty protected headers as bstr
          .add(cppbor::Map())             // empty unprotected headers
          .add(cppbor::Null())            // nil for the payload
          .add(std::vector<uint8_t>(0));  // MAC as returned from the HAL

  cppbor::Array deviceInfo = cppbor::Array()
                                 .add(cppbor::EncodedItem(verifiedDeviceInfo))
                                 .add(cppbor::Map());  // Empty device info

  cppbor::Array certificateRequest =
      cppbor::Array()
          .add(std::move(deviceInfo))
          .add(challenge)
          .add(cppbor::EncodedItem(protectedData))
          .add(std::move(macedKeysToSign));
  return certificateRequest;
}

cppbor::Array getCsr(widevine::WidevineProvisioner& provisioner) {
  const std::vector<uint8_t> challenge = generateChallenge();
  std::vector<uint8_t> verifiedDeviceInfo;
  std::vector<uint8_t> protectedData;

  if (!provisioner.GenerateCertificateRequest(
          false, getEekChain(), verifiedDeviceInfo, protectedData)) {
    std::cerr << "Failed to generate certificate request." << std::endl;
    exit(-1);
  }
  auto csr =
      composeCertificateRequest(protectedData, verifiedDeviceInfo, challenge);
  return csr;
}

std::unique_ptr<cppbor::Array> composeCertificateRequestV3(
    const std::vector<uint8_t>& csr) {
  auto [parsedCsr, _, csrErrMsg] = cppbor::parse(csr);
  if (!parsedCsr) {
    LOGE("Failed to parse input CSR.");
    return nullptr;
  }
  if (!parsedCsr->asArray()) {
    LOGE("Input CSR is not a CBOR array.");
    return nullptr;
  }
  std::string fingerPrint;
  if (!wvcdm::Properties::GetBuildInfo(&fingerPrint)) {
    LOGE("Failed to get finger print.");
    return nullptr;
  }

  cppbor::Map unverifiedDeviceInfo =
      cppbor::Map().add("fingerprint", cppbor::Tstr(fingerPrint));
  parsedCsr->asArray()->add(std::move(unverifiedDeviceInfo));
  return std::unique_ptr<cppbor::Array>(parsedCsr.release()->asArray());
}

std::unique_ptr<cppbor::Array> getCsrV3(
    widevine::WidevineProvisioner& provisioner) {
  const std::vector<uint8_t> challenge = generateChallenge();
  std::vector<uint8_t> csr;
  if (!provisioner.GenerateCertificateRequestV2(challenge, &csr)) {
    std::cerr << "Failed to generate certificate request v2." << std::endl;
    exit(-1);
  }
  return composeCertificateRequestV3(csr);
}

int main(int argc, char** argv) {
  if (argc < 2) {
    fprintf(stderr, "%s <bcc|bcc_str|device_info|csr|csr_v3>\n", argv[0]);
    return 0;
  }
  widevine::WidevineProvisioner provisioner;
  if (!std::strcmp(argv[1], "bcc")) {
    auto bcc = provisioner.GetBcc();
    fwrite(bcc.data(), 1, bcc.size(), stdout);
    fflush(stdout);
  } else if (!std::strcmp(argv[1], "bcc_str")) {
    auto bcc = provisioner.GetBcc();
    widevine::BccParser bcc_parser;
    std::string parsed_bcc = bcc_parser.Parse(bcc);
    std::copy(parsed_bcc.begin(), parsed_bcc.end(),
              std::ostream_iterator<char>(std::cout));
  } else if (!std::strcmp(argv[1], "device_info")) {
    std::vector<uint8_t> deviceInfo;
    if (provisioner.GetDeviceInfo(deviceInfo)) {
      fwrite(deviceInfo.data(), 1, deviceInfo.size(), stdout);
      fflush(stdout);
    }
  } else if (!std::strcmp(argv[1], "csr")) {
    auto csr = getCsr(provisioner);
    auto bytes = csr.encode();
    std::copy(bytes.begin(), bytes.end(),
              std::ostream_iterator<char>(std::cout));
  } else if (!std::strcmp(argv[1], "csr_v3")) {
    auto csr = getCsrV3(provisioner);
    if (csr != nullptr) {
      auto bytes = csr->encode();
      std::copy(bytes.begin(), bytes.end(),
                std::ostream_iterator<char>(std::cout));
    }
  }
  return 0;
}