////////////////////////////////////////////////////////////////////////////////
// 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 <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"

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

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

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

const size_t kTsPacketSize = 188;

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

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);

  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, &ecm);
  } else {
    status = wv_cas_ecm.GenerateSingleKeyEcm(content_keys[0],
                                             kDefaultTrackTypeSd, &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[kTsPacketSize];
  uint8_t continuity_counter;  // not used.
  status = wv_cas_ecm.GenerateTsPacket(ecm, kEcmPid, kTableId,
                                       &continuity_counter, packet);
  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 < kTsPacketSize; 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), kTsPacketSize);
  file.close();

  return 0;
}