#include "kg_policy.h"

int jsoneq(const char *json, jsmntok_t *tok, const char *s) {
  if (tok->type == JSMN_STRING && (int)strlen(s) == tok->end - tok->start &&
      strncmp(json + tok->start, s, tok->end - tok->start) == 0) {
    return 0;
  }
  return -1;
}

static uint32_t process_policy(const char *data, uint32_t len) {
    KG_LOG("Process policy\n");
    int i;
    int r;
    jsmn_parser p;
    jsmntok_t t[512]; /* We expect no more than 128 tokens */

    jsmn_init(&p);
    TEE_MemFill(t, 0x0, sizeof(jsmntok_t)*512);
    r = jsmn_parse(&p, data, len, t, sizeof(t) / sizeof(t[0]));
    if (r < 0) {
        KG_LOG("Failed to parse JSON: %d\n", r);
        return KG_POLICY_TOKEN_FAIL;
    }

    if (r < 1 || t[0].type != JSMN_OBJECT) {
        KG_LOG("Object expected\n");
        return KG_POLICY_FORMAT_FAIL;
    }

    if (r > 511){
        KG_LOG("Possible read out of array bounds\n");
        return KG_POLICY_TOKEN_FAIL; 
    }

    for (i = 1; i < r; i++) {
        if (jsoneq(data, &t[i], KG_POLICY_POLICYTIME) == 0) {
            KG_LOG_DBG("- timeStamp: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            i++;
        } else if (jsoneq(data, &t[i], KG_POLICY_ACTION) == 0) {
            KG_LOG_DBG("- action: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            i++;
        } else if (jsoneq(data, &t[i], KG_POLICY_KGSTATE) == 0) {
            KG_LOG_DBG("- kgState: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            i++;
        }
    }
    return KG_SUCCESS;
}

static uint32_t process_metadata(const char *data, uint32_t len, kg_metadata_t* metadata) {
    KG_LOG("Process metadata\n");
    uint32_t ret = KG_POLICY_PARSE_FAIL;
    int i;
    int r;
    jsmn_parser p;
    jsmntok_t t[512]; /* We expect no more than 128 tokens */
    
    uint8_t *metadata_buf = NULL;
    uint32_t metadata_buf_len = KG_BUF_LEN;

    metadata_buf = TEE_Malloc(metadata_buf_len, 0);
    if (metadata_buf == NULL) {
        KG_LOG("Failed to alloc buffer for metadata\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    jsmn_init(&p);
    TEE_MemFill(t, 0x0, sizeof(jsmntok_t)*512);
    r = jsmn_parse(&p, data, len, t, sizeof(t) / sizeof(t[0]));
    if (r < 0) {
        KG_LOG("Failed to parse JSON: %d\n", r);
        ret = KG_POLICY_TOKEN_FAIL;
        goto exit;
    }

    if (r < 1 || t[0].type != JSMN_OBJECT) {
        KG_LOG("Object expected\n");
        ret = KG_POLICY_FORMAT_FAIL;
        goto exit;
    }

    if (r > 511){
        KG_LOG("Possible read out of array bounds\n");
        return KG_POLICY_TOKEN_FAIL; 
    }

    for (i = 1; i < r; i++) {
        if (jsoneq(data, &t[i], KG_POLICY_VERSION) == 0) {
            KG_LOG_DBG("- version: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            TEE_MemFill(metadata_buf, 0x0, metadata_buf_len);
            if(t[i + 1].end - t[i + 1].start > metadata_buf_len){
                KG_LOG("Data size overflow\n");
                ret = KG_BUFFER_SIZE_FAIL;
                goto exit;
            }
            TEE_MemMove(metadata_buf, (void *)(data + t[i + 1].start), t[i + 1].end - t[i + 1].start);
#if defined(TEEGRIS_SDK_VER) && (TEEGRIS_SDK_VER >= 4)
            uint32_t policy_version = atoi((char *)metadata_buf);
#else
            char *ptr;
            uint32_t policy_version = strtol((char *)metadata_buf, &ptr, 10);
#endif
            KG_LOG_DBG("policy_version integer is %d\n", policy_version);
            
            if (metadata->policy_version <= policy_version) {
                metadata->policy_version = policy_version;
            } else {
                KG_LOG_DBG("policy version check failed against rpmb value\n");
                ret = KG_POLICY_VER_CHECK_FAIL;
                goto exit;
            }
            
            i++;
        } else if (jsoneq(data, &t[i], KG_POLICY_COUNTER) == 0) {
            KG_LOG_DBG("- counter: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            TEE_MemFill(metadata_buf, 0x0, metadata_buf_len);
            if(t[i + 1].end - t[i + 1].start > metadata_buf_len){
                KG_LOG("Data size overflow\n");
                ret = KG_BUFFER_SIZE_FAIL;
                goto exit;
            }
            TEE_MemMove(metadata_buf, (void *)(data + t[i + 1].start), t[i + 1].end - t[i + 1].start);
#if defined(TEEGRIS_SDK_VER) && (TEEGRIS_SDK_VER >= 4)
            uint32_t policy_counter = atoi((char *)metadata_buf);
#else
            char *ptr;
            uint32_t policy_counter = strtol((char *)metadata_buf, &ptr, 10);
#endif
            KG_LOG_DBG("policy_counter integer is %d\n", policy_counter);
            
            if (metadata->policy_counter <= policy_counter) {
                metadata->policy_counter = policy_counter;
            } else {
                KG_LOG_DBG("policy counter check failed against rpmb value\n");
                ret = KG_POLICY_CTR_CHECK_FAIL;
                goto exit;
            }
            
            i++;
        } else if (jsoneq(data, &t[i], KG_POLICY_KGID) == 0) {
            KG_LOG_DBG("- kgId: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            if (t[i + 1].end - t[i + 1].start > 24) {
                KG_LOG_DBG("KG policy kgid length check failed\n");
                ret = KG_POLICY_KGID_CHECK_FAIL;
                goto exit;
            }
            uint8_t kg_id[KG_ID_LEN] = {0x0};
            TEE_MemMove(kg_id, (void *)(data + t[i + 1].start), t[i + 1].end - t[i + 1].start);
            // check if kgid is the same
            
            uint32_t idx = 0;
            bool check_flag = true;
            while (idx < KG_ID_LEN) {
                if (metadata->kg_id[idx] != kg_id[idx]) {
                    check_flag = false;
                }
                idx++;
            }
            
            if (false == check_flag) {
                KG_LOG_DBG("KG policy kgid check failed\n");
                ret = KG_POLICY_KGID_CHECK_FAIL;
                goto exit;
            }

            i++;
        } else if (jsoneq(data, &t[i], KG_POLICY_SERVERTIME) == 0) {
            KG_LOG_DBG("- serverTimestamp: %.*s\n", t[i + 1].end - t[i + 1].start, data + t[i + 1].start);
            TEE_MemFill(metadata_buf, 0x0, metadata_buf_len);
            if(t[i + 1].end - t[i + 1].start > metadata_buf_len){
                KG_LOG("Data size overflow\n");
                ret = KG_BUFFER_SIZE_FAIL;
                goto exit;
            }
            TEE_MemMove(metadata_buf, (void *)(data + t[i + 1].start), t[i + 1].end - t[i + 1].start);
            char *ptr;
            uint64_t server_time = strtoll((char *)metadata_buf, &ptr, 10);
            KG_LOG_DBG("server_time long is %llu\n", server_time);
            
            if (metadata->svr_timestamp <= server_time) {
                metadata->svr_timestamp = server_time;
            } else {
                KG_LOG_DBG("policy svr timestamp check failed\n");
                ret = KG_POLICY_TIME_CHECK_FAIL;
                goto exit;
            }
            i++;
        }
    }

    ret = KG_SUCCESS;
exit:
    if (metadata_buf != NULL) {
        TEE_Free(metadata_buf);
        metadata_buf = NULL;
    }
    return ret;
}

uint32_t parse_policy(const char *policy, uint32_t policy_len, kg_metadata_t* metadata)
{
    int i;
    int r;
    jsmn_parser p;
    jsmntok_t t[1024];

    jsmn_init(&p);
    TEE_MemFill(t, 0x0, sizeof(jsmntok_t)*1024);

    r = jsmn_parse(&p, policy, strlen(policy), t, sizeof(t) / sizeof(t[0]));
    if (r < 0) {
        KG_LOG("Failed to parse JSON: %d\n", r);
        return KG_POLICY_TOKEN_FAIL;
    }

    if (r < 1 || t[0].type != JSMN_OBJECT) {
        KG_LOG("Object expected\n");
        return KG_POLICY_FORMAT_FAIL;
    }

    if (r > 1023){
        KG_LOG("Possible read out of array bounds\n");
        return KG_POLICY_TOKEN_FAIL; 
    }

    for (i = 1; i < r; i++) {
        if (jsoneq(policy, &t[i], KG_POLICY_METADATA) == 0) {
            KG_LOG_DBG("- metadata: %.*s\n", t[i + 1].end - t[i + 1].start, policy + t[i + 1].start);
            if (KG_SUCCESS != process_metadata(policy + t[i + 1].start, t[i + 1].end - t[i + 1].start, metadata)) {
                KG_LOG_DBG("Failed to parse policy metadata");
                return KG_POLICY_PARSE_FAIL;
            }
            i++;
        } else if (jsoneq(policy, &t[i], KG_POLICY_POLICY) == 0) {
            int j;
            KG_LOG_DBG("- Policies:\n");
            KG_LOG_DBG("list size: %d\n", t[i + 1].size);
            KG_LOG_DBG("- policies: %.*s\n", t[i + 1].end - t[i + 1].start, policy + t[i + 1].start);
            for (j = 0; j < t[i + 1].size * 4; j += 5) {
                if( i + j + 2 >= r){  
                    KG_LOG_DBG("uninitialized data usage\n");
                    return KG_BUFFER_SIZE_FAIL;
                }
                jsmntok_t *g = &t[i + j + 2];
                if (KG_SUCCESS != process_policy(policy + g->start, g->end - g->start)) {
                    KG_LOG_DBG("Failed to process policy\n");
                    return KG_POLICY_PARSE_FAIL;
                }
                //KG_LOG_DBG("  * %.*s\n", g->end - g->start, policy + g->start);
            }
            i += t[i + 1].size + 1;
        }
    }
    return KG_SUCCESS;
}