#include <gtest/gtest.h>

extern "C" {
  #include "crypto.h"
}

static const uint8_t kSha256ZeroVector[] =
  {0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14,
   0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
   0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c,
   0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55};

static const HashType kIncorrectHashType = (HashType)250;


class CryptoTest : public ::testing::Test {
protected:
  const uint8_t data_[kSha256Size] = {0};
  uint8_t hash_[kSha256Size] = {0};

  virtual void SetUp() {
  }

  virtual void TearDown() {
  }
};

TEST_F(CryptoTest, GetHashSize_Sha1) {
  size_t size = CryptoGetHashSize(kSha1);

  EXPECT_EQ(kSha1Size, size);
}

TEST_F(CryptoTest, GetHashSize_Sha224) {
  size_t size = CryptoGetHashSize(kSha224);

  EXPECT_EQ(kSha224Size, size);
}

TEST_F(CryptoTest, GetHashSize_Sha256) {
  size_t size = CryptoGetHashSize(kSha256);

  EXPECT_EQ(kSha256Size, size);
}

TEST_F(CryptoTest, GetHashSize_Sha384) {
  size_t size = CryptoGetHashSize(kSha384);

  EXPECT_EQ(kSha384Size, size);
}

TEST_F(CryptoTest, GetHashSize_Sha512) {
  size_t size = CryptoGetHashSize(kSha512);

  EXPECT_EQ(kSha512Size, size);
}

TEST_F(CryptoTest, GetHashSize_Invalid) {
  size_t size = CryptoGetHashSize(kIncorrectHashType);

  EXPECT_EQ(size, 0);
}

TEST_F(CryptoTest, GetRsaSignatureSize_Rsa1024) {
  size_t size = CryptoGetRsaSignatureSize(kRSA1024);

  EXPECT_GT(size, 0);
}

TEST_F(CryptoTest, GetRsaSignatureSize_Rsa2048) {
  size_t size = CryptoGetRsaSignatureSize(kRSA2048);

  EXPECT_GT(size, 0);
}

TEST_F(CryptoTest, GetRsaSignatureSize_Rsa3072) {
  size_t size = CryptoGetRsaSignatureSize(kRSA3072);

  EXPECT_GT(size, 0);
}

TEST_F(CryptoTest, GetRsaSignatureSize_Invalid) {
  size_t size = CryptoGetRsaSignatureSize((RsaKeyType)250);

  EXPECT_EQ(size, 0);
}

TEST_F(CryptoTest, Sha256_NullData) {
  PaTzResult result = CryptoSha256(NULL, 10, hash_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoTest, Sha256_NullDataAndZeroSize) {
  PaTzResult result = CryptoSha256(NULL, 0, hash_);

  bool is_hash_correct = memcmp(hash_, kSha256ZeroVector, sizeof(kSha256ZeroVector))
      ? false : true;

  EXPECT_EQ(result, PA_TZ_SUCCESS);
  EXPECT_TRUE(is_hash_correct);
}

TEST_F(CryptoTest, Sha256_ZeroSize) {
  PaTzResult result = CryptoSha256(data_, 0, hash_);

  bool is_hash_correct = memcmp(hash_, kSha256ZeroVector, sizeof(kSha256ZeroVector))
        ? false : true;

  EXPECT_EQ(result, PA_TZ_SUCCESS);
  EXPECT_TRUE(is_hash_correct);
}

TEST_F(CryptoTest, Sha256_NullHash) {
  PaTzResult result = CryptoSha256(data_, sizeof(data_), NULL);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}


class CryptoHmacSignatureGenerateTest : public ::testing::Test {
protected:
  const uint8_t data_[32] = {0};
  uint8_t signature_[256] = {0};
  size_t signature_size_ = sizeof(signature_);

  virtual void SetUp() {
    signature_size_ = sizeof(signature_);
  }

  virtual void TearDown() {
  }
};

TEST_F(CryptoHmacSignatureGenerateTest, NullData) {
  PaTzResult result = CryptoHmacSignatureGenerate(NULL, sizeof(data_), kSha256, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, NullDataZeroSize) {
  PaTzResult result = CryptoHmacSignatureGenerate(NULL, 0, kSha256, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, IncorrectHashType) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kIncorrectHashType, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, NullSignature) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, NULL, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, NullSignatureSize) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, NULL);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, Sha1Hash) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha1, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, SmallSignatureSize) {
  signature_size_ = CryptoGetHashSize(kSha256) - 1;

  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureGenerateTest, ValidInput) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, &signature_size_);

  EXPECT_EQ(result, PA_TZ_SUCCESS);
  EXPECT_EQ(signature_size_, CryptoGetHashSize(kSha256));
}


class CryptoHmacSignatureVerificationTest : public ::testing::Test {
protected:
  const uint8_t data_[32] = {0};
  uint8_t signature_[256] = {0};
  size_t signature_size_ = sizeof(signature_);

  virtual void SetUp() {
    signature_size_ = sizeof(signature_);
  }

  virtual void TearDown() {
  }
};

TEST_F(CryptoHmacSignatureVerificationTest, NullData) {
  PaTzResult result = CryptoHmacSignatureVerification(NULL, sizeof(data_), kSha256, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, NullDataZeroSize) {
  PaTzResult result = CryptoHmacSignatureVerification(NULL, 0, kSha256, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, NullSignature) {
  PaTzResult result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha256, NULL, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, IncorrectHashType) {
  PaTzResult result = CryptoHmacSignatureVerification(data_, sizeof(data_), kIncorrectHashType, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, UnsupportedHashType) {
  PaTzResult result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha1, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, IncorrectSignature) {
  PaTzResult result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha256, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, IncorrectSignatureSize) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, &signature_size_);
  EXPECT_EQ(result, PA_TZ_SUCCESS);

  result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha256, signature_, signature_size_ - 1);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, ModifiedSignature) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, &signature_size_);
  EXPECT_EQ(result, PA_TZ_SUCCESS);

  // Modify signature
  signature_[0] += 1;

  result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha256, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(CryptoHmacSignatureVerificationTest, ValidInput) {
  PaTzResult result = CryptoHmacSignatureGenerate(data_, sizeof(data_), kSha256, signature_, &signature_size_);
  EXPECT_EQ(result, PA_TZ_SUCCESS);

  result = CryptoHmacSignatureVerification(data_, sizeof(data_), kSha256, signature_, signature_size_);

  EXPECT_EQ(result, PA_TZ_SUCCESS);
}


class RsaSignatureVerificationTest : public ::testing::Test {
protected:
  const uint8_t data_[32] = {0};
  uint8_t signature_[256] = {0};
  const size_t signature_size_ = sizeof(signature_);
  const RsaPublicKey rsa_key_ = {0};

  virtual void SetUp() {
  }

  virtual void TearDown() {
  }
};

TEST_F(RsaSignatureVerificationTest, NullData) {
  PaTzResult result = RsaSignatureVerification(NULL, sizeof(data_), signature_, signature_size_, &rsa_key_, kSha256);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(RsaSignatureVerificationTest, NullSignature) {
  PaTzResult result = RsaSignatureVerification(data_, sizeof(data_), NULL, signature_size_, &rsa_key_, kSha256);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(RsaSignatureVerificationTest, NullKey) {
  PaTzResult result = RsaSignatureVerification(data_, sizeof(data_), signature_, signature_size_, NULL, kSha256);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(RsaSignatureVerificationTest, InvalidHashType) {
  PaTzResult result = RsaSignatureVerification(data_, sizeof(data_), signature_, signature_size_, &rsa_key_, kIncorrectHashType);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(RsaSignatureVerificationTest, UnsupportedHashType) {
  PaTzResult result = RsaSignatureVerification(data_, sizeof(data_), signature_, signature_size_, &rsa_key_, kSha1);

  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(RsaSignatureVerificationTest, IncorrectSignature) {
  PaTzResult result = RsaSignatureVerification(data_, sizeof(data_), signature_, signature_size_, &rsa_key_, kSha256);

  EXPECT_EQ(result, PA_TZ_SINGATURE_VALIDATION_FAILURE);
}