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

#include "ecm_parser_v3.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <vector>

#include "media_cas.pb.h"

namespace wvcas {
namespace {

using video_widevine::EcmGroupKeyData;
using video_widevine::EcmMetaData;
using video_widevine::EcmPayload;
using video_widevine::SignedEcmPayload;

constexpr int kEcmVersion = 3;
constexpr uint16_t kWidevineCasId = 0x4AD4;
constexpr int kEcmHeaderSizeByte = 3;

constexpr char kWrappedKeyIv[] = "wrapped_key_iv..";
constexpr char kWrappedKeyIv2[] = "wrapped_key_iv.2";
constexpr char kEntitlementId[] = "entitlement_id..";
constexpr char kEntitlementId2[] = "entitlement_id.2";
constexpr char kContentIv[] = "c_iv....c_iv....";
constexpr char kContentIv2[] = "c_iv....c_iv...2";
constexpr char kWrappedContentKey[] = "wrapped_key.....";
constexpr char kWrappedContentKey2[] = "wrapped_key....2";

void WriteEcmHeader(std::vector<uint8_t>* ecm) {
  ecm->push_back(kWidevineCasId >> 8);
  ecm->push_back(kWidevineCasId & 0xff);
  ecm->push_back(kEcmVersion);
}

std::vector<uint8_t> GenerateEcm(const SignedEcmPayload& signed_ecm_payload) {
  std::vector<uint8_t> ecm;
  WriteEcmHeader(&ecm);
  ecm.resize(kEcmHeaderSizeByte + signed_ecm_payload.ByteSize());
  signed_ecm_payload.SerializeToArray(ecm.data() + kEcmHeaderSizeByte,
                                      signed_ecm_payload.ByteSize());
  return ecm;
}

TEST(EcmParserV3Test, CreateWithEmptyEcmFail) {
  std::vector<uint8_t> ecm;
  EXPECT_TRUE(EcmParserV3::Create(ecm) == nullptr);
}

TEST(EcmParserV3Test, CreateWithOnlyEcmHeaderFail) {
  std::vector<uint8_t> ecm;
  WriteEcmHeader(&ecm);
  EXPECT_TRUE(EcmParserV3::Create(ecm) == nullptr);
}

TEST(EcmParserV3Test, CreateWithInvalidSignedEcmPayloadFail) {
  std::vector<uint8_t> ecm;
  WriteEcmHeader(&ecm);
  // appends some chars as payload
  ecm.resize(100, 'x');

  EXPECT_TRUE(EcmParserV3::Create(ecm) == nullptr);
}

TEST(EcmParserV3Test, CreateWithInvalidSerializedEcmFail) {
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload("invalid");
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  EXPECT_TRUE(EcmParserV3::Create(ecm) == nullptr);
}

TEST(EcmParserV3Test, CreateWithEvenKeySuccess) {
  EcmPayload ecm_payload;
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_data(kWrappedContentKey);
  ecm_payload.mutable_even_key_data()->set_content_iv(kContentIv);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_EQ(parser->version(), kEcmVersion);
  EXPECT_EQ(parser->age_restriction(), 0);
  EXPECT_EQ(parser->crypto_mode(), CryptoMode::kInvalid);
  EXPECT_FALSE(parser->has_fingerprinting());
  EXPECT_FALSE(parser->has_service_blocking());
  EXPECT_EQ(parser->ecm_serialized_payload(), ecm_payload.SerializeAsString());
  EXPECT_TRUE(parser->signature().empty());
  EXPECT_FALSE(parser->rotation_enabled());
  EXPECT_EQ(parser->content_iv_size(), 16);
  std::vector<uint8_t> result =
      parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey);
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv);
}

TEST(EcmParserV3Test, CreateWithEvenOddKeysSuccess) {
  EcmPayload ecm_payload;
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_data(kWrappedContentKey);
  ecm_payload.mutable_even_key_data()->set_content_iv(kContentIv);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv);
  ecm_payload.mutable_odd_key_data()->set_entitlement_key_id(kEntitlementId2);
  ecm_payload.mutable_odd_key_data()->set_wrapped_key_data(kWrappedContentKey2);
  ecm_payload.mutable_odd_key_data()->set_content_iv(kContentIv2);
  ecm_payload.mutable_odd_key_data()->set_wrapped_key_iv(kWrappedKeyIv2);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->rotation_enabled());
  std::vector<uint8_t> result =
      parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey);
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv);
  result = parser->entitlement_key_id(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId2);
  result = parser->wrapped_key_data(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey2);
  result = parser->content_iv(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv2);
  result = parser->wrapped_key_iv(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv2);
}

TEST(EcmParserV3Test, CreateWithOmittedOddKeyFieldsSuccess) {
  EcmPayload ecm_payload;
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_data(kWrappedContentKey);
  ecm_payload.mutable_even_key_data()->set_content_iv(kContentIv);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv);
  ecm_payload.mutable_odd_key_data()->set_wrapped_key_data(kWrappedContentKey2);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->rotation_enabled());
  std::vector<uint8_t> result =
      parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey);
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv);
  result = parser->entitlement_key_id(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey2);
  result = parser->content_iv(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv);
  result = parser->wrapped_key_iv(KeySlotId::kOddKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv);
}

TEST(EcmParserV3Test, AgeRestrictionSuccess) {
  const int expected_age_restriction = 3;
  EcmPayload ecm_payload;
  ecm_payload.mutable_meta_data()->set_age_restriction(
      expected_age_restriction);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_EQ(parser->age_restriction(), expected_age_restriction);
}

class EcmParserV3AgeRestrictionTest
    : public testing::Test,
      public testing::WithParamInterface<uint8_t> {};

TEST_P(EcmParserV3AgeRestrictionTest, ExpectedAgeRestriction) {
  const uint8_t expected_age_restriction = GetParam();
  EcmPayload ecm_payload;
  ecm_payload.mutable_meta_data()->set_age_restriction(
      expected_age_restriction);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_EQ(parser->age_restriction(), expected_age_restriction);
}

INSTANTIATE_TEST_SUITE_P(EcmParserV3AgeRestrictionTest,
                         EcmParserV3AgeRestrictionTest,
                         testing::Values(0, 3, 18));

class EcmParserV3CipherModeTest
    : public testing::Test,
      public testing::WithParamInterface<
          testing::tuple<CryptoMode, EcmMetaData::CipherMode>> {};

TEST_P(EcmParserV3CipherModeTest, ExpectedCipherMode) {
  const CryptoMode expected = std::get<0>(GetParam());
  EcmPayload ecm_payload;
  ecm_payload.mutable_meta_data()->set_cipher_mode(std::get<1>(GetParam()));
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_EQ(parser->crypto_mode(), expected);
}

INSTANTIATE_TEST_SUITE_P(
    EcmParserV3CipherModeTest, EcmParserV3CipherModeTest,
    testing::Values(
        std::make_tuple(CryptoMode::kAesCBC, EcmMetaData::AES_CBC),
        std::make_tuple(CryptoMode::kAesCTR, EcmMetaData::AES_CTR),
        std::make_tuple(CryptoMode::kDvbCsa2, EcmMetaData::DVB_CSA2),
        std::make_tuple(CryptoMode::kDvbCsa3, EcmMetaData::DVB_CSA3),
        std::make_tuple(CryptoMode::kAesOFB, EcmMetaData::AES_OFB),
        std::make_tuple(CryptoMode::kAesSCTE, EcmMetaData::AES_SCTE52),
        std::make_tuple(CryptoMode::kAesECB, EcmMetaData::AES_ECB)));

TEST(EcmParserV3Test, FingerprintingSuccess) {
  EcmPayload ecm_payload;
  ecm_payload.mutable_fingerprinting()->set_control("control");
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->has_fingerprinting());
  EXPECT_EQ(parser->fingerprinting().SerializeAsString(),
            ecm_payload.fingerprinting().SerializeAsString());
}

TEST(EcmParserV3Test, ServiceBlockingSuccess) {
  EcmPayload ecm_payload;
  ecm_payload.mutable_service_blocking()->add_device_groups("group");
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->has_service_blocking());
  EXPECT_EQ(parser->service_blocking().SerializeAsString(),
            ecm_payload.service_blocking().SerializeAsString());
}

TEST(EcmParserV3Test, SignatureSuccess) {
  const std::string expected_signature = "signature";
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_signature(expected_signature);
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_EQ(parser->signature(), expected_signature);
}

TEST(EcmParserV3Test, SetGroupIdSuccess) {
  const std::string group_id = "group_id";
  const std::string group_id2 = "another_group";
  EcmPayload ecm_payload;
  EcmGroupKeyData* group_key_data = ecm_payload.add_group_key_data();
  group_key_data->set_group_id(group_id);
  group_key_data = ecm_payload.add_group_key_data();
  group_key_data->set_group_id(group_id2);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->set_group_id(group_id));
  EXPECT_TRUE(parser->set_group_id(group_id2));
}

TEST(EcmParserV3Test, SetUnknownGroupIdFail) {
  EcmPayload ecm_payload;
  EcmGroupKeyData* group_key_data = ecm_payload.add_group_key_data();
  group_key_data->set_group_id("group_id");
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_FALSE(parser->set_group_id("unknown"));
}

TEST(EcmParserV3Test, ParserWithGroupIdSuccess) {
  const std::string group_id = "group_id";
  EcmPayload ecm_payload;
  EcmGroupKeyData* group_key_data = ecm_payload.add_group_key_data();
  group_key_data->set_group_id(group_id);
  group_key_data->mutable_even_key_data()->set_entitlement_key_id(
      kEntitlementId);
  group_key_data->mutable_even_key_data()->set_wrapped_key_data(
      kWrappedContentKey);
  group_key_data->mutable_even_key_data()->set_content_iv(kContentIv);
  group_key_data->mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv);
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId2);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_data(
      kWrappedContentKey2);
  ecm_payload.mutable_even_key_data()->set_content_iv(kContentIv2);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv2);

  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  // If group Id is not set, the normal keys will be returned.
  std::vector<uint8_t> result =
      parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId2);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey2);
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv2);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv2);

  // Now set the group id.
  EXPECT_TRUE(parser->set_group_id(group_id));
  result = parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey);
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv);
}

TEST(EcmParserV3Test, ParserGroupKeysWithOmittedFieldsSuccess) {
  const std::string group_id = "group_id";
  EcmPayload ecm_payload;
  EcmGroupKeyData* group_key_data = ecm_payload.add_group_key_data();
  group_key_data->set_group_id(group_id);
  group_key_data->mutable_even_key_data()->set_entitlement_key_id(
      kEntitlementId);
  group_key_data->mutable_even_key_data()->set_wrapped_key_data(
      kWrappedContentKey);
  // Content IV and wrapped key iv is omitted in |group_key_data|/
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId2);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_data(
      kWrappedContentKey2);
  ecm_payload.mutable_even_key_data()->set_content_iv(kContentIv2);
  ecm_payload.mutable_even_key_data()->set_wrapped_key_iv(kWrappedKeyIv2);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);
  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->set_group_id(group_id));

  std::vector<uint8_t> result =
      parser->entitlement_key_id(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kEntitlementId);
  result = parser->wrapped_key_data(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedContentKey);
  // Content IV and wrapped key iv are from normal non-group key.
  result = parser->content_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kContentIv2);
  result = parser->wrapped_key_iv(KeySlotId::kEvenKeySlot);
  EXPECT_EQ(std::string(result.begin(), result.end()), kWrappedKeyIv2);
}

TEST(EcmParserV3Test, EntitlementRotationEnabledSuccess) {
  const uint32_t entitlement_period_index = 10;
  const uint32_t entitlement_rotation_window_left = 100;
  EcmPayload ecm_payload;
  ecm_payload.mutable_meta_data()->set_entitlement_period_index(
      entitlement_period_index);
  ecm_payload.mutable_meta_data()->set_entitlement_rotation_window_left(
      entitlement_rotation_window_left);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_TRUE(parser->is_entitlement_rotation_enabled());
  EXPECT_EQ(parser->entitlement_period_index(), entitlement_period_index);
  EXPECT_EQ(parser->entitlement_rotation_window_left(),
            entitlement_rotation_window_left);
}

TEST(EcmParserV3Test, EntitlementRotationDefaultDisabledSuccess) {
  EcmPayload ecm_payload;
  // Put something in the payload just to make the ECM valid.
  ecm_payload.mutable_even_key_data()->set_entitlement_key_id(kEntitlementId);
  SignedEcmPayload signed_ecm_payload;
  signed_ecm_payload.set_serialized_payload(ecm_payload.SerializeAsString());
  std::vector<uint8_t> ecm = GenerateEcm(signed_ecm_payload);

  std::unique_ptr<EcmParserV3> parser = EcmParserV3::Create(ecm);

  ASSERT_TRUE(parser != nullptr);
  EXPECT_FALSE(parser->is_entitlement_rotation_enabled());
}

}  // namespace
}  // namespace wvcas