#include <gtest/gtest.h>

#include <stdint.h>

extern "C" {
  #include "mock.h"

  #include "access_control.h"
  #include "api.h"
  #include "authentication.h"
  #include "config.h"
  #include "provisioning.h"
  #include "task.h"
  #include "task_access.h"
  #include "PaFlagBits.h"
}

static const uint32_t kValidPid = 125;
static const uint32_t kValidHmacPid = 135;

static const uint32_t kUnauthenticatedPid = 196;

static const PaHandler kValidPidHandler = 
                  {kValidPid, kPaHandlerTidUnusedMark, 0};
static const PaHandler kValidHmacPidHandler =
                  {kValidHmacPid, kPaHandlerTidUnusedMark, 0};

static const char kValidAppName[] = "valid_app_name";
static const char kValidHmacAppName[] = "valid_hmac_app_name";

static const char kUnauthenticatedAppName[] = "unauthenticated_app_name";

static const PaHandler kValidAppNameHandler = {0, kPaHandlerTidUnusedMark, 0};

static const PaHandler kInvalidHandler =
                  {kValidPid, kPaHandlerTidUnusedMark + 1, 0};
static const PaHandler kUnauthenticatedPidHandler =
                  {kUnauthenticatedPid, kPaHandlerTidUnusedMark, 0};

static const uint8_t kValidPaId[kPaIdLength] = {0xFF, 0xAA};
static const uint64_t kPaFlags = (1 << PaFlagBits_bitAndroid) | (1 << PaFlagBits_bitThirdParty);

static PaCertificate_t kValidCertificate = {0};

static uint32_t g_is_five_enabled = 1;
static uint32_t g_access_control_return = 1;
static int32_t g_is_process_addr_readable = 1;
static int32_t g_is_process_addr_writable = 1;
static int g_platform_sys_map_trustlet_success = 1;
static PaTzResult g_task_access_get_bytes_return = PA_TZ_SUCCESS;
static PaTzResult g_task_access_put_bytes_return = PA_TZ_SUCCESS;

static const size_t kBufferSize = 8;
static const uint8_t kInputBuf[kBufferSize] = {};
static const ProcessAddress kUserAddress = (ProcessAddress)&kInputBuf;

extern "C" PaTzResult CheckIntegrityWeak(const TaskInfo *task) {
  if (task->integrity == kIntegrityMixed) {
      return PA_TZ_SUCCESS;
    } else {
      return PA_TZ_GENERAL_ERROR;
    }
}

extern "C" void FreeTaskInfo(TaskInfo *task) {

}

extern "C" PaTzResult PlatformSysMapTrustlet(const void *virt_trustlet,
                                             size_t size, MemoryAccessType type,
                                             void **virt_driver) {
  if (!g_platform_sys_map_trustlet_success) {
    return PA_TZ_GENERAL_ERROR;
  }

  *virt_driver = (void *)virt_trustlet;

  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult PlatformSysUnmapTrustlet(const void *virt_driver, size_t size) {
  return PA_TZ_SUCCESS;
}

extern "C" uint32_t AccessControlIsAllowedOperation(Operation operation) {
  return g_access_control_return;
}

extern "C" PaTzResult TaskFindByPid(uint32_t pid, TaskInfo *out_task) {

  if (pid == kValidPid || pid == kValidHmacPid) {
    if (pid == kValidHmacPid) {
      out_task->integrity = kIntegrityMixed;
    } else {
      out_task->integrity = kIntegrityPreload;
    }
    out_task->certificate = &kValidCertificate;
    out_task->certificate->paData.paId.buf = (uint8_t *)kValidPaId;
    out_task->certificate->paData.paId.size = kPaIdLength;
    out_task->certificate->paData.paFlags = kPaFlags;
    return PA_TZ_SUCCESS;
  }

  if (pid == kUnauthenticatedPid) {
    out_task->integrity = kIntegrityNone;
    return PA_TZ_SUCCESS;
  }

  return PA_TZ_GENERAL_ERROR;
}

extern "C" PaTzResult TaskFindByAppName(const char *app_name,
                                        const size_t app_name_size,
                                        TaskInfo *out_task) {
  if (!strncmp(app_name, kValidAppName, app_name_size) ||
      !strncmp(app_name, kValidHmacAppName, app_name_size)) {
    if (!strncmp(app_name, kValidHmacAppName, app_name_size)) {
      out_task->integrity = kIntegrityMixed;
    } else {
      out_task->integrity = kIntegrityPreload;
    }
    out_task->certificate = &kValidCertificate;
    out_task->certificate->paData.paId.buf = (uint8_t *)kValidPaId;
    out_task->certificate->paData.paId.size = kPaIdLength;
    out_task->certificate->paData.paFlags = kPaFlags;
    return PA_TZ_SUCCESS;
  }

  if (!strncmp(app_name, kUnauthenticatedAppName, app_name_size)) {
    out_task->integrity = kIntegrityNone;
    return PA_TZ_SUCCESS;
  }

  return PA_TZ_GENERAL_ERROR;
}

extern "C" PaTzResult ProcessAuthentication(const TaskInfo *task,
                           const char *process_names, const size_t process_names_size,
                           const PaTzMemoryRange *memory) {
  return ((task->integrity != kIntegrityNone) ? PA_TZ_SUCCESS : PA_TZ_AUTHENTICATION_FAILED);
}

extern "C" int32_t IsProcessAddressReadable(const TaskInfo *task, ProcessAddress address) {
  return g_is_process_addr_readable;
}

extern "C" int32_t IsProcessAddressWritable(const TaskInfo *task, ProcessAddress address) {
  return g_is_process_addr_writable;
}

extern "C" PaTzResult TaskAccessGetBytes(const TaskInfo *task,
                                         ProcessAddress address,
                                         size_t size, void *out) {
  return g_task_access_get_bytes_return;
}

extern "C" PaTzResult TaskAccessPutBytes(const TaskInfo *task, const void *in,
                                         size_t size, ProcessAddress address) {
  return g_task_access_put_bytes_return;
}

extern "C" PaTzResult CreateNewCertificate(
    const TaskInfo *issuer,
    ProcessAddress mapped_apk,
    const char *package_name,
    const uint8_t *rsa, size_t rsa_size,
    PaCertificate_t **new_certificate) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult CreateCopyOfCertificate(
    const TaskInfo *issuer,
    ProcessAddress mapped_apk,
    ProcessAddress apk_xattr,
    ProcessAddress mapped_file,
    PaCertificate_t **new_certificate) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskPutPaCertificate(
    const TaskInfo *task,
    const PaCertificate_t *certificate,
    ProcessAddress out_xattr, size_t out_xattr_size) {
  return PA_TZ_SUCCESS;
}

extern "C" void PaCertificateDestroy(PaCertificate_t *certificate) {
}

extern "C" uint32_t IsFiveEnabled() {
  return g_is_five_enabled;
}

class AuthenticateTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    InitMocks();
    g_access_control_return = 1;
    memset(&process_info, 0, sizeof(process_info));
  }

  virtual void TearDown() {
    DeinitMocks();
  }

  ProcessInfo process_info;
};

class TaskReadToTrustletTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    InitMocks();
    g_access_control_return = 1;
    g_is_process_addr_readable = 1;
    g_platform_sys_map_trustlet_success = 1;
    g_task_access_get_bytes_return = PA_TZ_SUCCESS;
  }

  virtual void TearDown() {
    DeinitMocks();
  }
};

class TaskWriteFromTrustletTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    InitMocks();
    g_access_control_return = 1;
    g_is_process_addr_writable = 1;
    g_platform_sys_map_trustlet_success = 1;
    g_task_access_put_bytes_return = PA_TZ_SUCCESS;
  }

  virtual void TearDown() {
    DeinitMocks();
  }
};

TEST_F(AuthenticateTest, CheckFlagsProcessInfo) {
  PaTzResult result = Authenticate(kValidPidHandler, NULL, 0, NULL, &process_info);

  EXPECT_EQ(PA_TZ_SUCCESS, result);
  EXPECT_EQ(process_info.is_android_app, 1);
  EXPECT_EQ(process_info.is_thirdparty_app, 1);
  EXPECT_EQ(process_info.is_weak_integrity, 0);
}

TEST_F(AuthenticateTest, CheckWeakIntegrityPidProcessInfo) {
  PaTzResult result = Authenticate(kValidHmacPidHandler, NULL, 0, NULL, &process_info);

  EXPECT_EQ(PA_TZ_SUCCESS, result);
  EXPECT_EQ(process_info.is_android_app, 1);
  EXPECT_EQ(process_info.is_thirdparty_app, 1);
  EXPECT_EQ(process_info.is_weak_integrity, 1);
}

TEST_F(AuthenticateTest, ValidPid) {
  PaTzResult result = Authenticate(kValidPidHandler, NULL, 0, NULL, NULL);

  EXPECT_EQ(PA_TZ_SUCCESS, result);
}

TEST_F(AuthenticateTest, CheckWeakIntegrityAppNameProcessInfo) {
  PaTzResult result = Authenticate(kValidAppNameHandler, kValidHmacAppName,
                                   sizeof(kValidHmacAppName) - 1, NULL, &process_info);

  EXPECT_EQ(PA_TZ_SUCCESS, result);
  EXPECT_EQ(process_info.is_android_app, 1);
  EXPECT_EQ(process_info.is_thirdparty_app, 1);
  EXPECT_EQ(process_info.is_weak_integrity, 1);
}

TEST_F(AuthenticateTest, ValidAppName) {
  PaTzResult result = Authenticate(kValidAppNameHandler, kValidAppName,
                                   sizeof(kValidAppName) - 1, NULL, NULL);

  EXPECT_EQ(PA_TZ_SUCCESS, result);
}

TEST_F(AuthenticateTest, UnauthenticatedPidHandler) {
  PaTzResult result = Authenticate(kUnauthenticatedPidHandler, NULL, 0, NULL, NULL);

  EXPECT_EQ(PA_TZ_AUTHENTICATION_FAILED, result);
}

TEST_F(AuthenticateTest, UnauthenticatedAppName) {
  PaTzResult result = Authenticate(kValidAppNameHandler, kUnauthenticatedAppName,
                                   sizeof(kUnauthenticatedAppName) - 1, NULL, NULL);

  EXPECT_EQ(PA_TZ_AUTHENTICATION_FAILED, result);
}

TEST_F(AuthenticateTest, InvalidHandler) {
  PaTzResult result = Authenticate(kInvalidHandler, NULL, 0, NULL, NULL);

  EXPECT_EQ(PA_TZ_INVALID_HANDLER, result) <<
      "Authentication SHOULD fail if flags in handler are invalid";
}

TEST_F(AuthenticateTest, Forbiden) {
  g_access_control_return = 0;

  PaTzResult result = Authenticate(kValidPidHandler, NULL, 0, NULL, NULL);

  EXPECT_TRUE(PA_TZ_CALLER_IS_FORBIDEN == result);
}

TEST_F(AuthenticateTest, ReturnValidInfo) {
  ProcessInfo info = {0};

  PaTzResult result = Authenticate(kValidPidHandler, NULL, 0, NULL, &info);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(info.id, &kValidPaId, sizeof(kValidPaId)) == 0);
}

TEST_F(AuthenticateTest, NotReturnInvalidInfo) {
  const ProcessInfo zero_info = {0};
  ProcessInfo info = {0};

  PaTzResult result = Authenticate(kInvalidHandler, NULL, 0, NULL, &info);

  EXPECT_EQ(PA_TZ_INVALID_HANDLER, result) <<
      "Authentication SHOULD fail if flags in handler are invalid";
  EXPECT_TRUE(memcmp(&zero_info, &info, sizeof(info)) == 0) <<
      "Output data SHOULD not be changed if authentication is failed";
}


TEST_F(TaskReadToTrustletTest, AllParamAreValid) {
    PaTzResult result;
    uint8_t out_trustlet;

    result = TaskReadToTrustlet(kValidPidHandler, kUserAddress, kBufferSize, &out_trustlet);
    EXPECT_TRUE(PA_TZ_SUCCESS == result);
}

TEST_F(TaskReadToTrustletTest, OutTrustletNull) {
    PaTzResult result;
    uint8_t *out_trustlet_null = NULL;

    result = TaskReadToTrustlet(kValidPidHandler, kUserAddress, kBufferSize, out_trustlet_null);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskReadToTrustletTest, NoAccessControl) {
    PaTzResult result;
    uint8_t out_trustlet;
    g_access_control_return = 0;

    result = TaskReadToTrustlet(kValidPidHandler, kUserAddress, kBufferSize, &out_trustlet);
    EXPECT_TRUE(PA_TZ_CALLER_IS_FORBIDEN == result);
}

TEST_F(TaskReadToTrustletTest, InvalidHandler) {
    PaTzResult result;
    uint8_t out_trustlet;
    PaHandler wrong_handler = {0};

    result = TaskReadToTrustlet(kInvalidHandler, kUserAddress, kBufferSize, &out_trustlet);
    EXPECT_EQ(PA_TZ_INVALID_HANDLER, result) <<
          "Reading SHOULD fail if flags in handler are invalid";
}

TEST_F(TaskReadToTrustletTest, FailedPlatformSysMapTrustlet) {
    PaTzResult result;
    uint8_t out_trustlet;
    g_platform_sys_map_trustlet_success = 0;

    result = TaskReadToTrustlet(kValidPidHandler, kUserAddress, kBufferSize, &out_trustlet);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskReadToTrustletTest, FailedTaskAccessGetBytes) {
    PaTzResult result;
    uint8_t out_trustlet;
    g_task_access_get_bytes_return = PA_TZ_GENERAL_ERROR;

    result = TaskReadToTrustlet(kValidPidHandler, kUserAddress, kBufferSize, &out_trustlet);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}


TEST_F(TaskWriteFromTrustletTest, AllParamAreValid) {
    PaTzResult result;
    uint8_t output_buffer;

    result = TaskWriteFromTrustlet(kValidPidHandler, (const void *)&kInputBuf, 
                                   kBufferSize, (ProcessAddress)&output_buffer);
    EXPECT_TRUE(PA_TZ_SUCCESS == result);
}

TEST_F(TaskWriteFromTrustletTest, InTrustletNull) {
    PaTzResult result;
    uint8_t output_buffer;
    const void *in_trustlet_null = NULL;

    result = TaskWriteFromTrustlet(kValidPidHandler, in_trustlet_null, 0, 
                                   (ProcessAddress)&output_buffer);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskWriteFromTrustletTest, NoAccessControl) {
    PaTzResult result;
    uint8_t output_buffer;
    g_access_control_return = 0;

    result = TaskWriteFromTrustlet(kValidPidHandler, (const void *)&kInputBuf, 
                                   kBufferSize, (ProcessAddress)&output_buffer);
    EXPECT_TRUE(PA_TZ_CALLER_IS_FORBIDEN == result);
}

TEST_F(TaskWriteFromTrustletTest, InvalidHandler) {
    PaTzResult result;
    uint8_t output_buffer;

    result = TaskWriteFromTrustlet(kInvalidHandler, (const void *)&kInputBuf, 
                                   kBufferSize, (ProcessAddress)&output_buffer);
    EXPECT_EQ(PA_TZ_INVALID_HANDLER, result) <<
         "Writing SHOULD fail if flags in handler are invalid";
}

TEST_F(TaskWriteFromTrustletTest, FailedPlatformSysMapTrustlet) {
    PaTzResult result;
    uint8_t output_buffer;
    g_platform_sys_map_trustlet_success = 0;

    result = TaskWriteFromTrustlet(kValidPidHandler, (const void *)&kInputBuf, 
                                   kBufferSize, (ProcessAddress)&output_buffer);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskWriteFromTrustletTest, FailedTaskAccessPutBytes) {
    PaTzResult result;
    uint8_t output_buffer;
    g_task_access_put_bytes_return = PA_TZ_GENERAL_ERROR;

    result = TaskWriteFromTrustlet(kValidPidHandler, (const void *)&kInputBuf, 
                                   kBufferSize, (ProcessAddress)&output_buffer);
    EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}