////////////////////////////////////////////////////////////////////////////////
// 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 <cassert>
#include <fstream>
#include <iostream>
#include <string>

#include "gflags/gflags.h"
#include "media_cas_packager_sdk/public/wv_cas_ecm.h"
#include "media_cas_packager_sdk/public/wv_cas_types.h"

DEFINE_int32(content_iv_size, 8, "Content IV size");
DEFINE_bool(key_rotation, true, "Whether key rotation is enabled");
DEFINE_string(crypto_mode, "CSA2", "Only CBC, CTR, or CSA2 is allowed");
DEFINE_int32(ecm_pid, 149, "PID for the ECM packet");
DEFINE_string(output_file, "",
              "If specified, generated ECM TS packet will be written to the "
              "specified output file path");

const char kCsaEvenKey[] = "even_key";                    // 8 bytes
const char kEvenContentIv8Bytes[] = "even_iv.";           // 8 bytes
const char kEvenEntitlementKeyId[] = "fake_key_id1....";  // 16 bytes
const char kEvenEntitlementKey[] =
    "fakefakefakefakefakefakefake1...";                  // 32 bytes
const char kCsaOddKey[] = "odd_key.";                    // 8 bytes
const char kOddContentIv8Bytes[] = "odd_iv..";           // 8 bytes
const char kOddEntitlementKeyId[] = "fake_key_id2....";  // 16 bytes
const char kOddEntitlementKey[] =
    "fakefakefakefakefakefakefake2...";  // 32 bytes
const size_t kTsPacketSize = 188;

widevine::cas::CryptoMode GetCryptoMode() {
  if (FLAGS_crypto_mode.compare("CBC") == 0) {
    return widevine::cas::CryptoMode::kAesCbc;
  }
  if (FLAGS_crypto_mode.compare("CTR") == 0) {
    return widevine::cas::CryptoMode::kAesCtr;
  }
  return widevine::cas::CryptoMode::kDvbCsa2;
}

int main(int argc, char **argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  // Generate ECM.
  widevine::cas::WvCasEcm wv_cas_ecm;
  widevine::cas::WvCasStatus status = wv_cas_ecm.Initialize(
      FLAGS_content_iv_size, FLAGS_key_rotation, GetCryptoMode());
  if (status != widevine::cas::OK) {
    std::cerr << "Failed to initialize WV CAS ECM, error: "
              << widevine::cas::GetWvCasStatusMessage(status)
              << std::endl;
  }
  std::string ecm;
  if (FLAGS_key_rotation) {
    status = wv_cas_ecm.GenerateEcm(
        kCsaEvenKey, kEvenContentIv8Bytes, kEvenEntitlementKeyId,
        kEvenEntitlementKey, kCsaOddKey, kOddContentIv8Bytes,
        kOddEntitlementKeyId, kOddEntitlementKey, &ecm);
  } else {
    status = wv_cas_ecm.GenerateSingleKeyEcm(kCsaEvenKey, kEvenContentIv8Bytes,
                                             kEvenEntitlementKeyId,
                                             kEvenEntitlementKey, &ecm);
  }
  if (status != widevine::cas::OK) {
    std::cerr << "Failed to generate WV CAS ECM, error: "
              << widevine::cas::GetWvCasStatusMessage(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, FLAGS_ecm_pid,
                                       /* table_id= */ 0x80,
                                       &continuity_counter, packet);
  if (status != widevine::cas::OK) {
    std::cerr << "Failed to create ECM TS packet" << 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.
  if (!FLAGS_output_file.empty()) {
    std::ofstream file;
    file.open(FLAGS_output_file.c_str(), std::ios_base::binary);
    assert(file.is_open());
    file.write(reinterpret_cast<char *>(packet), kTsPacketSize);
    file.close();
  }

  return 0;
}