#include "safebrowsing.pb.h"
#include "gtest/gtest.h"
#include "nsUrlClassifierUtils.h"
#include "mozilla/Base64.h"

using namespace mozilla;
using namespace mozilla::safebrowsing;

namespace {

// |Base64EncodedStringArray| and |MakeBase64EncodedStringArray|
// works together to make us able to do things "literally" and easily.

// Given a nsCString array, construct an object which can be implicitly
// casted to |const char**|, where all owning c-style strings have been
// base64 encoded. The memory life cycle of what the "cast operator"
// returns is just as the object itself.
class Base64EncodedStringArray
{
public:
  Base64EncodedStringArray(nsCString aArray[], size_t N);
  operator const char** () const { return (const char**)&mArray[0]; }

private:
  // Since we can't guarantee the layout of nsCString (can we?),
  // an additional nsTArray<nsCString> is required to manage the
  // allocated string.
  nsTArray<const char*> mArray;
  nsTArray<nsCString> mStringStorage;
};

// Simply used to infer the fixed-array size automatically.
template<size_t N>
Base64EncodedStringArray
MakeBase64EncodedStringArray(nsCString (&aArray)[N])
{
  return Base64EncodedStringArray(aArray, N);
}

} // end of unnamed namespace.


TEST(UrlClassifierFindFullHash, Request)
{
  nsCOMPtr<nsIUrlClassifierUtils> urlUtil =
    do_GetService("@mozilla.org/url-classifier/utils;1");

  const char* listNames[] = { "test-phish-proto", "test-unwanted-proto" };

  nsCString listStates[] = { nsCString("sta\x00te1", 7),
                             nsCString("sta\x00te2", 7) };

  nsCString prefixes[] = { nsCString("\x00\x00\x00\x01", 4),
                           nsCString("\x00\x00\x00\x00\x01", 5),
                           nsCString("\x00\xFF\x00\x01", 4),
                           nsCString("\x00\xFF\x00\x01\x11\x23\xAA\xBC", 8),
                           nsCString("\x00\x00\x00\x01\x00\x01\x98", 7) };

  nsCString requestBase64;
  nsresult rv;
  rv = urlUtil->MakeFindFullHashRequestV4(listNames,
                                          MakeBase64EncodedStringArray(listStates),
                                          MakeBase64EncodedStringArray(prefixes),
                                          ArrayLength(listNames),
                                          ArrayLength(prefixes),
                                          requestBase64);
  ASSERT_TRUE(NS_SUCCEEDED(rv));

  // Base64 URL decode first.
  FallibleTArray<uint8_t> requestBinary;
  rv = Base64URLDecode(requestBase64,
                       Base64URLDecodePaddingPolicy::Require,
                       requestBinary);
  ASSERT_TRUE(NS_SUCCEEDED(rv));

  // Parse the FindFullHash binary and compare with the expected values.
  FindFullHashesRequest r;
  ASSERT_TRUE(r.ParseFromArray(&requestBinary[0], requestBinary.Length()));

  // Compare client states.
  ASSERT_EQ(r.client_states_size(), (int)ArrayLength(listStates));
  for(int i = 0; i < r.client_states_size(); i++) {
    auto s = r.client_states(i);
    ASSERT_TRUE(listStates[i].Equals(nsCString(s.c_str(), s.size())));
  }

  auto threatInfo = r.threat_info();

  // Compare threat types.
  ASSERT_EQ(threatInfo.threat_types_size(), (int)ArrayLength(listStates));
  for (int i = 0; i < threatInfo.threat_types_size(); i++) {
    uint32_t expectedThreatType;
    rv = urlUtil->ConvertListNameToThreatType(nsCString(listNames[i]),
                                              &expectedThreatType);
    ASSERT_TRUE(NS_SUCCEEDED(rv));
    ASSERT_EQ(threatInfo.threat_types(i), expectedThreatType);
  }

  // Compare prefixes.
  ASSERT_EQ(threatInfo.threat_entries_size(), (int)ArrayLength(prefixes));
  for (int i = 0; i < threatInfo.threat_entries_size(); i++) {
    auto p = threatInfo.threat_entries(i).hash();
    ASSERT_TRUE(prefixes[i].Equals(nsCString(p.c_str(), p.size())));
  }
}

/////////////////////////////////////////////////////////////
// Following is to test parsing the gethash response.

namespace {

// safebrowsing::Duration manipulation.
struct MyDuration {
  uint32_t mSecs;
  uint32_t mNanos;
};
void PopulateDuration(Duration& aDest, const MyDuration& aSrc)
{
  aDest.set_seconds(aSrc.mSecs);
  aDest.set_nanos(aSrc.mNanos);
}

// The expected match data.
static MyDuration EXPECTED_MIN_WAIT_DURATION = { 12, 10 };
static MyDuration EXPECTED_NEG_CACHE_DURATION = { 120, 9 };
static const struct ExpectedMatch {
  nsCString mCompleteHash;
  ThreatType mThreatType;
  MyDuration mPerHashCacheDuration;
} EXPECTED_MATCH[] = {
  { nsCString("01234567890123456789012345678901"), SOCIAL_ENGINEERING_PUBLIC, { 8, 500 } },
  { nsCString("12345678901234567890123456789012"), SOCIAL_ENGINEERING_PUBLIC, { 7, 100} },
  { nsCString("23456789012345678901234567890123"), SOCIAL_ENGINEERING_PUBLIC, { 1, 20 } },
};

class MyParseCallback final :
  public nsIUrlClassifierParseFindFullHashCallback {
public:
  NS_DECL_ISUPPORTS

  explicit MyParseCallback(uint32_t& aCallbackCount)
    : mCallbackCount(aCallbackCount)
  {
  }

  NS_IMETHOD
  OnCompleteHashFound(const nsACString& aCompleteHash,
                      const nsACString& aTableNames,
                      uint32_t aPerHashCacheDuration) override
  {
    Verify(aCompleteHash,
           aTableNames,
           aPerHashCacheDuration);

    return NS_OK;
  }

  NS_IMETHOD
  OnResponseParsed(uint32_t aMinWaitDuration,
                   uint32_t aNegCacheDuration) override
  {
    VerifyDuration(aMinWaitDuration / 1000, EXPECTED_MIN_WAIT_DURATION);
    VerifyDuration(aNegCacheDuration, EXPECTED_NEG_CACHE_DURATION);

    return NS_OK;
  }

private:
  void
  Verify(const nsACString& aCompleteHash,
         const nsACString& aTableNames,
         uint32_t aPerHashCacheDuration)
  {
    auto expected = EXPECTED_MATCH[mCallbackCount];

    ASSERT_TRUE(aCompleteHash.Equals(expected.mCompleteHash));

    // Verify aTableNames
    nsCOMPtr<nsIUrlClassifierUtils> urlUtil =
      do_GetService("@mozilla.org/url-classifier/utils;1");
    nsCString tableNames;
    nsresult rv = urlUtil->ConvertThreatTypeToListNames(expected.mThreatType, tableNames);
    ASSERT_TRUE(NS_SUCCEEDED(rv));
    ASSERT_TRUE(aTableNames.Equals(tableNames));

    VerifyDuration(aPerHashCacheDuration, expected.mPerHashCacheDuration);

    mCallbackCount++;
  }

  void
  VerifyDuration(uint32_t aToVerify, const MyDuration& aExpected)
  {
    ASSERT_TRUE(aToVerify == aExpected.mSecs);
  }

  ~MyParseCallback() {}

  uint32_t& mCallbackCount;
};

NS_IMPL_ISUPPORTS(MyParseCallback, nsIUrlClassifierParseFindFullHashCallback)

} // end of unnamed namespace.

TEST(UrlClassifierFindFullHash, ParseRequest)
{
  // Build response.
  FindFullHashesResponse r;

  // Init response-wise durations.
  auto minWaitDuration = r.mutable_minimum_wait_duration();
  PopulateDuration(*minWaitDuration, EXPECTED_MIN_WAIT_DURATION);
  auto negCacheDuration = r.mutable_negative_cache_duration();
  PopulateDuration(*negCacheDuration, EXPECTED_NEG_CACHE_DURATION);

  // Init matches.
  for (uint32_t i = 0; i < ArrayLength(EXPECTED_MATCH); i++) {
    auto expected = EXPECTED_MATCH[i];
    auto match = r.mutable_matches()->Add();
    match->set_threat_type(expected.mThreatType);
    match->mutable_threat()->set_hash(expected.mCompleteHash.BeginReading(),
                                      expected.mCompleteHash.Length());
    auto perHashCacheDuration = match->mutable_cache_duration();
    PopulateDuration(*perHashCacheDuration, expected.mPerHashCacheDuration);
  }
  std::string s;
  r.SerializeToString(&s);

  uint32_t callbackCount = 0;
  nsCOMPtr<nsIUrlClassifierParseFindFullHashCallback> callback
    = new MyParseCallback(callbackCount);

  nsCOMPtr<nsIUrlClassifierUtils> urlUtil =
    do_GetService("@mozilla.org/url-classifier/utils;1");
  nsresult rv = urlUtil->ParseFindFullHashResponseV4(nsCString(s.c_str(), s.size()),
                                                     callback);
  NS_ENSURE_SUCCESS_VOID(rv);

  ASSERT_EQ(callbackCount, ArrayLength(EXPECTED_MATCH));
}


/////////////////////////////////////////////////////////////
namespace {

Base64EncodedStringArray::Base64EncodedStringArray(nsCString aArray[],
                                                   size_t N)
{
  for (size_t i = 0; i < N; i++) {
    nsCString encoded;
    nsresult rv = Base64Encode(aArray[i], encoded);
    NS_ENSURE_SUCCESS_VOID(rv);
    mStringStorage.AppendElement(encoded);
    mArray.AppendElement(encoded.get());
  }
}

}