#include <gtest/gtest.h>

#include <stdint.h>

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

#include "config.h"
#include "crypto.h"
#include "task.h"
#include "tlb.h"
#include "task_parser.h"
}

static const KernelAddress kKernelValidStart = 0x1000;
static const TaskInfo kEmptyTask = {0};
static const ProcessAddress kVmStart = 0x1000;
static const ProcessAddress kVmEnd = 0x9000;
static const KernelAddress kVmNext = 0x2000;
static const PhysicalAddress kPhysAddrToFind = 0x3000;
static const uint32_t kMaxNumHop  = 0xFFFF;
static const uint32_t kMaxNumVma = 4096;
static const int kDefaultTaskCount = 10;
static const uint64_t kInvalidHandlerData = 0;

const uint32_t kExistentPid = 42;
const uint32_t kNonExistentPid = 24;

const char kExistentAppName[] = "existent_app_name";
const char kNonExistentAppName[] = "nonexistent_app_name";

static const int kSignBufferSize = 8;
static const uint8_t kSignBuffer [kSignBufferSize] = {0, 1, 2, 3, 4, 5, 6, 7};
static const OCTET_STRING_t kFiveSignature = {(uint8_t *)&kSignBuffer, kSignBufferSize, {}};
static const PaData_t kPaData = {0, 0, 0, {}, (OCTET_STRING_t)kFiveSignature, {}};
static const PaCertificate_t kCertificate = {kPaData, {}, {}};

static VmaInfo g_init_vma;
static RbRoot g_rb_root = {(KernelAddress)&g_init_vma + 16};
static TaskInfo g_init_task = {(KernelAddress)&g_init_task, 0, 0,
                               0x0, 0x0, 0x0, 0, NULL, g_rb_root, 0};
static TaskInfo g_out_task;
static PaConfig g_pa_config = {0};

static int g_vma_count = 1;
static int g_find_vma_num = 0;
static int g_not_parse_five_sign = 0;
static uint8_t *g_sign_buf = (uint8_t *)&kSignBuffer;
static int g_sign_buf_size = kSignBufferSize;

static PaTzResult g_parse_task_integrity_result = PA_TZ_SUCCESS;
TaskDescriptor g_pa_task_descriptor = {0};

/*
 * Mocked functions
 */
extern "C" PaTzResult TaskParsePaCertificate(const TaskInfo *task,
                                             PaCertificate_t **certificate) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskParseVmaAddress(const TaskInfo *task, ProcessAddress address,
                                          ProcessAddress *vma_address) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskDescriptorFindByPid(uint32_t pid, TaskDescriptor *out_descr) {
  if (pid == kExistentPid) {
    memset(out_descr, 0xff, sizeof(*out_descr));
    return PA_TZ_SUCCESS;
  } else {
    return PA_TZ_AF_TASK_IS_NOT_FOUND;
  }
}

extern "C" PaTzResult TaskDescriptorFindByAppName(const char *app_name_str,
                                       const size_t app_name_len,
                                       TaskDescriptor *out_descr) {
  if (app_name_len != sizeof(kExistentAppName) - 1) {
    return PA_TZ_AF_TASK_IS_NOT_FOUND;
  }

  if (!strncmp(app_name_str, kExistentAppName, sizeof(kExistentAppName) - 1)) {
    memset(out_descr, 0xff, sizeof(*out_descr));
    return PA_TZ_SUCCESS;
  } else {
    return PA_TZ_AF_TASK_IS_NOT_FOUND;
  }
}

extern "C" PaTzResult TaskParseBasicInfo(KernelAddress virt, TaskInfo *g_out_task) {
  if (virt == 0) {
    return PA_TZ_GENERAL_ERROR;
  }
  return PA_TZ_SUCCESS;
}

extern "C" uint32_t IsKernelVirtualAddress(KernelAddress kernel_virt) {
  if (kernel_virt >= kKernelValidStart) {
    return 1;
  }

  return 0;
}

extern "C" PaTzResult TaskParseIntegrity(TaskInfo *g_out_task) {
  return g_parse_task_integrity_result;
}

extern "C" PaTzResult TaskAddressToPhysical(TlbConverterInfo* context, KernelAddress pgd_address,
                                            ProcessAddress user_virt,
                                            PhysicalAddress *phys,
                                            AccessPermissionFlags* flags) {
  *phys = user_virt;
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskParseVmaInfo(KernelAddress virt, VmaInfo *out_vma) {
  if (!virt) {
    return PA_TZ_GENERAL_ERROR;
  }

  memcpy(out_vma, (void *)virt, sizeof(VmaInfo));
  return PA_TZ_SUCCESS;
}

extern "C" const PaConfig *GetConfig(void) {
  return &g_pa_config;
}

extern "C" PaTzResult TaskParseFiveSignature(KernelAddress file_struct, void *signature, 
                                             size_t *signature_len) {
  if (g_not_parse_five_sign) {
    return PA_TZ_GENERAL_ERROR;
  }
  memcpy(signature, g_sign_buf, *signature_len * sizeof(uint8_t));
  *signature_len = g_sign_buf_size;

  return PA_TZ_SUCCESS;
}

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

extern "C" PaTzResult CryptoSha256(const uint8_t *data, size_t data_length, uint8_t *out_md) {
  return PA_TZ_SUCCESS;
}
/* End */

// CheckPhysicalCommandBufferTest ==============================================

class CheckPhysicalCommandBufferTest : public ::testing::Test {
protected:
  virtual void SetUp() {
  }
  virtual void TearDown() {
  }
};

TEST_F(CheckPhysicalCommandBufferTest, AllParamAreValidP) {
  uint32_t result;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, MinAddrP) {
  uint32_t result;
  PhysicalAddress min_phys_addr_to_find = 0x0;
  const ProcessAddress kVmMinStart = 0x0;
  const VmaInfo kInitVma = {kVmMinStart, kVmEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &min_phys_addr_to_find);
  ASSERT_EQ(result, kEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, BigTest_MaxAddrP) {
  uint32_t result;
  PhysicalAddress max_phys_addr_to_find = 0x007FFFFFE000;
  const ProcessAddress kVmMaxEnd = 0x007FFFFFF000;
  const VmaInfo kInitVma = {kVmStart, kVmMaxEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &max_phys_addr_to_find);
  ASSERT_EQ(result, kEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, TaskNullN) {
  uint32_t result;
  TaskInfo *task = NULL;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(task, &kInitVma, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, VmaNullN) {
  uint32_t result;
  const VmaInfo *kInitVmaNull = NULL;

  result = CheckPhysicalCommandBuffer(&kEmptyTask, kInitVmaNull, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, PhysAddrNULL) {
  uint32_t result;
  PhysicalAddress *phys_addr_to_find_NULL = NULL;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *)phys_addr_to_find_NULL);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, AbsentVmWriteFlagP) {
  uint32_t result;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_READ), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, AbsentVmReadFlagP) {
  uint32_t result;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, PresentVmExecFlagP) {
  uint32_t result;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_EXEC | VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &kPhysAddrToFind);
  ASSERT_EQ(result, kNonEqual);
}

TEST_F(CheckPhysicalCommandBufferTest, WrongPhysAddrToFindP) {
  uint32_t result;
  PhysicalAddress wrong_phys_addr = 0x10000;
  const VmaInfo kInitVma = {kVmStart, kVmEnd, kVmNext, (VM_READ | VM_WRITE), 0};

  result = CheckPhysicalCommandBuffer(&kEmptyTask, &kInitVma, (const void *) &wrong_phys_addr);
  ASSERT_EQ(result, kNonEqual);
}

// TaskFindVmaTest =============================================================

static FindResult TestCustomFind(const TaskInfo *task, const VmaInfo *vma, const void* param) {
  int required_vma = *((int *)param);
  if (required_vma == vma->vm_flags) {
    return kEqual;
  }

  return kNonEqual;
}

static void InitVmaInfo() {
  g_init_vma = {0x0, 0x0, 0x0, 0x0, 0x1};
  VmaInfo *prev_vma = &g_init_vma;
  VmaInfo *curr_vma = NULL;

  for (int i = 1; i < g_vma_count; i++) {
    curr_vma = (VmaInfo *)malloc(sizeof(VmaInfo));
    *curr_vma = {0x0, 0x0, 0x0, 0x0, 0x1};
    curr_vma->vm_flags = i;
    prev_vma->vm_next = (KernelAddress)curr_vma;
    prev_vma = curr_vma;
  }
}

static void FreeVmaInfo() {
  VmaInfo *curr_vma = (VmaInfo *)g_init_vma.vm_next;
  VmaInfo *next_vma = NULL;

  while (curr_vma) {
    next_vma = (VmaInfo *)curr_vma->vm_next;
    free(curr_vma);
    curr_vma = (VmaInfo *)next_vma;
  }
}

class TaskFindVmaTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_vma_count = 5;
    g_find_vma_num = 3;
    g_sign_buf = (uint8_t *)&kSignBuffer;
    g_sign_buf_size = kSignBufferSize;
    InitVmaInfo();
    g_init_task.virt = (KernelAddress)&g_init_task;
    g_init_task.mmap = (KernelAddress)&g_init_vma;
    g_init_task.certificate = (PaCertificate_t *)&kCertificate;
    g_not_parse_five_sign = 0;
  }
  virtual void TearDown() {
    FreeVmaInfo();
  }
};

TEST_F(TaskFindVmaTest, AllParamAreValidP) {
  PaTzResult result;
  VmaInfo out_vma;
  g_find_vma_num = 3;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_SUCCESS);
}

TEST_F(TaskFindVmaTest, BigTest_MoreThenMaxNumberOfVmaInListP) {
  PaTzResult result;
  VmaInfo out_vma;
  g_vma_count = kMaxNumVma + 2;
  g_find_vma_num = g_vma_count - 1;

  InitVmaInfo();

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, NoRequiredVmaInListP) {
  PaTzResult result;
  VmaInfo out_vma;
  g_find_vma_num = g_vma_count;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, TaskIsNullN) {
  PaTzResult result;
  VmaInfo out_vma;
  const TaskInfo *kNullTask = NULL;

  result = TaskFindVmaLinear(kNullTask, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, CustomFindIsNullN) {
  PaTzResult result;
  VmaInfo out_vma;
  const CustomFind kNullCustomFind = NULL;

  result = TaskFindVmaLinear(&g_init_task, kNullCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, OutVmaIsNullN) {
  PaTzResult result;
  VmaInfo *null_out_vma = NULL;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, null_out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, WrongMmapAddrN) {
  PaTzResult result;
  VmaInfo out_vma;
  g_init_task.mmap = 0x0;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);
}

TEST_F(TaskFindVmaTest, VmNextIsNotKernelVirtAddrN) {
  PaTzResult result;
  VmaInfo out_vma;
  VmaInfo *changed_vma = ((VmaInfo *)(g_init_vma.vm_next));
  KernelAddress vm_next = changed_vma->vm_next;
  changed_vma->vm_next = 0x1;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);

  changed_vma->vm_next = vm_next;
}

TEST_F(TaskFindVmaTest, LoopOfVmaListN) {
  PaTzResult result;
  VmaInfo out_vma;
  g_find_vma_num = g_vma_count;
  VmaInfo *temp_vma = &g_init_vma;
  while (temp_vma->vm_next) {
    temp_vma = (VmaInfo *)temp_vma->vm_next;
  }
  temp_vma->vm_next = (KernelAddress)&g_init_vma;

  result = TaskFindVmaLinear(&g_init_task, TestCustomFind, (const void *)&g_find_vma_num, &out_vma);
  ASSERT_EQ(result, PA_TZ_GENERAL_ERROR);

  temp_vma->vm_next = (KernelAddress)NULL;
}

// TaskFindByPidTest ====================================================

class TaskFindByPidTest : public ::testing::Test {
protected:

  virtual void SetUp() {
    g_parse_task_integrity_result = PA_TZ_SUCCESS;
  }

  virtual void TearDown() {
  }
};

TEST_F(TaskFindByPidTest, ParseIntegrityFailed_Error) {
  g_parse_task_integrity_result = PA_TZ_GENERAL_ERROR;

  PaTzResult result = TaskFindByPid(kExistentPid, &g_out_task);
  ASSERT_EQ(PA_TZ_GENERAL_ERROR, result);
}

TEST_F(TaskFindByPidTest, NonExistentPidSearch_Error) {
  PaTzResult result = TaskFindByPid(kNonExistentPid, &g_out_task);
  ASSERT_EQ(PA_TZ_AF_TASK_IS_NOT_FOUND, result);
}

TEST_F(TaskFindByPidTest, ExistentPidSearch_Success) {
  PaTzResult result = TaskFindByPid(kExistentPid, &g_out_task);
  ASSERT_EQ(PA_TZ_SUCCESS, result);
}

// TaskFindByAppNameTest ====================================================

class TaskFindByAppNameTest : public ::testing::Test {
protected:

  virtual void SetUp() {
  }

  virtual void TearDown() {
  }
};

TEST_F(TaskFindByAppNameTest, NonExistentAppNameSearch_Error) {
  PaTzResult result = TaskFindByAppName(kNonExistentAppName,
                                        sizeof(kNonExistentAppName) - 1,
                                        &g_out_task);
  ASSERT_EQ(PA_TZ_AF_TASK_IS_NOT_FOUND, result);
}

TEST_F(TaskFindByAppNameTest, ExistentAppNameSearch_Success) {
  PaTzResult result = TaskFindByAppName(kExistentAppName,
                                        sizeof(kExistentAppName) - 1,
                                        &g_out_task);
  ASSERT_EQ(PA_TZ_SUCCESS, result);
}