//
// 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 <algorithm>
#include <cstdio>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <iterator>
#include <utility>
#include <vector>

#include <sys/random.h>

#include "WidevineProvisioner.h"

constexpr size_t kChallengeSize = 16;

// 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;
}

int main(int argc, char** argv) {
  if (argc<2) {
    fprintf(stderr,"%s <bcc|device_info>\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],"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));
  }
  return 0;
}