#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;
    uuid_t rule_id;
    int do_log;
    int ref_cnt;
    int action_is_intercept;
    uuid_t keyring_for_trusted;
    uuid_t keyring_for_untrusted;
    uuid_t decryption_profile;
    uuid_t tcp_option_profile;
};

struct intercept_policy_enforcer
{
    struct maat *maat;
    char table_name[32];
    void *logger;
};

static void intercept_param_new_cb(const char *table_name, const char *key, const char *table_line, void **ad, long argl, void *argp)
{
    cJSON *json_root = NULL;
    cJSON *json_subroot = NULL;
    cJSON *item = NULL;
    struct intercept_param *param = NULL;
    struct intercept_policy_enforcer *enforcer = (struct intercept_policy_enforcer *)argp;

    char *json_str = strdup(table_line);
    json_root = cJSON_Parse(json_str);
    if (json_root == NULL)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule: (invlad json format) %s", table_line);
        goto error_out;
    }

    param = ALLOC(struct intercept_param, 1);
    param->ref_cnt = 1;
    uuid_parse(key, param->rule_id);
    uuid_clear(param->keyring_for_trusted);
    uuid_clear(param->keyring_for_untrusted);
    uuid_clear(param->decryption_profile);
    uuid_clear(param->tcp_option_profile);

    // action
    item = cJSON_GetObjectItem(json_root, "action");
    if (!item || !cJSON_IsString(item))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid action format) %s.", key, table_line);
        goto error_out;
    }
    if (strcmp(item->valuestring, "intercept") == 0)
    {
        param->action_is_intercept = 1;
    }
    else if (strcmp(item->valuestring, "no_intercept") == 0)
    {
        param->action_is_intercept = 0;
    }
    else
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid action format) %s.", key, table_line);
        goto error_out;
    }

    // log_option
    item = cJSON_GetObjectItem(json_root, "log_option");
    if (!item || !cJSON_IsString(item))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid log_option format) %s.", key, table_line);
        goto error_out;
    }
    if (0 == strcasecmp(item->valuestring, "none"))
    {
        param->do_log = 0;
    }
    else if (0 == strcasecmp(item->valuestring, "metadata"))
    {
        param->do_log = 1;
    }
    else if (0 == strcasecmp(item->valuestring, "all"))
    {
        param->do_log = 1;
    }
    else
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid log_option format) %s.", key, table_line);
        goto error_out;
    }

    // action_parameter
    json_subroot = cJSON_GetObjectItem(json_root, "action_parameter");
    if (!json_subroot || !cJSON_IsObject(json_subroot))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid action_parameter format) %s.", key, table_line);
        goto error_out;
    }

    // vsys_id
    item = cJSON_GetObjectItem(json_subroot, "vsys_id");
    if (!item || !cJSON_IsNumber(item))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid vsys_id format) %s.", key, table_line);
        goto error_out;
    }
    param->vsys_id = item->valueint;

    // keyring_for_trusted
    item = cJSON_GetObjectItem(json_subroot, "keyring_for_trusted");
    if (item)
    {
        if (cJSON_IsString(item))
        {
            uuid_parse(item->valuestring, param->keyring_for_trusted);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule: %s (invalid keyring_for_trusted format) %s.", key, table_line);
            goto error_out;
        }
    }
    else
    {
        // maybe not exist
    }

    item = cJSON_GetObjectItem(json_subroot, "keyring_for_untrusted");
    if (item)
    {
        if (cJSON_IsString(item))
        {
            uuid_parse(item->valuestring, param->keyring_for_untrusted);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule: %s (invalid keyring_for_untrusted format) %s", key, table_line);
            goto error_out;
        }
    }
    else
    {
        // maybe not exist
    }

    item = cJSON_GetObjectItem(json_subroot, "decryption_profile");
    if (item)
    {
        if (cJSON_IsString(item))
        {
            uuid_parse(item->valuestring, param->decryption_profile);
        }
        else
        {
            TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule: %s (invalid decryption_profile format) %s", key, table_line);
            goto error_out;
        }
    }
    else
    {
        // maybe not exist
    }

    item = cJSON_GetObjectItem(json_subroot, "tcp_option_profile");
    if (!item || !cJSON_IsString(item))
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept rule:%s (invalid tcp_option_profile format) %s.", key, table_line);
        goto error_out;
    }
    uuid_parse(item->valuestring, param->tcp_option_profile);

    *ad = param;
    TFE_LOG_INFO(enforcer->logger, "Add intercept rule: %s", key);

    cJSON_Delete(json_root);
    free(json_str);
    return;

error_out:
    if (json_root)
    {
        cJSON_Delete(json_root);
        json_root = NULL;
    }
    if (json_str)
    {
        free(json_str);
        json_str = NULL;
    }
    if (param)
    {
        free(param);
        param = NULL;
    }
}

static void intercept_param_free_cb(const char *table_name, void **ad, long argl, void *argp)
{
    char str_rule_id[UUID_STRING_SIZE] = {0};
    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))
    {
        uuid_unparse(param->rule_id, str_rule_id);
        TFE_LOG_INFO(enforcer->logger, "Del intercept policy %s", str_rule_id);
        free(param);
        *ad = NULL;
    }
}

static void intercept_param_dup_cb(const char *table_name, 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(NULL, (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 = tfe_get_maat_handle();
    enforcer->logger = logger;
    snprintf(enforcer->table_name, sizeof(enforcer->table_name), "PXY_INTERCEPT_RULE");

    ret = maat_plugin_table_ex_schema_register(enforcer->maat,
                                               enforcer->table_name,
                                               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_RULE, 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_select(struct intercept_policy_enforcer *enforcer, uuid_t *rule_id_array, int rule_id_num, uuid_t *selected_rule_id)
{
    uuid_t *curr_rule_id;
    char str_rule_id[UUID_STRING_SIZE] = {0};
    uint8_t is_hit_intercept_rule = 0;
    uint8_t is_hit_no_intercept_rule = 0;
    uuid_t max_intercept_rule_id;
    uuid_t max_no_intercept_rule_id;

    uuid_clear(*selected_rule_id);
    uuid_clear(max_intercept_rule_id);
    uuid_clear(max_no_intercept_rule_id);

    struct intercept_param *param = NULL;

    for (int i = 0; i < rule_id_num; i++)
    {
        curr_rule_id = &rule_id_array[i];
        uuid_unparse(*curr_rule_id, str_rule_id);
        param = (struct intercept_param *)maat_plugin_table_get_ex_data(enforcer->maat, enforcer->table_name, (const char *)str_rule_id, strlen(str_rule_id));
        if (param == NULL)
        {
            TFE_LOG_INFO(enforcer->logger, "Failed to get intercept parameter of policy %s.", str_rule_id);
            continue;
        }

        // intercept
        if (param->action_is_intercept)
        {
            is_hit_intercept_rule = 1;
            if (uuid_compare(max_intercept_rule_id, *curr_rule_id) < 0)
            {
                uuid_copy(max_intercept_rule_id, *curr_rule_id);
            }
            TFE_LOG_INFO(enforcer->logger, "rule[%d/%d]: %s is intercept.", i, rule_id_num, str_rule_id);
        }
        // not intercept
        else
        {
            is_hit_no_intercept_rule = 1;
            if (uuid_compare(max_no_intercept_rule_id, *curr_rule_id) < 0)
            {
                uuid_copy(max_no_intercept_rule_id, *curr_rule_id);
            }
            TFE_LOG_INFO(enforcer->logger, "rule[%d/%d]: %s is no intercept.", i, rule_id_num, str_rule_id);
        }
    }

    if (is_hit_no_intercept_rule)
    {
        uuid_copy(*selected_rule_id, max_no_intercept_rule_id);
        return 0;
    }

    if (is_hit_intercept_rule)
    {
        uuid_copy(*selected_rule_id, max_intercept_rule_id);
        return 0;
    }

    // no policy get, passthrough
    return -1;
}

// 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;
    uuid_t rule_id;
    char str_rule_id[UUID_STRING_SIZE] = {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(uuid_t), &size);
    if (ret < 0)
    {
        TFE_LOG_ERROR(g_default_logger, "Failed at fetch intercept rule_id from cmsg: %s", strerror(-ret));
        goto error_passthrough;
    }

    uuid_unparse(rule_id, str_rule_id);
    param = (struct intercept_param *)maat_plugin_table_get_ex_data(enforcer->maat, enforcer->table_name, (const char *)str_rule_id, strlen(str_rule_id));
    if (param == NULL)
    {
        TFE_LOG_INFO(enforcer->logger, "Failed to get intercept parameter of policy %s.", str_rule_id);
        goto error_passthrough;
    }

    // intercept
    if (param->action_is_intercept)
    {
        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_DO_LOG, (const unsigned char *)&param->do_log, sizeof(param->do_log));
    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;
}