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

#include <memory>
#include <string>

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

#include "crypto_session.h"
#include "google/protobuf/text_format.h"
#include "key_session.h"
#include "license_protocol.pb.h"
#include "log.h"
#include "metrics_collections.h"
#include "platform.h"
#include "test_base.h"
#include "test_printers.h"
#include "wv_cdm_types.h"
#include "wv_metrics.pb.h"

using ::testing::AllOf;
using ::testing::Ge;
using ::testing::Le;

namespace wvcdm {
class CryptoSessionForTest : public TestCryptoSession, public WvCdmTestBase {
 public:
  CryptoSessionForTest() : TestCryptoSession(metrics_.GetCryptoMetrics()) {}

  void SetUp() override {}

  KeySession* key_session() { return key_session_.get(); }

 private:
  static metrics::SessionMetrics metrics_;
};

metrics::SessionMetrics CryptoSessionForTest::metrics_;

class CryptoSessionMetricsTest : public WvCdmTestBase {
 protected:
  uint32_t FindKeyboxSystemID() {
    if (CryptoSession::needs_keybox_provisioning()) {
      return NULL_SYSTEM_ID;
    }
    uint8_t key_data[256];
    size_t key_data_len = sizeof(key_data);
    const OEMCryptoResult sts =
        OEMCrypto_GetKeyData(key_data, &key_data_len, kLevelDefault);
    if (sts != OEMCrypto_SUCCESS) return NULL_SYSTEM_ID;
    const uint32_t* data = reinterpret_cast<uint32_t*>(key_data);
    return htonl(data[1]);
  }
};

TEST_F(CryptoSessionMetricsTest, OpenSessionValidMetrics) {
  metrics::CryptoMetrics crypto_metrics;
  std::unique_ptr<CryptoSession> session(
      CryptoSession::MakeCryptoSession(&crypto_metrics));
  session->Open(wvcdm::kLevelDefault);
  // Exercise a method that will touch a metric.
  bool supports_usage_table;
  ASSERT_TRUE(session->HasUsageTableSupport(&supports_usage_table));

  drm_metrics::WvCdmMetrics::CryptoMetrics metrics_proto;
  crypto_metrics.Serialize(&metrics_proto);

  // TODO(blueeyes): If we an convert to use full proto rather than proto
  // lite, convert these tests to use Message-based convenience functions.

  // Validate common metrics regardless of provisioning type.
  ASSERT_EQ(1, metrics_proto.oemcrypto_initialize_time_us().size());
  EXPECT_TRUE(metrics_proto.oemcrypto_initialize_time_us(0)
                  .attributes()
                  .has_oem_crypto_result());
  EXPECT_EQ(1u,
            metrics_proto.oemcrypto_initialize_time_us(0).operation_count());
  EXPECT_TRUE(metrics_proto.oemcrypto_initialize_time_us(0).has_mean());

  const CdmUsageSupportType usage_type =
      supports_usage_table ? kUsageEntrySupport : kNonSecureUsageSupport;
  EXPECT_EQ(usage_type,
            metrics_proto.oemcrypto_usage_table_support().int_value());

  // Validate metrics that differ based on provisioning type.
  CdmClientTokenType token_type = session->GetPreProvisionTokenType();

  if (token_type == kClientTokenKeybox) {
    EXPECT_EQ(OEMCrypto_Keybox,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else if (token_type == kClientTokenOemCert) {
    EXPECT_EQ(OEMCrypto_OEMCertificate,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else if (token_type == kClientTokenDrmCert) {
    // TODO(blueeyes): Add support for getting the system id from a
    // pre-installed DRM certificate..
  } else if (token_type == kClientTokenBootCertChain) {
    EXPECT_EQ(OEMCrypto_BootCertificateChain,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else if (token_type == kClientTokenDrmCertificateReprovisioning) {
    EXPECT_EQ(OEMCrypto_DrmReprovisioning,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else {
    FAIL() << "Unexpected token type: " << token_type;
  }
}

TEST_F(CryptoSessionMetricsTest, GetProvisioningTokenValidMetrics) {
  metrics::CryptoMetrics crypto_metrics;
  std::unique_ptr<CryptoSession> session(
      CryptoSession::MakeCryptoSession(&crypto_metrics));

  ASSERT_EQ(NO_ERROR, session->Open(wvcdm::kLevelDefault));

  CdmClientTokenType token_type = session->GetPreProvisionTokenType();
  LOGI("token_type: %d", token_type);

  // DRM Certificate provisioning method does not support a provisioning
  // token. Otherwise, we should be able to fetch the token.
  std::string token;
  std::string additional_token;
  if (token_type != kClientTokenDrmCert) {
    ASSERT_EQ(NO_ERROR,
              session->GetProvisioningToken(&token, &additional_token));
  }

  drm_metrics::WvCdmMetrics::CryptoMetrics metrics_proto;
  crypto_metrics.Serialize(&metrics_proto);

  if (token_type == kClientTokenKeybox) {
    EXPECT_EQ(OEMCrypto_Keybox,
              metrics_proto.oemcrypto_provisioning_method().int_value());
    ASSERT_GE(metrics_proto.crypto_session_get_token().size(), 1);
    EXPECT_GE(metrics_proto.crypto_session_get_token(0).count(), 1);
  } else if (token_type == kClientTokenOemCert) {
    EXPECT_EQ(OEMCrypto_OEMCertificate,
              metrics_proto.oemcrypto_provisioning_method().int_value());
    ASSERT_GE(metrics_proto.oemcrypto_get_oem_public_certificate().size(), 1);
    EXPECT_THAT(metrics_proto.oemcrypto_get_oem_public_certificate(0).count(),
                AllOf(Ge(1), Le(2)));
  } else if (token_type == kClientTokenBootCertChain) {
    EXPECT_EQ(OEMCrypto_BootCertificateChain,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else if (token_type == kClientTokenDrmCertificateReprovisioning) {
    EXPECT_EQ(OEMCrypto_DrmReprovisioning,
              metrics_proto.oemcrypto_provisioning_method().int_value());
  } else {
    ASSERT_EQ(0, metrics_proto.crypto_session_get_token().size());
  }
}

}  // namespace wvcdm