#include <assert.h>
#include <cjson/cJSON.h>
#include <tfe_stream.h>
#include <tfe_plugin.h>
#include <tfe_proxy.h>
#include "../include/traffic_mirror.h"

extern Maat_feather_t g_business_maat;

struct traffic_mirror_me
{
    struct profile_table_ex_data * profile_ex_data;
    struct traffic_mirror_rebuild * rebuild_ctx;
};

struct traffic_mirror_instance __g_traffic_mirror_instance;
struct traffic_mirror_instance * g_traffic_mirror_instance = &__g_traffic_mirror_instance;

void policy_table_ex_data_free(struct policy_table_ex_data * object)
{
    if ((__sync_sub_and_fetch(&object->atomic_refcnt, 1) == 0)) free(object);
}

void policy_table_ex_data_dup_cb(int table_id, MAAT_PLUGIN_EX_DATA * to,
    MAAT_PLUGIN_EX_DATA * from, long argl, void * argp)
{
    struct policy_table_ex_data * ex_data = (struct policy_table_ex_data *)from;
    __sync_add_and_fetch(&ex_data->atomic_refcnt, 1);
    *to = (void *)ex_data;
}

void policy_table_ex_data_free_cb(int table_id, MAAT_PLUGIN_EX_DATA * ad, long argl, void * argp)
{
    struct policy_table_ex_data * ex_data = (struct policy_table_ex_data *)argp;
    policy_table_ex_data_free(ex_data);
}

void policy_table_ex_data_new_cb(int table_id, const char * key, const char * table_line,
    MAAT_PLUGIN_EX_DATA * ad, long argl, void * argp)
{
    struct traffic_mirror_instance * instance = (struct traffic_mirror_instance *) argp;
    assert(instance != nullptr && instance->logger != nullptr);

    char * str_json = NULL;
    cJSON * json_root = NULL;
    cJSON * json_subroot = NULL;
    cJSON * json_item = NULL;

    struct policy_table_ex_data * ex_data = NULL;

    size_t user_region_offset;
    size_t user_region_len;

    int result = Maat_helper_read_column(table_line, 7, &user_region_offset, &user_region_len);
    if (unlikely(result < 0))
    {
        TFE_LOG_ERROR(instance->logger, "Failed at get policy table's user region.");
        goto ignore;
    }

    str_json = ALLOC(char, user_region_len + 1);
    memcpy(str_json, table_line + user_region_offset, user_region_len);

    json_root = cJSON_Parse(str_json);
    if (unlikely(!json_root))
    {
        TFE_LOG_ERROR(instance->logger, "failed at parsing user region as JSON format.");
        goto ignore;
    }

    json_subroot = cJSON_GetObjectItem(json_root, "decrypt_mirror");
    if (unlikely(!json_subroot))
    {
        TFE_LOG_ERROR(instance->logger, "invalid format, decrypt_mirror is not defined.");
        goto ignore;
    }

    ex_data = ALLOC(struct policy_table_ex_data, 1);
    ex_data->atomic_refcnt = 1;
    ex_data->enable = 0;
    ex_data->profile_id = 0;

    json_item = cJSON_GetObjectItem(json_subroot, "enable");
    if (unlikely(!json_item || cJSON_IsNumber(json_item)))
    {
        TFE_LOG_ERROR(instance->logger, "invalid JSON, decrypt_mirror->enable not existed or invalid type.");
        goto ignore;
    }

    ex_data->enable = json_item->valueint;
    if (!ex_data->enable)
    {
        goto success;
    }

    json_item = cJSON_GetObjectItem(json_subroot, "mirror_profile");
    if (unlikely(!json_item || cJSON_IsNumber(json_item)))
    {
        TFE_LOG_ERROR(instance->logger, "invalid JSON, decrypt_mirror->mirror_profile not existed or invalid type.");
        goto ignore;
    }

success:
    TFE_LOG_DEBUG(instance->logger, "table line in PXY_INTERCEPT_COMPILE added: %s", table_line);
    *ad = ex_data;
    ex_data = nullptr;
    goto out;

ignore:
    TFE_LOG_ERROR(instance->logger, "table line in PXY_INTERCEPT_COMPILE ignored: %s", table_line);
    goto out;

out:
    if (ex_data) policy_table_ex_data_free(ex_data);
    if (json_root) cJSON_Delete(json_root);
    if (str_json) free(str_json);
}

void profile_table_ex_data_free(struct profile_table_ex_data * object)
{
    if ((__sync_sub_and_fetch(&object->atomic_refcnt, 1) == 0)) free(object);
}

void profile_table_ex_data_dup_cb(int table_id, MAAT_PLUGIN_EX_DATA * to,
    MAAT_PLUGIN_EX_DATA * from, long argl, void * argp)
{
    struct profile_table_ex_data * ex_data = (struct profile_table_ex_data *)from;
    __sync_add_and_fetch(&ex_data->atomic_refcnt, 1);
    *to = (void *)ex_data;
}

void profile_table_ex_data_free_cb(int table_id, MAAT_PLUGIN_EX_DATA * ad, long argl, void * argp)
{
    struct profile_table_ex_data * ex_data = (struct profile_table_ex_data *)ad;
    profile_table_ex_data_free(ex_data);
}

void profile_table_ex_data_new_cb(int table_id, const char * key, const char * table_line,
    MAAT_PLUGIN_EX_DATA * ad, long argl, void * argp)
{
    struct traffic_mirror_instance * instance = (struct traffic_mirror_instance *) argp;
    assert(instance != nullptr && instance->logger != nullptr);

    const static struct ether_addr ether_addr_broadcast{0xff,0xff,0xff,0xff, 0xff, 0xff};
    char * str_json = NULL;
    cJSON * json_root = NULL;
    cJSON * json_item = NULL;

    struct profile_table_ex_data * ex_data = NULL;
    size_t addr_list_offset;
    size_t addr_list_len;

    int result = Maat_helper_read_column(table_line, 3, &addr_list_offset, &addr_list_len);
    if (unlikely(result < 0))
    {
        TFE_LOG_ERROR(instance->logger, "Failed at get profile table's addrlist.");
        goto ignore;
    }

    str_json = ALLOC(char, addr_list_len + 1);
    memcpy(str_json, table_line + addr_list_offset, addr_list_len);

    json_root = cJSON_Parse(str_json);
    if (unlikely(!json_root))
    {
        TFE_LOG_ERROR(instance->logger, "failed at parsing addrlist as JSON format.");
        goto ignore;
    }

    ex_data = ALLOC(struct profile_table_ex_data, 1);
    ex_data->atomic_refcnt = 1;

    json_item = cJSON_GetObjectItem(json_root, "vlan");
    if (json_item)
    {
        if (unlikely(!cJSON_IsArray(json_item)))
        {
            TFE_LOG_ERROR(instance->logger, "invalid JSON, mirror_profile->vlan is not a array.");
            goto ignore;
        }

        ex_data->nr_targets = cJSON_GetArraySize(json_item);
        ex_data->vlans = (unsigned int *)calloc(ex_data->nr_targets, sizeof(unsigned int));
        ex_data->ether_addrs = (struct ether_addr *)calloc(ex_data->nr_targets, sizeof(struct ether_addr));

        cJSON * element;
        unsigned int iter = 0;
        cJSON_ArrayForEach(element, json_item)
        {
            if (unlikely(!cJSON_IsNumber(element)))
            {
                TFE_LOG_ERROR(instance->logger, "invalid JSON, "
                                                "elements in mirror_profile->vlan is not a number");
                goto ignore;
            }

            ex_data->vlans[iter] = element->valueint;
            ex_data->ether_addrs[iter] = ether_addr_broadcast;
            iter++;
        }

        assert(iter + 1 == ex_data->nr_targets);
        goto success;
    }

    json_item = cJSON_GetObjectItem(json_item, "mac");
    if (json_item)
    {
        if (unlikely(!cJSON_IsArray(json_item)))
        {
            TFE_LOG_ERROR(instance->logger, "invalid JSON, mirror_profile->mac is not a array.");
            goto ignore;
        }

        ex_data->nr_targets = cJSON_GetArraySize(json_item);
        ex_data->vlans = (unsigned int *)calloc(ex_data->nr_targets, sizeof(unsigned int));
        ex_data->ether_addrs = (struct ether_addr *)calloc(ex_data->nr_targets, sizeof(struct ether_addr));

        cJSON * element;
        unsigned int iter = 0;
        cJSON_ArrayForEach(element, json_item)
        {
            if (unlikely(!cJSON_IsString(element)))
            {
                TFE_LOG_ERROR(instance->logger, "invalid JSON, "
                                                "elements in mirror_profile->mac is not a string");
                goto ignore;
            }

            struct ether_addr ether_addr_aton{};
            if (unlikely(!ether_aton_r(element->valuestring, &ether_addr_aton)))
            {
                TFE_LOG_ERROR(instance->logger, "invalid JSON, "
                                                "elements in mirror_profile->mac is not a valid ether address");
                goto ignore;
            }

            ex_data->ether_addrs[iter] = ether_addr_aton;
            ex_data->vlans[iter] = 0;
            iter++;
        }

        assert(iter + 1 == ex_data->nr_targets);
    }

success:
    *ad = (void *)ex_data;
    ex_data = nullptr;

    TFE_LOG_DEBUG(instance->logger, "table line in PXY_PROFILE_TRAFFIC_MIRROR added: %s", table_line);
    goto out;

ignore:
    TFE_LOG_ERROR(instance->logger, "table line in PXY_PROFILE_TRAFFIC_MIRROR ignored: %s", table_line);
    goto out;

out:
    if (ex_data)
    {
        profile_table_ex_data_free(ex_data);
    }

    if (str_json)
    {
        free(str_json);
    }

    if (json_root)
    {
        cJSON_Delete(json_root);
    }
}

int traffic_mirror_init(struct tfe_proxy * proxy)
{
	int result = 0;
	struct traffic_mirror_instance * instance = g_traffic_mirror_instance;

	/* INIT DECRYPT MIRROR INSTANCE */
	instance->maat_feather = g_business_maat;
	instance->logger = tfe_proxy_get_error_logger();
	instance->nr_threads = tfe_proxy_get_work_thread_count();

	/* REGISTER MAAT FEATHER */
	instance->policy_table_id = Maat_table_register(instance->maat_feather, "PXY_INTERCEPT_COMPILE");
	if (unlikely(instance->policy_table_id < 0))
    {
        TFE_LOG_ERROR(instance->logger, "failed at register table PXY_INTERCEPT_COMPILE, ret = %d",
            instance->policy_table_id); goto errout;
    }

    instance->profile_table_id = Maat_table_register(instance->maat_feather, "PXY_PROFILE_TRAFFIC_MIRROR");
	if (unlikely(instance->profile_table_id < 0))
    {
        TFE_LOG_ERROR(instance->logger, "failed at register table PXY_PROFILE_TRAFFIC_MIRROR, ret = %d",
            instance->profile_table_id); goto errout;
    }

    result = Maat_plugin_EX_register(instance->maat_feather, instance->policy_table_id,
        policy_table_ex_data_new_cb, policy_table_ex_data_free_cb, policy_table_ex_data_dup_cb,
        nullptr, 0, instance);

	if(unlikely(result < 0))
    {
        TFE_LOG_ERROR(instance->logger, "failed at Maat_plugin_EX_register(PXY_INTERCEPT_COMPILE), "
                                        "table_id = %d, ret = %d", instance->policy_table_id, result);
        goto errout;
    }

    result = Maat_plugin_EX_register(instance->maat_feather, instance->policy_table_id,
        profile_table_ex_data_new_cb, profile_table_ex_data_free_cb, policy_table_ex_data_dup_cb,
        nullptr, 0, instance);

	if (unlikely(result < 0))
    {
        TFE_LOG_ERROR(instance->logger, "failed at Maat_plugin_EX_register(PXY_PROFILE_TRAFFIC_MIRROR), "
                                        "table_id = %d, ret = %d", instance->policy_table_id, result);
    }

errout:
	return 0;
}

int traffic_mirror_on_open_cb(const struct tfe_stream * stream, unsigned int thread_id,
    enum tfe_conn_dir dir, void ** pme)
{
    /* Firstly, fetch destination address of traffic mirror */
    struct traffic_mirror_me * me = NULL;
    struct traffic_mirror_instance * instance = g_traffic_mirror_instance;
    struct tfe_cmsg * cmsg = tfe_stream_get0_cmsg(stream);

    assert(instance != NULL);
    assert(cmsg != NULL);

    char str_policy_id[TFE_SYMBOL_MAX];
    char str_profile_id[TFE_SYMBOL_MAX];

    unsigned int opt_val;
    uint16_t opt_out_size;

    struct policy_table_ex_data * policy_ex_data = NULL;
    struct profile_table_ex_data * profile_ex_data = NULL;

    int ret = tfe_cmsg_get_value(cmsg, TFE_CMSG_POLICY_ID, (unsigned char*)&opt_val, sizeof(opt_val), &opt_out_size);
    if (ret < 0)
    {
        TFE_LOG_ERROR(instance->logger, "failed at getting policy id from cmsg, detach the stream.");
        goto detach;
    }

    snprintf(str_policy_id, sizeof(str_policy_id), "%u", opt_val);
    policy_ex_data = (struct policy_table_ex_data *)Maat_plugin_get_EX_data(instance->maat_feather,
        instance->policy_table_id, str_policy_id);

    if (!policy_ex_data)
    {
        TFE_LOG_ERROR(instance->logger, "failed at getting policy %s's EXDATA, detach the stream", str_policy_id);
        goto detach;
    }

    if (!policy_ex_data->enable)
    {
        goto detach;
    }

    snprintf(str_profile_id, sizeof(str_policy_id), "%u", policy_ex_data->profile_id);
    profile_ex_data = (struct profile_table_ex_data *)Maat_plugin_get_EX_data(instance->maat_feather,
        instance->profile_table_id, str_profile_id);

    if (!profile_ex_data)
    {
        TFE_LOG_ERROR(instance->logger, "failed at getting policy %s's profile, profile id = %s, "
                                        "detach the stream", str_policy_id, str_profile_id);
        goto detach;
    }


    me = ALLOC(struct traffic_mirror_me, 1);
    me->rebuild_ctx = traffic_mirror_rebuild_create(stream->addr, profile_ex_data, NULL);
    me->profile_ex_data = profile_ex_data;

    /* profile_ex_data's ownership is transfer to me */
    profile_ex_data = NULL;
    traffic_mirror_rebuild_handshake(me->rebuild_ctx);

    *pme = (void *)me;
    return ACTION_FORWARD_DATA;

detach:
    if (me)
    {
        free(me);
    }

    if (policy_ex_data)
    {
        policy_table_ex_data_free(policy_ex_data);
    }

    if (profile_ex_data)
    {
        profile_table_ex_data_free(profile_ex_data);
    }

    tfe_stream_detach(stream);
	return ACTION_FORWARD_DATA;
}

enum tfe_stream_action traffic_mirror_on_data_cb(const struct tfe_stream * stream, unsigned int thread_id,
    enum tfe_conn_dir dir, const unsigned char * data, size_t len, void ** pme)
{
    struct traffic_mirror_me * me = (struct traffic_mirror_me *)(*pme);
    traffic_mirror_rebuild_data(me->rebuild_ctx, (const char *)data, (size_t)len, dir);
}

void traffic_mirror_on_close_cb(const struct tfe_stream * stream, unsigned int thread_id,
    enum tfe_stream_close_reason reason, void ** pme)
{
    struct traffic_mirror_me * me = (struct traffic_mirror_me *)(*pme);
    traffic_mirror_rebuild_farewell(me->rebuild_ctx);
    traffic_mirror_rebuild_destroy(me->rebuild_ctx);
    profile_table_ex_data_free(me->profile_ex_data);

    free(me);
    *pme = NULL;
}

void traffic_mirror_deinit(struct tfe_proxy * proxy){}

struct tfe_plugin traffic_mirror_plugin_desc =
{
    .symbol= "traffic_mirror",
    .type = TFE_PLUGIN_TYPE_BUSINESS,
    .on_init = traffic_mirror_init,
    .on_deinit = traffic_mirror_deinit,
    .on_open = traffic_mirror_on_open_cb,
    .on_data = traffic_mirror_on_data_cb,
    .on_close = traffic_mirror_on_close_cb
};

TFE_PLUGIN_REGISTER(traffic_mirror, traffic_mirror_plugin_desc)