////////////////////////////////////////////////////////////////////////////////
// Copyright 2018 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.
////////////////////////////////////////////////////////////////////////////////

// Example of how to use the wv_cas_ecm library.

#include <stddef.h>
#include <stdio.h>

#include <cassert>
#include <cstdint>
#include <fstream>
#include <iostream>
#include <string>

#include <cstdint>
#include "common/status.h"
#include "media_cas_packager_sdk/public/wv_cas_ecm.h"
#include "media_cas_packager_sdk/public/wv_cas_types.h"

namespace {

constexpr int kEcmVersion = 3;  // Choices are 2 and 3.

constexpr size_t kContentIvSize = 16;  // 8 or 16
constexpr bool kKeyRotation = true;    // whether key rotation is enabled
constexpr char kCryptoMode[] =
    "AesScte";  // "AesCbc", "AesCtr", "DvbCsa2", "DvbCsa3", "AesOfb", "AesScte"
constexpr int kEcmPid = 149;        // PID for the ECM packet
constexpr int kAgeRestriction = 0;  // Age restriction for the ECM
constexpr char kOutputFile[] =
    "/tmp/ecm.ts";                // ECM TS packet will be output to here
constexpr uint8_t kTableId = 0x80;  // 0x80 or 0x81
constexpr char kDefaultTrackTypeSd[] = "SD";

constexpr char kEvenKey[] = "even_key........";               // 16 bytes
constexpr char kEvenKeyId[] = "even_key_id.....";             // 16 bytes
constexpr char kEvenContentIv8Bytes[] = "even_iv.";           // 8 bytes
constexpr char kEvenContentIv16Bytes[] = "even_iv.even_iv.";  // 16 bytes
constexpr char kEvenEntitlementKeyId[] = "fake_key_id1....";  // 16 bytes
constexpr char kEvenEntitlementKey[] =
    "fakefakefakefakefakefakefake1...";             // 32 bytes
constexpr char kEvenWrapIv[] = "even_warp_iv....";  // 16 bytes

constexpr char kOddKey[] = "odd_key.........";               // 16 bytes
constexpr char kOddKeyId[] = "odd_key_id......";             // 16 bytes
constexpr char kOddContentIv8Bytes[] = "odd_iv..";           // 8 bytes
constexpr char kOddContentIv16Bytes[] = "odd_iv..odd_iv..";  // 16 bytes
constexpr char kOddEntitlementKeyId[] = "fake_key_id2....";  // 16 bytes
constexpr char kOddEntitlementKey[] =
    "fakefakefakefakefakefakefake2...";            // 32 bytes
constexpr char kOddWrapIv[] = "odd_warp_iv.....";  // 16 bytes

// Fingerprinting and service blocking are only available with ECM v3+.
constexpr char kFingerprintingControl[] = "ctr";
constexpr char kServiceBlockingDeviceGroup1[] = "group";
constexpr char kServiceBlockingDeviceGroup2[] = "g2";
constexpr unsigned char kTestECPrivateKey2Secp256r1[] = {
    0x30, 0x77, 0x02, 0x01, 0x01, 0x04, 0x20, 0x34, 0x9a, 0xf2, 0x95,
    0x94, 0xd4, 0xca, 0xb9, 0xa0, 0x81, 0xe4, 0x1c, 0xf5, 0xde, 0x8d,
    0x23, 0xf6, 0x79, 0xba, 0x3c, 0x6e, 0xc9, 0x0b, 0x56, 0x0f, 0x07,
    0x5e, 0x9f, 0xe9, 0x38, 0x18, 0xfc, 0xa0, 0x0a, 0x06, 0x08, 0x2a,
    0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0xa1, 0x44, 0x03, 0x42,
    0x00, 0x04, 0x7b, 0x2a, 0x61, 0x59, 0xe5, 0x1b, 0xb6, 0x30, 0x7b,
    0x59, 0x98, 0x42, 0x59, 0x37, 0xfb, 0x46, 0xfe, 0x53, 0xfe, 0x32,
    0xa1, 0x5c, 0x93, 0x36, 0x11, 0xb0, 0x5a, 0xae, 0xa4, 0x48, 0xe3,
    0x20, 0x12, 0xce, 0x78, 0xa7, 0x7f, 0xfd, 0x73, 0x5e, 0x09, 0x77,
    0x53, 0x77, 0x8f, 0xd6, 0x1b, 0x26, 0xfa, 0xc4, 0x2c, 0xc4, 0x11,
    0xfa, 0x72, 0x6a, 0xbe, 0x94, 0x78, 0x4d, 0x74, 0x20, 0x27, 0x86};

const size_t kTsPacketSize = 188;
const size_t kMaxPossibleTsPacketsSizeBytes = 6 * kTsPacketSize;

using widevine::cas::EcmVersion;
using widevine::cas::EntitlementKeyInfo;
using widevine::cas::WvCasContentKeyInfo;
using widevine::cas::WvCasEcmParameters;

WvCasEcmParameters CreateWvCasEcmParameters(bool key_rotation,
                                            int content_iv_size) {
  WvCasEcmParameters params;
  params.content_iv_size = content_iv_size == 8
                               ? widevine::cas::kIvSize8
                               : widevine::cas::kIvSize16;
  params.key_rotation_enabled = key_rotation;
  if (!widevine::cas::StringToCryptoMode(kCryptoMode,
                                               &params.crypto_mode)) {
    std::cerr << "Unsupported crypto mode " << kCryptoMode << std::endl;
  }
  params.age_restriction = kAgeRestriction;
  params.ecm_version = kEcmVersion <= 2 ? EcmVersion::kV2 : EcmVersion::kV3;
  params.ecc_private_signing_key.assign(std::begin(kTestECPrivateKey2Secp256r1),
                                        std::end(kTestECPrivateKey2Secp256r1));
  return params;
}

std::vector<EntitlementKeyInfo> CreateInjectedEntitlements(bool key_rotation) {
  std::vector<EntitlementKeyInfo> injected_entitlements;
  injected_entitlements.reserve(key_rotation ? 2 : 1);
  injected_entitlements.emplace_back();
  EntitlementKeyInfo* entitlement = &injected_entitlements.back();
  entitlement->key_id = kEvenEntitlementKeyId;
  entitlement->key_value = kEvenEntitlementKey;
  entitlement->is_even_key = true;
  entitlement->track_type = kDefaultTrackTypeSd;
  if (key_rotation) {
    injected_entitlements.emplace_back();
    EntitlementKeyInfo* entitlement = &injected_entitlements.back();
    entitlement->key_id = kOddEntitlementKeyId;
    entitlement->key_value = kOddEntitlementKey;
    entitlement->is_even_key = false;
    entitlement->track_type = kDefaultTrackTypeSd;
  }
  return injected_entitlements;
}

std::vector<WvCasContentKeyInfo> CreateContentKeyInfo(bool key_rotation,
                                                      int content_iv_size) {
  std::vector<WvCasContentKeyInfo> content_keys;
  content_keys.reserve(key_rotation ? 2 : 1);
  content_keys.emplace_back();
  WvCasContentKeyInfo* content_key = &content_keys.back();
  content_key->key = kEvenKey;
  content_key->key_id = kEvenKeyId;
  content_key->content_iv =
      content_iv_size == 8 ? kEvenContentIv8Bytes : kEvenContentIv16Bytes;
  content_key->wrapped_key_iv = kEvenWrapIv;
  if (key_rotation) {
    content_keys.emplace_back();
    WvCasContentKeyInfo* content_key = &content_keys.back();
    content_key->key = kOddKey;
    content_key->key_id = kOddKeyId;
    content_key->content_iv =
        content_iv_size == 8 ? kOddContentIv8Bytes : kOddContentIv16Bytes;
    content_key->wrapped_key_iv = kOddWrapIv;
  }

  return content_keys;
}

}  // namespace

int main(int argc, char** argv) {
  WvCasEcmParameters params =
      CreateWvCasEcmParameters(kKeyRotation, kContentIvSize);
  std::vector<EntitlementKeyInfo> entitlements =
      CreateInjectedEntitlements(kKeyRotation);
  widevine::cas::WvCasEcm wv_cas_ecm(params, entitlements);

  if (params.ecm_version == EcmVersion::kV3) {
    widevine::cas::EcmFingerprintingParams fingerprinting_params;
    fingerprinting_params.control = kFingerprintingControl;
    wv_cas_ecm.SetFingerprinting(&fingerprinting_params);

    widevine::cas::EcmServiceBlockingParams service_blocking_params;
    service_blocking_params.device_groups = {kServiceBlockingDeviceGroup1,
                                             kServiceBlockingDeviceGroup2};
    wv_cas_ecm.SetServiceBlocking(&service_blocking_params);
  }

  std::vector<WvCasContentKeyInfo> content_keys =
      CreateContentKeyInfo(kKeyRotation, kContentIvSize);
  std::string ecm;
  widevine::Status status;
  if (kKeyRotation) {
    status =
        wv_cas_ecm.GenerateEcm(content_keys[0], content_keys[1],
                               kDefaultTrackTypeSd, /*group_ids=*/{}, &ecm);
  } else {
    status = wv_cas_ecm.GenerateSingleKeyEcm(
        content_keys[0], kDefaultTrackTypeSd, /*group_ids=*/{}, &ecm);
  }

  if (!status.ok()) {
    std::cerr << "Failed to generate WV CAS ECM, error: " << status
              << std::endl;
    return -1;
  } else {
    std::cout << "ECM size: " << ecm.size() << std::endl;
    std::cout << "ECM bytes: ";
    for (size_t i = 0; i < ecm.size(); i++) {
      printf("'\\x%02x', ", static_cast<uint16_t>(ecm.at(i)));
    }
    std::cout << std::endl;
  }
  // Generate ECM TS Packet.
  uint8_t packet[kMaxPossibleTsPacketsSizeBytes];
  ssize_t packet_size = kMaxPossibleTsPacketsSizeBytes;
  uint8_t continuity_counter;  // not used.
  status = wv_cas_ecm.GenerateTsPacket(
      ecm, kEcmPid, kTableId, &continuity_counter, packet, &packet_size);
  if (!status.ok()) {
    std::cerr << "Failed to create ECM TS packet: " << status << std::endl;
    return -1;
  } else {
    std::cout << "TS packet bytes: ";
    for (size_t i = 0; i < packet_size; i++) {
      printf("'\\x%02x', ", static_cast<uint16_t>(packet[i]));
    }
    std::cout << std::endl;
  }
  // Write ECM TS Packet to a file.
  std::ofstream file;
  file.open(kOutputFile, std::ios_base::binary);
  assert(file.is_open());
  file.write(reinterpret_cast<char*>(packet), packet_size);
  file.close();

  return 0;
}