#include <rpmb.h>
#include <string.h>

#include "rpmb_storage.h"
#include "debug_log.h"

#define TZ_RPMB_ARCOUNTER_PARTITION_ID 13

int arcounter_rpmb_init()
{
  TEE_Result ret = TEES_RPMBCheckEnable();
  if (TEE_SUCCESS != ret) {
    TTY_LOG("RPMB is unavailable : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }
  return ARCOUNTER_RPMB_SUCCESS;
}

int arcounter_rpmb_read(uint32_t sector_pos, uint8_t* data)
{
  TEE_Result ret = TEE_SUCCESS;

  ret = TEES_RPMBRead(TZ_RPMB_ARCOUNTER_PARTITION_ID, sector_pos, data, 1, RPMB_TYPE_BLOCK);
  if (TEE_SUCCESS != ret) {
    TTY_LOG("Failed to read RPMB : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }
  TTY_LOG("Success to read RPMB : 0x%x", ret);
  return ARCOUNTER_RPMB_SUCCESS;
}

int arcounter_rpmb_write(uint32_t sector_pos, uint8_t* data)
{
  TEE_Result ret = TEE_SUCCESS;

  ret = TEES_RPMBWrite(TZ_RPMB_ARCOUNTER_PARTITION_ID, sector_pos, data, 1, RPMB_TYPE_BLOCK);
  if (TEE_SUCCESS != ret) {
    TTY_LOG("Failed to write RPMB : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }
  TTY_LOG("success to write RPMB : 0x%x", ret);
  return ARCOUNTER_RPMB_SUCCESS;
}

int arcounter_rpmb_set_flag(uint8_t flag)
{
  TEE_Result ret = TEE_SUCCESS;
  uint8_t curr_flag = 0x0;
  uint8_t buffer[TZ_RPMB_ARCOUNTER_BLOCK_SIZE] = {0,};

  TTY_LOG("arcounter flag : 0x%x", flag);
  ret = TEES_RPMBRead(TZ_RPMB_ARCOUNTER_PARTITION_ID, 0, buffer, 1, RPMB_TYPE_BLOCK);
  if (TEE_SUCCESS != ret) {
    TTY_LOG("Failed to read RPMB : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }

  memcpy(&curr_flag, buffer+ENABLE_FLAG_OFFSET, ENABLE_FLAG_SIZE);
  if (curr_flag == flag) {
    TTY_LOG("flag is already set: %d", flag);
    return ARCOUNTER_RPMB_SUCCESS;
  }
  if (flag == 0x0) {
    memset(buffer, 0x0, TZ_RPMB_ARCOUNTER_BLOCK_SIZE);
  }
  memcpy(buffer+ENABLE_FLAG_OFFSET, &flag, ENABLE_FLAG_SIZE);
  ret = TEES_RPMBWrite(TZ_RPMB_ARCOUNTER_PARTITION_ID, 0, buffer, 1, RPMB_TYPE_BLOCK);
  if (TEE_SUCCESS != ret) {
    TTY_LOG("Failed to write RPMB : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }
  TTY_LOG("arcounter_rpmb_set_flag successful : %d", flag);

  return ARCOUNTER_RPMB_SUCCESS;
}

int arcounter_check_flag(uint8_t *flag)
{
  TEE_Result ret = TEE_SUCCESS;
  uint8_t buffer[TZ_RPMB_ARCOUNTER_BLOCK_SIZE] = {0,};

  ret = TEES_RPMBRead(TZ_RPMB_ARCOUNTER_PARTITION_ID, 0, buffer, 1, RPMB_TYPE_BLOCK);
  if (TEE_SUCCESS != ret) {
    TTY_LOG("Failed to read RPMB : 0x%x", ret);
    return ARCOUNTER_RPMB_FAILURE;
  }

  memcpy(flag, buffer+ENABLE_FLAG_OFFSET, ENABLE_FLAG_SIZE);
  TTY_LOG("arcounter flag : 0x%x", *flag);

  return ARCOUNTER_RPMB_SUCCESS;
}