//
// Copyright 2013 Google Inc. All Rights Reserved.
//

#include <stdio.h>
#include <string>

#include "gtest/gtest.h"
#include "media/stagefright/foundation/ABase.h"
#include "media/stagefright/foundation/AString.h"
#include "MockCDM.h"
#include "wv_cdm_constants.h"
#include "WVDrmPlugin.h"

using namespace android;
using namespace std;
using namespace testing;
using namespace wvcdm;
using namespace wvdrm;

class WVDrmPluginTest : public Test {
 protected:
  static const uint32_t kSessionIdSize = 16;
  Vector<uint8_t> sessionId;
  CdmSessionId cdmSessionId;

  virtual void SetUp() {
    uint8_t sessionIdRaw[kSessionIdSize];
    FILE* fp = fopen("/dev/urandom", "r");
    fread(sessionIdRaw, sizeof(uint8_t), kSessionIdSize, fp);
    fclose(fp);

    sessionId.appendArray(sessionIdRaw, kSessionIdSize);
    cdmSessionId.assign(sessionId.begin(), sessionId.end());

    // Set default CdmResponseType value for gMock
    DefaultValue<CdmResponseType>::Set(wvcdm::NO_ERROR);
  }
};

TEST_F(WVDrmPluginTest, OpensSessions) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  uint8_t sessionIdRaw[kSessionIdSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(sessionIdRaw, sizeof(uint8_t), kSessionIdSize, fp);
  fclose(fp);

  CdmSessionId cdmSessionId(sessionIdRaw, sessionIdRaw + kSessionIdSize);

  EXPECT_CALL(cdm, OpenSession(StrEq("com.widevine"), _))
      .WillOnce(DoAll(SetArgPointee<1>(cdmSessionId),
                      Return(wvcdm::NO_ERROR)));

  status_t res = plugin.openSession(sessionId);

  ASSERT_EQ(OK, res);
  EXPECT_THAT(sessionId, ElementsAreArray(sessionIdRaw, kSessionIdSize));
}

TEST_F(WVDrmPluginTest, ClosesSessions) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  EXPECT_CALL(cdm, CloseSession(cdmSessionId))
      .Times(1);

  status_t res = plugin.closeSession(sessionId);

  ASSERT_EQ(OK, res);
}

TEST_F(WVDrmPluginTest, GeneratesKeyRequests) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kInitDataSize = 128;
  uint8_t initDataRaw[kInitDataSize];
  static const uint32_t kRequestSize = 256;
  uint8_t requestRaw[kRequestSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(initDataRaw, sizeof(uint8_t), kInitDataSize, fp);
  fread(requestRaw, sizeof(uint8_t), kRequestSize, fp);
  fclose(fp);

  Vector<uint8_t> initData;
  initData.appendArray(initDataRaw, kInitDataSize);

  CdmKeyMessage cdmRequest(requestRaw, requestRaw + kRequestSize);
  Vector<uint8_t> request;

  KeyedVector<String8, String8> parameters;
  CdmNameValueMap cdmParameters;

  parameters.add(String8("paddingScheme"), String8("PKCS7"));
  cdmParameters["paddingScheme"] = "PKCS7";
  parameters.add(String8("favoriteParticle"), String8("tetraquark"));
  cdmParameters["favoriteParticle"] = "tetraquark";
  parameters.add(String8("answer"), String8("42"));
  cdmParameters["answer"] = "42";

  String8 defaultUrl;

  {
    InSequence calls;

    EXPECT_CALL(cdm, GenerateKeyRequest(cdmSessionId,
                                        ElementsAreArray(initDataRaw,
                                                         kInitDataSize),
                                        kLicenseTypeOffline, cdmParameters, _))
        .WillOnce(DoAll(SetArgPointee<4>(cdmRequest),
                        Return(wvcdm::NO_ERROR)));

    EXPECT_CALL(cdm, GenerateKeyRequest(cdmSessionId,
                                        ElementsAreArray(initDataRaw,
                                                         kInitDataSize),
                                        kLicenseTypeStreaming, cdmParameters,
                                        _))
        .WillOnce(DoAll(SetArgPointee<4>(cdmRequest),
                        Return(wvcdm::NO_ERROR)));
  }

  status_t res = plugin.getLicenseRequest(sessionId, initData,
                                          String8("video/h264"),
                                          DrmPlugin::kLicenseType_Offline,
                                          parameters, request, defaultUrl);

  ASSERT_EQ(OK, res);
  EXPECT_THAT(request, ElementsAreArray(requestRaw, kRequestSize));
  EXPECT_TRUE(defaultUrl.isEmpty());

  res = plugin.getLicenseRequest(sessionId, initData, String8("video/h264"),
                                 DrmPlugin::kLicenseType_Streaming, parameters,
                                 request, defaultUrl);

  ASSERT_EQ(OK, res);
  EXPECT_THAT(request, ElementsAreArray(requestRaw, kRequestSize));
  EXPECT_TRUE(defaultUrl.isEmpty());
}

TEST_F(WVDrmPluginTest, AddsKeys) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kResponseSize = 256;
  uint8_t responseRaw[kResponseSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(responseRaw, sizeof(uint8_t), kResponseSize, fp);
  fclose(fp);

  Vector<uint8_t> response;
  response.appendArray(responseRaw, kResponseSize);

  EXPECT_CALL(cdm, AddKey(cdmSessionId, ElementsAreArray(responseRaw,
                                                         kResponseSize)))
      .Times(1);

  status_t res = plugin.provideLicenseResponse(sessionId, response);

  ASSERT_EQ(OK, res);
}

TEST_F(WVDrmPluginTest, CancelsKeyRequests) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  EXPECT_CALL(cdm, CancelKeyRequest(cdmSessionId))
      .Times(1);

  status_t res = plugin.removeLicense(sessionId);

  ASSERT_EQ(OK, res);
}

TEST_F(WVDrmPluginTest, QueriesKeyStatus) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  KeyedVector<String8, String8> expectedLicenseStatus;
  CdmNameValueMap cdmLicenseStatus;

  expectedLicenseStatus.add(String8("areTheKeysAllRight"), String8("yes"));
  cdmLicenseStatus["areTheKeysAllRight"] = "yes";
  expectedLicenseStatus.add(String8("isGMockAwesome"), String8("ohhhhhhYeah"));
  cdmLicenseStatus["isGMockAwesome"] = "ohhhhhhYeah";
  expectedLicenseStatus.add(String8("answer"), String8("42"));
  cdmLicenseStatus["answer"] = "42";

  EXPECT_CALL(cdm, QueryKeyStatus(cdmSessionId, _))
      .WillOnce(DoAll(SetArgPointee<1>(cdmLicenseStatus),
                Return(wvcdm::NO_ERROR)));

  KeyedVector<String8, String8> licenseStatus;

  status_t res = plugin.queryLicenseStatus(sessionId, licenseStatus);

  ASSERT_EQ(OK, res);

  ASSERT_EQ(expectedLicenseStatus.size(), licenseStatus.size());
  for (size_t i = 0; i < expectedLicenseStatus.size(); ++i) {
    const String8& key = expectedLicenseStatus.keyAt(i);
    EXPECT_NE(android::NAME_NOT_FOUND, licenseStatus.indexOfKey(key));
    EXPECT_EQ(expectedLicenseStatus.valueFor(key), licenseStatus.valueFor(key));
  }
}

TEST_F(WVDrmPluginTest, GetsProvisioningRequests) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kRequestSize = 256;
  uint8_t requestRaw[kRequestSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(requestRaw, sizeof(uint8_t), kRequestSize, fp);
  fclose(fp);

  CdmProvisioningRequest cdmRequest(requestRaw, requestRaw + kRequestSize);

  static const char* kDefaultUrl = "http://google.com/";

  EXPECT_CALL(cdm, GetProvisioningRequest(_, _))
      .WillOnce(DoAll(SetArgPointee<0>(cdmRequest),
                      SetArgPointee<1>(kDefaultUrl),
                      Return(wvcdm::NO_ERROR)));

  Vector<uint8_t> request;
  String8 defaultUrl;

  status_t res = plugin.getProvisionRequest(request, defaultUrl);

  ASSERT_EQ(OK, res);
  EXPECT_THAT(request, ElementsAreArray(requestRaw, kRequestSize));
  EXPECT_EQ(String8(kDefaultUrl), defaultUrl);
}

TEST_F(WVDrmPluginTest, HandlesProvisioningResponses) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kResponseSize = 512;
  uint8_t responseRaw[kResponseSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(responseRaw, sizeof(uint8_t), kResponseSize, fp);
  fclose(fp);

  Vector<uint8_t> response;
  response.appendArray(responseRaw, kResponseSize);

  EXPECT_CALL(cdm, HandleProvisioningResponse(ElementsAreArray(responseRaw,
                                                               kResponseSize)))
      .Times(1);

  status_t res = plugin.provideProvisionResponse(response);

  ASSERT_EQ(OK, res);
}

TEST_F(WVDrmPluginTest, GetsSecureStops) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kStopSize = 53;
  static const uint32_t kStopCount = 7;
  uint8_t stopsRaw[kStopCount][kStopSize];
  FILE* fp = fopen("/dev/urandom", "r");
  for (uint32_t i = 0; i < kStopCount; ++i) {
    fread(stopsRaw[i], sizeof(uint8_t), kStopSize, fp);
  }
  fclose(fp);

  CdmSecureStops cdmStops;
  for (uint32_t i = 0; i < kStopCount; ++i) {
    cdmStops.push_back(string(stopsRaw[i], stopsRaw[i] + kStopSize));
  }

  EXPECT_CALL(cdm, GetSecureStops(_))
      .WillOnce(DoAll(SetArgPointee<0>(cdmStops),
                      Return(wvcdm::NO_ERROR)));

  List<Vector<uint8_t> > stops;

  status_t res = plugin.getSecureStops(stops);

  ASSERT_EQ(OK, res);

  List<Vector<uint8_t> >::iterator iter = stops.begin();
  uint32_t rawIter = 0;
  while (rawIter < kStopCount && iter != stops.end()) {
    EXPECT_THAT(*iter, ElementsAreArray(stopsRaw[rawIter], kStopSize));

    ++iter;
    ++rawIter;
  }
  // Assert that both lists are the same length
  EXPECT_EQ(kStopCount, rawIter);
  EXPECT_EQ(stops.end(), iter);
}

TEST_F(WVDrmPluginTest, ReleasesSecureStops) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kMessageSize = 128;
  uint8_t messageRaw[kMessageSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(messageRaw, sizeof(uint8_t), kMessageSize, fp);
  fclose(fp);

  Vector<uint8_t> message;
  message.appendArray(messageRaw, kMessageSize);

  EXPECT_CALL(cdm, ReleaseSecureStops(ElementsAreArray(messageRaw,
                                                       kMessageSize)))
      .Times(1);

  status_t res = plugin.releaseSecureStops(message);

  ASSERT_EQ(OK, res);
}

TEST_F(WVDrmPluginTest, DoesNotGetStringProperties) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  String8 result;

  status_t res = plugin.getPropertyString(String8("property"), result);

  ASSERT_NE(OK, res);
  EXPECT_TRUE(result.isEmpty());
}

TEST_F(WVDrmPluginTest, DoesNotGetByteProperties) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  Vector<uint8_t> result;

  status_t res = plugin.getPropertyByteArray(String8("property"), result);

  ASSERT_NE(OK, res);
  EXPECT_TRUE(result.isEmpty());
}

TEST_F(WVDrmPluginTest, DoesNotSetStringProperties) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  status_t res = plugin.setPropertyString(String8("property"),
                                          String8("ignored"));

  ASSERT_NE(OK, res);
}

TEST_F(WVDrmPluginTest, DoesNotSetByteProperties) {
  MockCDM cdm;
  WVDrmPlugin plugin(&cdm);

  static const uint32_t kValueSize = 32;
  uint8_t valueRaw[kValueSize];
  FILE* fp = fopen("/dev/urandom", "r");
  fread(valueRaw, sizeof(uint8_t), kValueSize, fp);
  fclose(fp);

  Vector<uint8_t> value;
  value.appendArray(valueRaw, kValueSize);

  status_t res = plugin.setPropertyByteArray(String8("property"), value);

  ASSERT_NE(OK, res);
}