#include <gtest/gtest.h>

extern "C" {
  #include "api.h"
  #include "pa_tz_api.h"
  #include "serialize.h"
  #include "PaDriverCommand.h"
  #include "PaDriverCommandResponse.h"
  #include <tee_internal_api.h>
  #include <tees_sys.h>
}

static uint32_t g_platform_cmd = 0;
static uint32_t g_buffer_size = 128;
static uint32_t g_config_is_inited = 1;
static bool g_platform_is_device_compromised = false;
static bool g_running_convert_authenticate_result = false;
static PaTzResult g_arg_convert_authenticate_result = (PaTzResult)-1;

extern "C" uint32_t ConfigIsInited() {
  return g_config_is_inited;
}

extern "C" PaTzResult Authenticate(PaHandler handler, const char *process_name,
                                   const size_t process_name_size,
                                   const PaTzMemoryRange *memory, ProcessInfo *info) {
  if (g_running_convert_authenticate_result)
    return g_arg_convert_authenticate_result;
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskReadToTrustlet(PaHandler handler, ProcessAddress user_address,
                                         size_t size, void *out_trustlet) {
  PaTzResult result = PA_TZ_SUCCESS;
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult TaskWriteFromTrustlet(PaHandler handler, const void *in_trustlet,
                                            size_t size, ProcessAddress user_address) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult IssueNewCertificate(PaHandler handler,
                                          ProcessAddress mapped_apk,
                                          const char *package_name,
                                          const uint8_t *rsa, size_t rsa_size,
                                          uint8_t *out_xattr, uint32_t *out_xattr_size) {
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult IssueCopyOfCertificate(
    PaHandler handler,
    ProcessAddress source_file,
    ProcessAddress source_xattr,
    ProcessAddress destination_file,
    ProcessAddress out_xattr, size_t out_xattr_size) {
  return PA_TZ_SUCCESS;
}

extern "C" bool PlatformIsDeviceCompromised() {
  return g_platform_is_device_compromised;
}

class PlatformHandleCommandTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_config_is_inited = 1;
    g_platform_is_device_compromised = false;
  }

  virtual void TearDown() {
  }
};

TEST_F(PlatformHandleCommandTest, InputNull_BadParameters) {
  TEE_Result result;
  uint8_t output[128] = {0};

  result = TEES_HandleDriverCommand(g_platform_cmd,
      NULL, g_buffer_size, output, &g_buffer_size);

  ASSERT_EQ(TEE_ERROR_BAD_PARAMETERS, result) <<
      "Bad parameters SHOULD return if any input arguments are NULL";
}

TEST_F(PlatformHandleCommandTest, OutputNull_BadParameters) {
  TEE_Result result;
  uint8_t input[128] = {0};

  result = TEES_HandleDriverCommand(g_platform_cmd,
      input, sizeof(input), NULL, &g_buffer_size);

  ASSERT_EQ(TEE_ERROR_BAD_PARAMETERS, result) <<
      "Bad parameters SHOULD return if any input arguments are NULL";
}

TEST_F(PlatformHandleCommandTest, OutputSizeNull_BadParameters) {
  TEE_Result result;
  uint8_t buffer[128] = {0};

  result = TEES_HandleDriverCommand(g_platform_cmd,
      buffer, sizeof(buffer), buffer, NULL);

  ASSERT_EQ(TEE_ERROR_BAD_PARAMETERS, result) <<
      "Bad parameters SHOULD return if any input arguments are NULL";
}

TEST_F(PlatformHandleCommandTest, ConfigIsNotInited) {
  TEE_Result result;
  uint8_t buffer[128] = {0};
  uint32_t buffer_size = sizeof(buffer);

  g_config_is_inited = 0;

  result = TEES_HandleDriverCommand(g_platform_cmd,
      buffer, buffer_size, buffer, &buffer_size);

  ASSERT_EQ(TEE_ERROR_BAD_STATE, result) <<
      "Bad state error SHOULD return if config is not loaded previously";
}

class PlatformHandleCommandTestDriver : public PlatformHandleCommandTest {
 protected:
  virtual void SetUp() {
    PlatformHandleCommandTest::SetUp();
    response_ = NULL;
    init_buffers();
  }

  virtual void TearDown() {
    g_platform_is_device_compromised = false;
    PlatformHandleCommandTest::TearDown();
    free_response();
  }

  void free_response() {
    PaFreeDriverCommandResponse(response_);
    response_ = NULL;
  }

  void init_buffers() {
    memset(buffer, 0, sizeof(buffer));
    buffer_size = sizeof(buffer);

    memset(encoded_response_buffer, 0, sizeof(encoded_response_buffer));
    encoded_response_buffer_size = sizeof(encoded_response_buffer);
  }

  void set_driver_command(PaDriverCommand_PR aCommand) {
    memset(&driver_command, 0, sizeof(driver_command));
    driver_command.present = aCommand;
  }

  PaDriverCommandResponse_t *response_ = NULL;

  uint8_t buffer[128];
  uint32_t buffer_size;

  PaDriverCommand_t driver_command;

  uint8_t encoded_response_buffer[128];
  uint32_t encoded_response_buffer_size;
};

TEST_F(PlatformHandleCommandTestDriver, PlatformIsDeviceCompromisedFalse) {
  set_driver_command(PaDriverCommand_PR_authenticateCommand);

  TEE_Result result = PaEncodeDriverCommand(&driver_command,
                                            buffer,
                                            &buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error encoding command";

  g_platform_is_device_compromised = false;

  result = TEES_HandleDriverCommand(g_platform_cmd,
                                    buffer,
                                    buffer_size,
                                    encoded_response_buffer,
                                    &encoded_response_buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result)
    << "There should not be an error here if app could not be authenticated";

  result = PaDecodeDriverCommandResponse(encoded_response_buffer,
                                         encoded_response_buffer_size,
                                         &response_);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error decoding response";
  ASSERT_EQ(PaDriverCommandResponse_PR_authenticateResponse,
      response_->present)
    << "Unexpected response_->present";
  ASSERT_EQ(PaDriverAuthenticateResult_paAuthenticated,
      response_->choice.result)
    << "Unexpected response_->choice.result";
}

TEST_F(PlatformHandleCommandTestDriver, PlatformIsDeviceCompromisedTrue) {
  set_driver_command(PaDriverCommand_PR_authenticateCommand);

  TEE_Result result = PaEncodeDriverCommand(&driver_command, buffer,
      &buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error encoding command";

  g_platform_is_device_compromised = true;

  result = TEES_HandleDriverCommand(g_platform_cmd,
                                    buffer,
                                    buffer_size,
                                    encoded_response_buffer,
                                    &encoded_response_buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result) <<
      "There should not be an error here if app could not be authenticated";

  result = PaDecodeDriverCommandResponse(encoded_response_buffer,
                                         encoded_response_buffer_size,
                                         &response_);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error decoding response";
  ASSERT_EQ(PaDriverCommandResponse_PR_authenticateResponse,
            response_->present) << "Unexpected response_->present";
  ASSERT_EQ(PaDriverAuthenticateResult_paDeviceIsCompromised,
            response_->choice.result) << "Unexpected response_->choice.result";
}

struct ResponseMapper {
  PaTzResult                   arg; 
  e_PaDriverAuthenticateResult val;
};

struct ResponseMapperNamed {
  PaTzResult                   arg_;
  e_PaDriverAuthenticateResult val_;
  const char *arg_name_;
  const char *val_name_;

  ResponseMapperNamed(PaTzResult arg,
                      e_PaDriverAuthenticateResult val,
                      const char *arg_name,
                      const char *val_name)
      : arg_(arg), val_(val), arg_name_(arg_name), val_name_(val_name) {
  }

  friend std::ostream& operator << (std::ostream& os,
                                    const ResponseMapperNamed &src) {
    return os << "mapping(" << src.arg_name_ << ")=>" << src.val_name_;
  }
};

class PlatformHandleCommandTestDriverWithParam : 
            public ::testing::TestWithParam<ResponseMapperNamed> {
 protected:
  virtual void SetUp() {
    g_config_is_inited = 1;
    g_platform_is_device_compromised = false;
    response_ = NULL;
    init_buffers();
  }

  virtual void TearDown() {
    g_running_convert_authenticate_result = false;
    free_response();
  }

  void free_response() {
    PaFreeDriverCommandResponse(response_);
    response_ = NULL;
  }

  void init_buffers() {
    memset(encoded_command_buffer, 0, sizeof(encoded_command_buffer));
    encoded_command_buffer_size = sizeof(encoded_command_buffer);

    memset(encoded_response_buffer, 0, sizeof(encoded_response_buffer));
    encoded_response_buffer_size = sizeof(encoded_response_buffer);
  }

  void set_driver_command(PaDriverCommand_PR aCommand) {
    memset(&driver_command, 0, sizeof(driver_command));
    driver_command.present = aCommand;
  }

  PaDriverCommandResponse_t *response_ = NULL;

  uint8_t encoded_command_buffer[128];
  uint32_t encoded_command_buffer_size;

  PaDriverCommand_t driver_command;

  uint8_t encoded_response_buffer[128];
  uint32_t encoded_response_buffer_size;  
};

TEST_P(PlatformHandleCommandTestDriverWithParam, ConvertAuthenticateResult) {
  set_driver_command(PaDriverCommand_PR_authenticateCommand);

  TEE_Result result = PaEncodeDriverCommand(&driver_command,
      encoded_command_buffer, &encoded_command_buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error encoding command";

  g_running_convert_authenticate_result = true;

  g_arg_convert_authenticate_result = GetParam().arg_;
  result = TEES_HandleDriverCommand(g_platform_cmd, encoded_command_buffer,
      encoded_command_buffer_size, encoded_response_buffer,
      &encoded_response_buffer_size);
  ASSERT_EQ(TEE_SUCCESS, result)
    << "There should not be an error here if app could not be authenticated";

  result = PaDecodeDriverCommandResponse(encoded_response_buffer,
      encoded_response_buffer_size, &response_);
  ASSERT_EQ(TEE_SUCCESS, result) << "Error decoding response";
  ASSERT_EQ(PaDriverCommandResponse_PR_authenticateResponse, response_->present)
    << "Unexpected response_->present";

  auto PaDriver_response = GetParam().val_;

  ASSERT_EQ(PaDriver_response, response_->choice.result);
  free_response();
}

#define QUOTE(name) #name
#define bindResponseMapperNamed(arg, val) \
   ResponseMapperNamed(arg, val, QUOTE(arg), QUOTE(val))

constexpr PaTzResult PA_TZ_other_values = (PaTzResult)-1;

INSTANTIATE_TEST_CASE_P(ConvertAuthenticateResult,
                        PlatformHandleCommandTestDriverWithParam,
                        ::testing::Values(
  bindResponseMapperNamed(PA_TZ_SUCCESS, PaDriverAuthenticateResult_paAuthenticated),
  bindResponseMapperNamed(PA_TZ_AF_TASK_IS_NOT_FOUND, PaDriverAuthenticateResult_paTaskIsNotFound),
  bindResponseMapperNamed(PA_TZ_AF_INTEGRITY_IS_NONE, PaDriverAuthenticateResult_paIntegrityIsNone),
  bindResponseMapperNamed(PA_TZ_AF_INTEGRITY_IS_NOT_READY, PaDriverAuthenticateResult_paIntegrityIsNotready),
  bindResponseMapperNamed(PA_TZ_AF_CERTIFICATE_IS_ABSENT, PaDriverAuthenticateResult_paCertificateIsAbsent),
  bindResponseMapperNamed(PA_TZ_AF_CERTIFICATE_IS_INCORRECT, PaDriverAuthenticateResult_paCertificateIsIncorrect),
  bindResponseMapperNamed(PA_TZ_AF_CERTIFICATE_IS_NOT_MATCH, PaDriverAuthenticateResult_paCertificateIsNotMatch),
  bindResponseMapperNamed(PA_TZ_AF_APPNAME_IS_INCORRECT, PaDriverAuthenticateResult_paAppNameIsIncorrect),
  bindResponseMapperNamed(PA_TZ_AF_TASK_HAS_DUPLICATE, PaDriverAuthenticateResult_paAppNameIsNotUnique),
  bindResponseMapperNamed(PA_TZ_PROCA_NOT_SUPPORTED, PaDriverAuthenticateResult_paNotSupported),
  bindResponseMapperNamed(PA_TZ_other_values, PaDriverAuthenticateResult_paNotAuthenticated)));