#include <tfe_utils.h>
#include <MESA/maat.h>
#include <cjson/cJSON.h>
#include <tfe_resource.h>
#include <intercept_policy.h>

struct intercept_param
{
    int vsys_id;
    uint64_t rule_id;
    int ref_cnt;
    int action;
    int keyring_for_trusted;
    int keyring_for_untrusted;
    int decryption_profile;
    int tcp_option_profile;
};

struct intercept_policy_enforcer
{
    struct maat *maat;
    int table_id;
    void *logger;
};

static void intercept_param_new_cb(const char *table_name, int table_id, const char *key, const char *table_line, void **ad, long argl, void *argp)
{
    int action = 0;
    int vsys_id = 0;
    size_t len = 0;
    size_t offset = 0;
    char buffer[8] = {0};
    char *json_str = NULL;
    cJSON *json = NULL;
    cJSON *item = NULL;
    struct intercept_param *param = NULL;
    struct intercept_policy_enforcer *enforcer = (struct intercept_policy_enforcer *)argp;

    if (maat_helper_read_column(table_line, 3, &offset, &len) < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept action: %s", table_line);
        goto error_out;
    }

    memcpy(buffer, table_line + offset, MIN(sizeof(buffer), len));
    action = atoi(buffer);
    if (action != 2 && action != 3)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept action: %s", table_line);
        goto error_out;
    }

    if (maat_helper_read_column(table_line, 7, &offset, &len) < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept user region: %s", table_line);
        goto error_out;
    }

    json_str = ALLOC(char, len + 1);
    memcpy(json_str, table_line + offset, len);
    json = cJSON_Parse(json_str);
    if (json == NULL)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: id = %s", key);
        goto error_out;
    }

    item = cJSON_GetObjectItem(json, "vsys_id");
    if (!item || !cJSON_IsNumber(item))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %s invalid vsys_id format", key);
        goto error_out;
    }
    vsys_id = item->valueint;

    param = ALLOC(struct intercept_param, 1);
    param->vsys_id = vsys_id;
    param->rule_id = atoll(key);
    param->ref_cnt = 1;
    param->action = action;
    param->keyring_for_trusted = 1;
    param->keyring_for_untrusted = 0;
    param->decryption_profile = 0;
    param->tcp_option_profile = 0;

    item = cJSON_GetObjectItem(json, "keyring_for_trusted");
    if (item)
    {
        if (item->type == cJSON_Number)
        {
            param->keyring_for_trusted = item->valueint;
        }
        else if (item->type == cJSON_String)
        {
            param->keyring_for_trusted = atoi(item->valuestring);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %lu invalid keyring_for_trusted format", param->rule_id);
        }
    }

    item = cJSON_GetObjectItem(json, "keyring_for_untrusted");
    if (item)
    {
        if (item->type == cJSON_Number)
        {
            param->keyring_for_untrusted = item->valueint;
        }
        else if (item->type == cJSON_String)
        {
            param->keyring_for_untrusted = atoi(item->valuestring);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %lu invalid keyring_for_untrusted format", param->rule_id);
        }
    }

    item = cJSON_GetObjectItem(json, "decryption");
    if (item)
    {
        if (item->type == cJSON_Number)
        {
            param->decryption_profile = item->valueint;
        }
        else if (item->type == cJSON_String)
        {
            param->decryption_profile = atoi(item->valuestring);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %lu invalid decryption_profile format", param->rule_id);
        }
    }

    item = cJSON_GetObjectItem(json, "tcp_option_profile");
    if (item)
    {
        if (item->type == cJSON_Number)
        {
            param->tcp_option_profile = item->valueint;
        }
        else if (item->type == cJSON_String)
        {
            param->tcp_option_profile = atoi(item->valuestring);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %lu invalid tcp_option_profile format", param->rule_id);
        }
    }

    *ad = param;
    TFE_LOG_INFO(enforcer->logger, "Add intercept policy: %lu", param->rule_id);

error_out:
    if (json)
    {
        cJSON_Delete(json);
    }
    if (json_str)
    {
        free(json_str);
    }
}

static void intercept_param_free_cb(int table_id, void **ad, long argl, void *argp)
{
    struct intercept_policy_enforcer *enforcer = (struct intercept_policy_enforcer *)argp;
    struct intercept_param *param = (struct intercept_param *)*ad;
    if (param == NULL)
    {
        return;
    }

    if ((__sync_sub_and_fetch(&param->ref_cnt, 1) == 0))
    {
        TFE_LOG_INFO(enforcer->logger, "Del intercept policy %lu", param->rule_id);
        free(param);
        *ad = NULL;
    }
}

static void intercept_param_dup_cb(int table_id, void **to, void **from, long argl, void *argp)
{
    struct intercept_param *param = (struct intercept_param *)*from;
    if (param)
    {
        __sync_add_and_fetch(&(param->ref_cnt), 1);
        *to = param;
    }
    else
    {
        *to = NULL;
    }
}

static void intercept_param_free(struct intercept_param *param)
{
    intercept_param_free_cb(0, (void **)&param, 0, NULL);
}

struct intercept_policy_enforcer *intercept_policy_enforcer_create(void *logger)
{
    int ret = 0;
    struct intercept_policy_enforcer *enforcer = ALLOC(struct intercept_policy_enforcer, 1);
    enforcer->maat = (struct maat *)tfe_bussiness_resouce_get(STATIC_MAAT);
    enforcer->logger = logger;
    enforcer->table_id = maat_get_table_id(enforcer->maat, "PXY_INTERCEPT_COMPILE");

    if (enforcer->table_id < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register table of PXY_INTERCEPT_COMPILE, ret = %d", enforcer->table_id);
        goto error_out;
    }

    ret = maat_plugin_table_ex_schema_register(enforcer->maat,
                                               "PXY_INTERCEPT_COMPILE",
                                               intercept_param_new_cb,
                                               intercept_param_free_cb,
                                               intercept_param_dup_cb,
                                               0,
                                               enforcer);
    if (ret != 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register callback of PXY_INTERCEPT_COMPILE, ret = %d", ret);
        goto error_out;
    }

    return enforcer;

error_out:
    intercept_policy_enforce_destory(enforcer);
    return NULL;
}

void intercept_policy_enforce_destory(struct intercept_policy_enforcer *enforcer)
{
    if (enforcer)
    {
        free(enforcer);
        enforcer = NULL;
    }
}

// return  0 : success
// return -1 : error (need passthrough)
int intercept_policy_enforce(struct intercept_policy_enforcer *enforcer, struct tfe_cmsg *cmsg)
{
    int ret = 0;
    uint16_t size = 0;
    uint64_t rule_id = 0;
    char buff[16] = {0};
    struct intercept_param *param = NULL;
    uint8_t hit_no_intercept = 0;
    int tcp_passthrough = 0;
    char reason_hit_no_intercept[] = "Hit No Intercept";
    char reason_invalid_intercept_param[] = "Invalid Intercept Param";

    ret = tfe_cmsg_get_value(cmsg, TFE_CMSG_POLICY_ID, (unsigned char *)&rule_id, sizeof(rule_id), &size);
    if (ret < 0)
    {
        TFE_LOG_ERROR(g_default_logger, "Failed at fetch intercept rule_id from cmsg: %s", strerror(-ret));
        goto error_passthrough;
    }

    snprintf(buff, sizeof(buff), "%lu", rule_id);
    param = (struct intercept_param *)maat_plugin_table_get_ex_data(enforcer->maat, enforcer->table_id, buff, strlen(buff));
    if (param == NULL)
    {
        TFE_LOG_INFO(enforcer->logger, "Failed to get intercept parameter of policy %lu.", rule_id);
        goto error_passthrough;
    }

    // intercept
    if (param->action == 2)
    {
        tcp_passthrough = 0;
        hit_no_intercept = 0;
    }
    // not intercept
    else
    {
        tcp_passthrough = 1;
        hit_no_intercept = 1;
        tfe_cmsg_set(cmsg, TFE_CMSG_SSL_PASSTHROUGH_REASON, (const unsigned char *)&reason_hit_no_intercept, strlen(reason_hit_no_intercept));
    }

    tfe_cmsg_set(cmsg, TFE_CMSG_POLICY_VSYS_ID, (const unsigned char *)&param->vsys_id, sizeof(param->vsys_id));
    tfe_cmsg_set(cmsg, TFE_CMSG_TCP_PASSTHROUGH, (const unsigned char *)&tcp_passthrough, sizeof(tcp_passthrough));
    tfe_cmsg_set(cmsg, TFE_CMSG_HIT_NO_INTERCEPT, (const unsigned char *)&hit_no_intercept, sizeof(hit_no_intercept));
    tfe_cmsg_set(cmsg, TFE_CMSG_TCP_OPTION_PROFILE_ID, (const unsigned char *)&(param->tcp_option_profile), sizeof(param->tcp_option_profile));
    tfe_cmsg_set(cmsg, TFE_CMSG_DECRYPTION_PROFILE_ID, (const unsigned char *)&(param->decryption_profile), sizeof(param->decryption_profile));
    tfe_cmsg_set(cmsg, TFE_CMSG_KEYRING_FOR_TRUSTED_ID, (const unsigned char *)&(param->keyring_for_trusted), sizeof(param->keyring_for_trusted));
    tfe_cmsg_set(cmsg, TFE_CMSG_KEYRING_FOR_UNTRUSTED, (const unsigned char *)&(param->keyring_for_untrusted), sizeof(param->keyring_for_untrusted));

    intercept_param_free(param);

    return 0;

error_passthrough:
    tcp_passthrough = 1;
    tfe_cmsg_set(cmsg, TFE_CMSG_TCP_PASSTHROUGH, (const unsigned char *)&tcp_passthrough, sizeof(tcp_passthrough));
    tfe_cmsg_set(cmsg, TFE_CMSG_SSL_PASSTHROUGH_REASON, (const unsigned char *)&reason_invalid_intercept_param, strlen(reason_invalid_intercept_param));

    return -1;
}