#include <stdio.h>
#include <string.h>
#include <time.h>
#include "mesa_sts.h"
#include "cJSON.h"
#include "stellar/stellar.h"
#include "stellar/session.h"
#include "stellar/stellar_mq.h"
#include "stellar/stellar_exdata.h"
#include "stellar/log.h"

#include "session_flags_internal.h"

#define	_MAX_STR_LEN		64

#define UNUSED(x) (void)(x)

const char *CFG_FILE_PATH="./etc/session_flags/session_flags.toml";

uint32_t sts_random_switch = 0;

struct session_flags_init_conf g_sf_conf;

static const char* random_looking_judge_table[] = {
	"frequency",
	"block_frequency",
	"cumulative_sums",
	"runs",
	"longest_run",
	"rank",
	"non_overlapping_template_matching",
	"overlapping_template_matching",
	"universal",
	"random_excursions",
	"random_excursions_variant",
	"poker_detect",
	"runs_distribution",
	"self_correlation",
	"binary_derivative"
};

static void session_flags_topic_free_cb(void *msg, void *msg_free_arg)
{
	UNUSED(msg_free_arg);
	struct session_flags_message *flags = (struct session_flags_message *)msg;
	
	if (flags)
	{
		if (flags->packet_sequence_array)
		{
			free(flags->packet_sequence_array);
		}
		free(flags);
	}

	return;
}

static void session_flags_exdata_free_cb(int idx, void *ex_ptr, void *arg)
{
	UNUSED(idx);
	UNUSED(arg);

	if (ex_ptr == NULL)
	{
		return;
	}
	struct session_flags_ctx *ctx = (struct session_flags_ctx *)ex_ptr;
	tunneling_hs_stream_free(ctx);

	free(ex_ptr);
}

void session_flags_entry(struct session *session, int topic_id, const void *msg, void *per_session_ctx, void *plugin_env)
{
	UNUSED(per_session_ctx);
	if (msg == NULL)
	{
		return;
	}

	const struct packet *pkt = (const struct packet *)msg;
	struct session_flags_plugin_info *sf_plugin_info = (struct session_flags_plugin_info *)plugin_env;
	struct session_flags_ctx *ctx = (struct session_flags_ctx *)session_exdata_get(session, sf_plugin_info->sess_ctx_exdata_idx);

	if (ctx == NULL)
	{
		ctx = (struct session_flags_ctx *)calloc(1, sizeof(struct session_flags_ctx));
		session_exdata_set(session, sf_plugin_info->sess_ctx_exdata_idx, ctx);

		session_flags_stat_init(&ctx->stat, session_get_direction(session));
		if (g_sf_conf.tunneling_enabled)
		{
			tunneling_hs_stream_init(sf_plugin_info, ctx);
		}
	}

	struct session_flags_stat *stat = &ctx->stat;
	session_flags_result *flags_result = NULL;
	size_t pktlen = 0;

	pktlen = packet_get_raw_len(pkt);

	uint64_t curr_time_ms = 0;
	const struct timeval *tv;

	tv = packet_get_timeval(pkt);
	if (tv != NULL)
	{
		curr_time_ms = tv->tv_sec * 1000 + tv->tv_usec / 1000;
	}

	flags_result = session_flags(sf_plugin_info, ctx, session, topic_id, pktlen, session_get_flow_type(session), curr_time_ms);

	if ((stat->stream_live_time_ms > g_sf_conf.session_max_process_time_ms))
	{
		stellar_session_plugin_dettach_current_session(session);
	}

	uint32_t pkts_num = stat->c2s.pkts + stat->s2c.pkts;

	if (flags_result != NULL)
	{
		if (flags_result->flags != ctx->history_flags)
		{	
			struct session_flags_message *flags_msg = session_flags_generate_firewall_message(flags_result->flags, flags_result->identify);
			if (session_mq_publish_message(session, sf_plugin_info->session_flags_topic_id, flags_msg) < 0)
			{
				session_flags_topic_free_cb(flags_msg, NULL);

				STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "[%s] publish session flags message failed, flags:0x%x, pkts_num:%d", session_get0_readable_addr(session), flags_result->flags, pkts_num);
				return;
			}

			ctx->history_flags = flags_result->flags;

			STELLAR_LOG_DEBUG(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, 
							"[%s] flags:0x%x, pkts_num:%d\n \
							c2s_iter: bulky: %f, CBR: %f, download: %f, interactive: %f, pseudo_unidirectional: %f, streaming: %f, CBR_CV: %f\n \
							s2c_iter: bulky: %f, CBR: %f, download: %f, interactive: %f, pseudo_unidirectional: %f, streaming: %f, CBR_CV: %f\n \
							iter_cnt: %d, bidirectional: %f\n", session_get0_readable_addr(session), ctx->history_flags, pkts_num,
							stat->iter.c2s.bulky, stat->iter.c2s.CBR, stat->iter.c2s.download, stat->iter.c2s.interactive, stat->iter.c2s.pseudo_unidirectional, stat->iter.c2s.streaming, session_flags_calculate_CV(&stat->iter.c2s.omean),
							stat->iter.s2c.bulky, stat->iter.s2c.CBR, stat->iter.s2c.download, stat->iter.s2c.interactive, stat->iter.s2c.pseudo_unidirectional, stat->iter.s2c.streaming, session_flags_calculate_CV(&stat->iter.s2c.omean),
							stat->iter.iter_cnt, stat->iter.bidirectional);
		}
	}

	return;
}

static void session_flags_load_config(struct session_flags_plugin_info *sf_plugin_info, const char *cfg_file, struct session_flags_init_conf *g_sf_conf)
{
	char errbuf[256] = {0};

	FILE *fp = fopen(cfg_file, "r");
    if (NULL == fp)
    {
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "Can't open config file:%s", cfg_file);
        return;
    }

    toml_table_t *toml_root = toml_parse_file(fp, errbuf, sizeof(errbuf));
    fclose(fp);

    toml_table_t *cfg_tbl = toml_table_in(toml_root, "SESSION_FLAGS");
    if (NULL == cfg_tbl)
    {
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "config file:%s has no key: [SESSION_FLAGS]", cfg_file);
        toml_free(toml_root);
        return;
    }

	toml_datum_t toml_val = toml_int_in(cfg_tbl, "INTERACTIVE_STARTTIME_MS");
	if (toml_val.ok)
	{
		g_sf_conf->interactive_starttime_ms = toml_val.u.i;
	}
	else
	{
		g_sf_conf->interactive_starttime_ms = 15000;
	}

	toml_val = toml_int_in(cfg_tbl, "INTERACTIVE_PULSE_NUM");
	if (toml_val.ok)
	{
		g_sf_conf->interactive_pulse_num = toml_val.u.i;
	}
	else
	{
		g_sf_conf->interactive_pulse_num = 4;
	}

	toml_val = toml_int_in(cfg_tbl, "INTERACTIVE_LATENCY_MS");
	if (toml_val.ok)
	{
		g_sf_conf->interactive_latency_ms = toml_val.u.i;
	}
	else
	{
		g_sf_conf->interactive_latency_ms = 1000;
	}

	toml_val = toml_int_in(cfg_tbl, "MAIN_DIR_FRONT_N_PKTS");
	if (toml_val.ok)
	{
		g_sf_conf->main_dir_front_n_pkts = toml_val.u.i;
	}
	else
	{
		g_sf_conf->main_dir_front_n_pkts = 100;
	}

	toml_val = toml_int_in(cfg_tbl, "LARGE_PKTS_INIT_SIZE");
	if (toml_val.ok)
	{
		g_sf_conf->large_ptks_init_size = toml_val.u.i;
	}
	else
	{
		g_sf_conf->large_ptks_init_size = 1000;
	}

	toml_val = toml_int_in(cfg_tbl, "SESSION_MAX_PROCESS_TIME_MS");
	if (toml_val.ok)
	{
		g_sf_conf->session_max_process_time_ms = toml_val.u.i;
	}
	else
	{
		g_sf_conf->session_max_process_time_ms = 30000;
	}

	toml_val = toml_int_in(cfg_tbl, "FET_ENABLED");
	if (toml_val.ok)
	{
		g_sf_conf->fet_enabled = toml_val.u.b;
	}
	else
	{
		g_sf_conf->fet_enabled = 1;
	}

	toml_val = toml_int_in(cfg_tbl, "RANDOM_LOOKING_UDP_IGNORE_PKTS");
	if (toml_val.ok)
	{
		g_sf_conf->random_looking_udp_ignore_pkts = toml_val.u.i;
	}
	else
	{
		g_sf_conf->random_looking_udp_ignore_pkts = 3;
	}

	toml_val = toml_int_in(cfg_tbl, "TUNNELING_TLS_IGNORE_PKTS");
	if (toml_val.ok)
	{
		g_sf_conf->tunneling_tls_ignore_pkts = toml_val.u.i;
	}
	else
	{
		g_sf_conf->tunneling_tls_ignore_pkts = 4;
	}

	toml_val = toml_int_in(cfg_tbl, "TUNNELING_MAX_SCAN_PKTS");
	if (toml_val.ok)
	{
		g_sf_conf->tunneling_max_scan_pkts = toml_val.u.i;
	}
	else
	{
		g_sf_conf->tunneling_max_scan_pkts = 15;
	}

	toml_val = toml_string_in(cfg_tbl, "TUNNELING_PCRE_LIST");
	if (toml_val.ok)
	{
		strncpy(g_sf_conf->tunneling_pcre_list, toml_val.u.s, sizeof(g_sf_conf->tunneling_pcre_list) - 1);
		free(toml_val.u.s);
		toml_val.u.s = NULL;
	}
	else
	{
		strncpy(g_sf_conf->tunneling_pcre_list, "{\"tunneling_pcre_list\":[]}", sizeof(g_sf_conf->tunneling_pcre_list) - 1);
	}

	toml_val = toml_string_in(cfg_tbl, "RANDOM_LOOKING_JUDGE_LIST");
	if (toml_val.ok)
	{
		strncpy(g_sf_conf->random_looking_judge_list, toml_val.u.s, sizeof(g_sf_conf->random_looking_judge_list) - 1);
		free(toml_val.u.s);
		toml_val.u.s = NULL;
	}
	else
	{
		strncpy(g_sf_conf->random_looking_judge_list, "{\"random_looking_judge_list\":[]}", sizeof(g_sf_conf->random_looking_judge_list) - 1);
	}

	toml_free(toml_root);
}

extern "C" void *session_flags_plugin_init(struct stellar *st)
{
	char random_looking_list_str[2048] = {0};
	struct session_flags_plugin_info *sf_plugin_info = (struct session_flags_plugin_info *)calloc(1, sizeof(struct session_flags_plugin_info));
	cJSON *json = NULL, *item = NULL;
	int array_num;

	sf_plugin_info->st = st;
	sf_plugin_info->sess_ctx_exdata_idx = stellar_exdata_new_index(st, "SESSION_FLAGS_SESS_CTX", session_flags_exdata_free_cb, NULL);

	sf_plugin_info->log_handle = stellar_get_logger(st);
	if(sf_plugin_info->log_handle==NULL)
	{
		printf("stellar_get_logger object failed ...\n");
		goto ERROR;
	}

	memset(&g_sf_conf, 0, sizeof(g_sf_conf));
	session_flags_load_config(sf_plugin_info, CFG_FILE_PATH, &g_sf_conf);
	tunneling_hyperscan_engine_init(sf_plugin_info, &g_sf_conf);

	json = cJSON_Parse(g_sf_conf.random_looking_judge_list);
	if (json == NULL)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "cJSON_Parse failed, random_looking_list_str:%s", random_looking_list_str);
		goto ERROR;
	}
	item = cJSON_GetObjectItem(json, "random_looking_judge_list");
	if (item == NULL)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "cJSON_GetObjectItem failed, random_looking_list_str:%s", random_looking_list_str);
		goto ERROR;
	}
	array_num = cJSON_GetArraySize(item);
	if (array_num > STS_RANDOM_JUDGE_NUM || array_num < 0)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "array size error, array_num:%d", array_num);
		goto ERROR;
	}
	for (int i = 0; i < array_num; i++)
	{
		for (int j = 0; j < STS_RANDOM_JUDGE_NUM; j++)
		{
			if (strcmp(cJSON_GetArrayItem(item, i)->valuestring, random_looking_judge_table[j]) == 0)
			{
				STS_SET_FLAG(sts_random_switch, j);
				g_sf_conf.random_judge_flags_cnt++;
				break;
			}
		}
	}
	cJSON_Delete(json);

	sf_plugin_info->plugin_id = stellar_session_plugin_register(st, NULL, NULL, sf_plugin_info);
	if (sf_plugin_info->plugin_id < 0)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "stellar_session_plugin_register failed");
		goto ERROR;
	
	}

	sf_plugin_info->session_flags_topic_id = stellar_mq_create_topic(st, SESSION_FLAGS_MESSAGE_TOPIC, session_flags_topic_free_cb, NULL);
	if (sf_plugin_info->session_flags_topic_id < 0)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "stellar_session_mq_create_topic failed");
		goto ERROR;
	}

	sf_plugin_info->tcp_topic_id = stellar_mq_get_topic_id(st, TOPIC_TCP_INPUT);
	sf_plugin_info->udp_topic_id = stellar_mq_get_topic_id(st, TOPIC_UDP_INPUT);
	if (sf_plugin_info->tcp_topic_id < 0 || sf_plugin_info->udp_topic_id < 0)
	{
		STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "stellar_session_mq_get_topic_id failed");
		goto ERROR;
	}

	stellar_session_mq_subscribe(st, sf_plugin_info->tcp_topic_id, session_flags_entry, sf_plugin_info->plugin_id);
	stellar_session_mq_subscribe(st, sf_plugin_info->udp_topic_id, session_flags_entry, sf_plugin_info->plugin_id);

	return sf_plugin_info;

ERROR:
	if (sf_plugin_info != NULL)
	{
		free(sf_plugin_info);
	}

	if (json != NULL)
	{
		cJSON_Delete(json);
	}

	perror("session_flags init failed");
	exit(-1);
}

extern "C" void session_flags_plugin_exit(void *plugin_ctx)
{
	if (plugin_ctx == NULL)
	{
		return;
	}

	struct session_flags_plugin_info *sf_plugin_info = (struct session_flags_plugin_info *)plugin_ctx;
	
	tunneling_hyperscan_engine_exit(sf_plugin_info->tunneling_hs_db);

	free(plugin_ctx);

	return;
}