#include <gtest/gtest.h>
#include <activation.h>
#include <server_emulator.h>
#include <stdio.h>
#include <string.h>

#define kEncryptedMaxSize (256)
#define kTagSize (16)
#define kIvSize (12)

TEST(ActivationGenerateSessionCertificate, Negative) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);

  ASSERT_NE(kActivationSuccess,
            ActivationGenerateSessionCertificate(NULL, NULL, NULL, NULL));
  ASSERT_NE(kActivationSuccess,
            ActivationGenerateSessionCertificate(NULL, &drk_cert_size,
                                                 NULL, &rsa_cert_size));
  ASSERT_NE(kActivationSuccess,
            ActivationGenerateSessionCertificate(drk_cert, NULL, rsa_cert, NULL));
}

TEST(ActivationGenerateSessionCertificate, Simple) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);

  ASSERT_EQ(kActivationSuccess,
            ActivationGenerateSessionCertificate(drk_cert, &drk_cert_size,
                                                 rsa_cert, &rsa_cert_size));
}

TEST(ActivationStoreServerKey, Simple) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);
  uint8_t plain[] = { 0xfe, 0xff, 0xe9, 0x92, 0x86, 0x65, 0x73, 0x1c,
         0x6d, 0x6a, 0x8f, 0x94, 0x67, 0x30, 0x83, 0x08 };
  size_t plain_size = sizeof(plain);
  uint8_t encrypted[kEncryptedMaxSize];
  size_t encrypted_size = sizeof(encrypted);

  // TODO: Add a test that checks RSA certificate validity against DRK certificate.
  ASSERT_EQ(kActivationSuccess,
            ActivationGenerateSessionCertificate(drk_cert, &drk_cert_size,
                                                 rsa_cert, &rsa_cert_size));
  ASSERT_EQ(kServerEmulatorSuccess,
            ServerEmulatorEncrypt(rsa_cert, rsa_cert_size, plain, plain_size, encrypted,
                                  &encrypted_size));
  ASSERT_EQ(kActivationSuccess,
            ActivationStoreServerKey(encrypted, encrypted_size));
}

TEST(ActivationStoreServerKey, Negative) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);
  uint8_t plain[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
  size_t plain_size = sizeof(plain);
  uint8_t encrypted[kEncryptedMaxSize];
  size_t encrypted_size = sizeof(encrypted);

  ASSERT_EQ(kActivationSuccess,
            ActivationGenerateSessionCertificate(drk_cert, &drk_cert_size,
                                                 rsa_cert, &rsa_cert_size));
  ASSERT_EQ(kServerEmulatorSuccess,
            ServerEmulatorEncrypt(rsa_cert, rsa_cert_size, plain, plain_size, encrypted,
                                  &encrypted_size));
  // Corrupt encrypted data.
  encrypted[0] = 0;
  ASSERT_EQ(kActivationErrorTee,
            ActivationStoreServerKey(encrypted, encrypted_size));
}

TEST(ActivationGenerateCredentials, Negative) {
  uint8_t data[60] = {
    0xd9, 0x31, 0x32, 0x25, 0xf8, 0x84, 0x06, 0xe5,
    0xa5, 0x59, 0x09, 0xc5, 0xaf, 0xf5, 0x26, 0x9a,
    0x86, 0xa7, 0xa9, 0x53, 0x15, 0x34, 0xf7, 0xda,
    0x2e, 0x4c, 0x30, 0x3d, 0x8a, 0x31, 0x8a, 0x72,
    0x1c, 0x3c, 0x0c, 0x95, 0x95, 0x68, 0x09, 0x53,
    0x2f, 0xcf, 0x0e, 0x24, 0x49, 0xa6, 0xb5, 0x25,
    0xb1, 0x6a, 0xed, 0xf5, 0xaa, 0x0d, 0xe6, 0x57,
    0xba, 0x63, 0x7b, 0x39 };
  const size_t data_size = sizeof(data);
  size_t encrypted_data_size = data_size;
  uint8_t encrypted_data[encrypted_data_size];
  uint8_t iv[kIvSize];
  size_t iv_size = sizeof(iv);
  uint8_t tag[kTagSize];
  size_t tag_size = sizeof(tag);

  ASSERT_EQ(kActivationErrorBadParameters,
            ActivationGenerateCredentials(NULL, data_size, encrypted_data,
                                          &encrypted_data_size, iv, &iv_size,
                                          tag, &tag_size));
  ASSERT_EQ(kActivationErrorBadParameters,
            ActivationGenerateCredentials(data, 0, encrypted_data,
                                          &encrypted_data_size, iv, &iv_size,
                                          tag, &tag_size));
  ASSERT_EQ(kActivationErrorBadParameters,
            ActivationGenerateCredentials(data, data_size, NULL,
                                          &encrypted_data_size, iv, &iv_size,
                                          tag, &tag_size));
  ASSERT_EQ(kActivationErrorBadParameters,
            ActivationGenerateCredentials(data, data_size, encrypted_data, NULL,
                                          iv, &iv_size, tag, &tag_size));
}

TEST(ActivationGenerateCredentials, Simple) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);
  uint8_t plain[] = { 0xfe, 0xff, 0xe9, 0x92, 0x86, 0x65, 0x73, 0x1c,
         0x6d, 0x6a, 0x8f, 0x94, 0x67, 0x30, 0x83, 0x08 };
  const size_t plain_size = sizeof(plain) / sizeof(plain[0]);
  uint8_t encrypted[kEncryptedMaxSize];
  size_t encrypted_size = sizeof(encrypted);

  ASSERT_EQ(kActivationSuccess,
            ActivationGenerateSessionCertificate(drk_cert, &drk_cert_size,
                                                 rsa_cert, &rsa_cert_size));
  ASSERT_EQ(kServerEmulatorSuccess,
            ServerEmulatorEncrypt(rsa_cert, rsa_cert_size, plain, plain_size, encrypted,
                                  &encrypted_size));
  ASSERT_EQ(kActivationSuccess,
            ActivationStoreServerKey(encrypted, encrypted_size));

  uint8_t data[60] = {
    0xd9, 0x31, 0x32, 0x25, 0xf8, 0x84, 0x06, 0xe5,
    0xa5, 0x59, 0x09, 0xc5, 0xaf, 0xf5, 0x26, 0x9a,
    0x86, 0xa7, 0xa9, 0x53, 0x15, 0x34, 0xf7, 0xda,
    0x2e, 0x4c, 0x30, 0x3d, 0x8a, 0x31, 0x8a, 0x72,
    0x1c, 0x3c, 0x0c, 0x95, 0x95, 0x68, 0x09, 0x53,
    0x2f, 0xcf, 0x0e, 0x24, 0x49, 0xa6, 0xb5, 0x25,
    0xb1, 0x6a, 0xed, 0xf5, 0xaa, 0x0d, 0xe6, 0x57,
    0xba, 0x63, 0x7b, 0x39
  };
  const size_t data_size = sizeof(data);
  size_t encrypted_data_size = data_size;
  size_t decrypted_data_size = data_size;
  uint8_t encrypted_data[encrypted_data_size];
  uint8_t decrypted_data[decrypted_data_size];
  uint8_t iv[kIvSize];
  size_t iv_size = sizeof(iv);
  uint8_t tag[kTagSize];
  size_t tag_size = sizeof(tag);

  ASSERT_EQ(kActivationSuccess,
            ActivationGenerateCredentials(data, data_size, encrypted_data,
                                          &encrypted_data_size, iv, &iv_size, 
                                          tag, &tag_size));
  ASSERT_EQ(data_size,
            encrypted_data_size);

  ASSERT_EQ(kServerEmulatorSuccess,
            ServerEmulatorAesDecrypt(128, encrypted_data, encrypted_data_size,
                                    plain, sizeof(plain), iv, iv_size,
                                    tag, tag_size, decrypted_data,
                                    decrypted_data_size));

  ASSERT_EQ(0,
            memcmp(data, decrypted_data, decrypted_data_size));
}