#include <cctype>
#include <stdlib.h>
#include <stdio.h>
#include "cJSON.h"
#include "session_flags_internal.h"
#include "tunneling.h"
#include "stellar/log.h"

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

thread_local hs_scratch_t *hs_scratch = NULL;
extern struct session_flags_init_conf g_sf_conf;

static char tunneling_length_to_character(enum flow_type flow_type, size_t len)
{
    char ret;

    switch(len)
    {
        case 1 ... 200:
            ret = 'A';
            break;
        case 201 ... 600:
            ret = 'B';
            break;
        case 601 ... 1000:
            ret = 'C';
            break;
        case 1001 ... 1460:
            ret = 'D';
            break;
        default:
            ret = 'Z';
            break;
    }

    if (flow_type == FLOW_TYPE_C2S)
    {
        return ret;
    }
    else
    {
        return tolower(ret);
    }
}

static int tunneling_match_event_handler(unsigned int id, unsigned long long from, unsigned long long to, unsigned int flags, void *context) {
    UNUSED(id);
    UNUSED(from);
    UNUSED(to);
    UNUSED(flags);

    struct session_flags_ctx *ctx = (struct session_flags_ctx *)context;
	ctx->stat.result.flags |= SESSION_FLAGS_TUNNELING;
    return 0;
}

int tunneling_scan_sequence(struct session_flags_plugin_info *sf_plugin_info, struct session *session, struct session_flags_ctx *ctx, size_t payload_len, enum flow_type flow_type, uint64_t pkts_cnt)
{
    if (ctx->stat.result.flags & SESSION_FLAGS_TUNNELING)
    {
        return 0;
    }

    ctx->stat.tunneling_stat.payload_pkt_num++;
    if (ctx->stat.result.is_tls && ctx->stat.tunneling_stat.payload_pkt_num <= g_sf_conf.tunneling_tls_ignore_pkts)
    {
        return 0;
    }

    if((ctx->stat.result.is_tls==0) && (ctx->stat.tunneling_stat.payload_pkt_num > g_sf_conf.tunneling_max_scan_pkts))
    {
        return 0;
    }

    if((ctx->stat.result.is_tls) && (ctx->stat.tunneling_stat.payload_pkt_num > g_sf_conf.tunneling_max_scan_pkts+g_sf_conf.tunneling_tls_ignore_pkts))
    {
        return 0;
    }

    if (hs_scratch == NULL)
    {
        hs_error_t err = hs_alloc_scratch(sf_plugin_info->tunneling_hs_db, &hs_scratch);
        if (err != HS_SUCCESS)
        {
            STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "hs_alloc_scratch failed, err:%d", err);
            return -1;
        }
    }

    char tunneling_seq_char = tunneling_length_to_character(flow_type, payload_len);
    STELLAR_LOG_DEBUG(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "session: %s, is tls:%s, total_num: %d, payload_pkt_num: %d, tunneling_seq_char:%c, payload_len:%d", 
                            session_get0_readable_addr(session), ctx->stat.result.is_tls == true ? "yes":"no", pkts_cnt, ctx->stat.tunneling_stat.payload_pkt_num, tunneling_seq_char, payload_len);

    hs_error_t err = hs_scan_stream(ctx->tunneling_hs_stream, &tunneling_seq_char, 1, 0, hs_scratch, tunneling_match_event_handler, ctx);
    if (err != HS_SUCCESS)
    {
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "hs_scan_stream failed, err:%d", err);
        return -1;
    }
    if (ctx->stat.result.flags & SESSION_FLAGS_TUNNELING)
    {
        ctx->stat.result.identify[session_flags_tunneling_mask] = pkts_cnt;
    }

    return 0;
}

void tunneling_hs_stream_init(struct session_flags_plugin_info *sf_plugin_info, struct session_flags_ctx *ctx)
{
    hs_error_t err = hs_open_stream(sf_plugin_info->tunneling_hs_db, 0, &ctx->tunneling_hs_stream);
    if (err != HS_SUCCESS)
    {
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "hs_open_stream failed, err:%d", err);
        return;
    }
}

void tunneling_hs_stream_free(struct session_flags_ctx *ctx)
{
    if (ctx->tunneling_hs_stream == NULL)
    {
        return;
    }

    hs_close_stream(ctx->tunneling_hs_stream, hs_scratch, NULL, NULL);
}

int tunneling_hyperscan_engine_init(struct session_flags_plugin_info *sf_plugin_info, struct session_flags_init_conf *g_sf_conf)
{
    cJSON *json = NULL, *item = NULL;
    int array_num;
    char **pcre = NULL;
    hs_compile_error_t *compile_err;
    hs_error_t err;
    unsigned int *flags = NULL;
    unsigned int *ids = NULL;
    int ret = 0;

    g_sf_conf->tunneling_enabled = 0;

    json = cJSON_Parse(g_sf_conf->tunneling_pcre_list);
	if (json == NULL)
	{
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "cJSON_Parse failed, tunneling_pcre_list:%s", g_sf_conf->tunneling_pcre_list);
		goto END;
	}
	item = cJSON_GetObjectItem(json, "tunneling_pcre_list");
	if (item == NULL)
	{
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "cJSON_GetObjectItem failed, tunneling_pcre_list:%s", g_sf_conf->tunneling_pcre_list);
		goto END;
	}
	array_num = cJSON_GetArraySize(item);
	if (array_num < 0)
	{
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "array size error, array_num:%d", array_num);
		goto END;
	}

    if (array_num == 0)
    {
        goto END;
    }

    g_sf_conf->tunneling_enabled = 1;
    pcre = (char **)calloc(array_num, sizeof(char *));
	for (int i = 0; i < array_num; i++)
	{
		pcre[i] = cJSON_GetArrayItem(item, i)->valuestring;
    }

    flags = (unsigned int *)calloc(array_num, sizeof(unsigned int));
    ids = (unsigned int *)calloc(array_num, sizeof(unsigned int));
    for (int i = 0; i < array_num; i++)
    {
        flags[i] = HS_FLAG_DOTALL;
        ids[i] = i;
    }

    err = hs_compile_multi(pcre, flags, ids, array_num, HS_MODE_STREAM, NULL, &sf_plugin_info->tunneling_hs_db, &compile_err);
    if (err != HS_SUCCESS)
    {
        STELLAR_LOG_FATAL(sf_plugin_info->log_handle, SESSION_FLAGS_LOG_MODULE, "hs_compile_multi failed, err:%d, pattern id: %d, err_msg: %s, pattern: %s", err, compile_err->expression, compile_err->message, pcre[compile_err->expression]);
        cJSON_Delete(json);
        free(pcre);
        ret = -1;
        goto END;
    }

END:
    if (json != NULL)
    {
        cJSON_Delete(json);
    }
    if (pcre != NULL)
    {
        free(pcre);
    }
    if (flags != NULL)
    {
        free(flags);
    }
    if (ids != NULL)
    {   
        free(ids);
    }
    return ret;
}

void tunneling_hyperscan_engine_exit(hs_database_t *tunneling_hs_db)
{
    if (tunneling_hs_db != NULL)
    {
        hs_free_database(tunneling_hs_db);
    }
}