/*
 * dpkt_basic_proto_main.cpp
 *
 *  Created on: Aug 7, 2020
 *      Author: iie
 */

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <assert.h>

#include "uthash/uthash.h"

#include "stellar/utils.h"
#include "stellar/packet.h"
#include "stellar/session.h"

#include "lpi_plus_internal.h"
#include "lpip_extend.h"


const char *l7_protocol_file = (char *)"./tsgconf/tsg_l7_protocol.conf";

struct lpi_plus_appid_info
{
    int appid;
    char lpi_name[255];
    char app_name[255];
    UT_hash_handle hh_appid;
    UT_hash_handle hh_lpi_name;
};

struct lpi_plus_mapper
{
	struct lpi_plus_appid_info *hash_name2appid;
	struct lpi_plus_appid_info *hash_appid2name;
};

struct lpi_plus
{
	unsigned int max_pkts;
	int lpip_session_exdata_idx;
	int topic_appid;
	struct module_manager *mod_mgr;
	struct lpi_plus_mapper *mapper;
};

struct lpi_plus_exdata
{
	struct lpi_plus_detect_context ctx;
	int appid[MAX_APPID_NUM];
	size_t appid_num;
};

static void lpi_plus_get_host_order_port(struct session *sess __unused, unsigned short *sport, unsigned short *dport)
{
	*sport=0;
	*dport=0;
//get host order port from stellar session api

	const struct packet *pkt = session_get_current_packet(sess);
	enum flow_type flow_type=session_get_flow_type(sess);
	if(pkt && (flow_type==FLOW_TYPE_C2S || flow_type==FLOW_TYPE_S2C))
	{
		int layer_cnt=packet_get_layer_count(pkt);
		const struct layer *layer = packet_get_layer_by_idx(pkt, layer_cnt - 1);
		if (layer && layer->proto == LAYER_PROTO_TCP)
		{
			*sport = ntohs(layer->hdr.tcp->th_sport);
			*dport = ntohs(layer->hdr.tcp->th_dport);
		}
		if (layer && layer->proto == LAYER_PROTO_UDP)
		{
			*sport = ntohs(layer->hdr.udp->uh_sport);
			*dport = ntohs(layer->hdr.udp->uh_dport);
		}
		//S2C, swap sport and dport
        if(flow_type == FLOW_TYPE_S2C)
        {
            unsigned short tmp = *sport;
            *sport = *dport;
            *dport = tmp;
        }
    }
	return;
}

//TODO: fill packet sequence
static struct appid_message *lpi_plus_message_new(struct session *sess, int *id_array, size_t id_num)
{
	struct appid_message *result=CALLOC(struct appid_message, 1);
	result->sess=sess;
	result->appid_num=id_num;
	for(unsigned int i=0; i<result->appid_num; i++)
	{
		result->appid[i]=(int)(id_array[i]);
	}

	return result;
}

static int lpi_plus_appid_update(int current_id_array[], size_t *current_id_num, int incoming_id)
{
	size_t i=0;

	if(*current_id_num>=MAX_APPID_NUM)
	{
		return 0;
	}

	for(i=0; i<(*current_id_num); i++)
	{
		if(incoming_id==current_id_array[i])
		{
			return 0;
		}
	}

	current_id_array[(*current_id_num)]=incoming_id;
	(*current_id_num)++;

	return 1;
}

static int lpi_plus_name2appid(struct lpi_plus_mapper *mapper, const char *lpi_name, size_t name_sz)
{
	struct lpi_plus_appid_info *out=NULL;
	HASH_FIND(hh_lpi_name, mapper->hash_appid2name, lpi_name, name_sz, out);
	if(out==NULL)return 0;
	return out->appid;
}

const char *lpi_plus_appid2name(struct lpi_plus *lpip, int appid)
{
	struct lpi_plus_appid_info *out=NULL;
	HASH_FIND(hh_appid, lpip->mapper->hash_appid2name, &appid, sizeof(int), out);
	if(out==NULL)return NULL;
	return out->app_name;
}

static struct lpi_plus_mapper *lpi_plus_mapper_new(const char *filename)
{
	size_t ret=0;
	FILE *fp=NULL;
	char line[1024]={0};
	char type_name[32]={0};
	struct lpi_plus_appid_info *appid_info=NULL;
	struct lpi_plus_appid_info *tmp=NULL;

	fp=fopen(filename, "r");
	if(fp==NULL)
	{
		printf("Open %s failed ...", filename);
		return NULL;
	}
	struct lpi_plus_mapper * mapper=(struct lpi_plus_mapper *)calloc(1, sizeof(struct lpi_plus_mapper));
	while((fgets(line, sizeof(line), fp))!=NULL)
	{
		if(line[0]=='#' || line[0]=='\n' || line[0]=='\r' ||line[0]=='\0')
		{
			continue;
		}
		
	    appid_info=(struct lpi_plus_appid_info *)calloc(1, sizeof(struct lpi_plus_appid_info));
		ret=sscanf(line, "%31s %63s %d %63s", type_name, appid_info->lpi_name, &appid_info->appid, appid_info->app_name);
		
		if(ret==3)
		{
			strcpy(appid_info->app_name, appid_info->lpi_name);
		}

		HASH_FIND(hh_lpi_name, mapper->hash_name2appid, appid_info->lpi_name, strlen(appid_info->lpi_name), tmp);

		//deduplicate
		if(tmp)
		{
			free(appid_info);
			continue;
		}
	    HASH_ADD_KEYPTR(hh_lpi_name, mapper->hash_name2appid, appid_info->lpi_name, strlen(appid_info->lpi_name), appid_info);
	    HASH_ADD(hh_appid, mapper->hash_appid2name, appid, sizeof(int), appid_info);
		memset(line, 0, sizeof(line));
	}
	fclose(fp);
	return mapper;
}

static  void lpi_plus_mapper_free(struct lpi_plus_mapper *mapper)
{
	if(mapper==NULL)return;

	struct lpi_plus_appid_info *out=NULL, *tmp=NULL;
	HASH_CLEAR(hh_appid, mapper->hash_appid2name);
	HASH_ITER(hh_lpi_name, mapper->hash_name2appid, out ,tmp)
	{		
		HASH_DELETE(hh_lpi_name, mapper->hash_name2appid, out);
		free(out);
	}
	free(mapper);
}

void lpi_plus_context_update(struct session *sess, struct lpi_plus_detect_context *ctx, 
        const char *scan_data, int scan_data_len)
{
	lpi_data_t *data = &ctx->lpi_data;
    int l4_proto = 0;
    enum session_type type = session_get_type(sess);
    if (type == SESSION_TYPE_TCP)
    {
        l4_proto = IPPROTO_TCP;
    }
    if (type == SESSION_TYPE_UDP)
    {
        l4_proto = IPPROTO_UDP;
    }
    int cur_pkt_dir = session_get_flow_type(sess);
    ctx->detected_pkt_cnt++;
    (cur_pkt_dir == FLOW_TYPE_C2S) ? (ctx->detected_c2s_pkt++) : (ctx->detected_s2c_pkt++);
    ctx->current_is_c2s_flow = ((cur_pkt_dir == FLOW_TYPE_C2S) ? 1 : 0);
    uint32_t dir = 0;
    if(ctx->current_is_c2s_flow)
    {
        dir = 0;
    }
    else 
    {
        dir = 1;
	}
    uint32_t four_bytes;

    data->observed[dir] += scan_data_len;

    if (data->trans_proto == 0)data->trans_proto = l4_proto;

	if (scan_data_len < 4)
	{
		memcpy((char *)&four_bytes, scan_data, scan_data_len);
		four_bytes = (ntohl(four_bytes)) >> (8 * (4 - scan_data_len));
		four_bytes = htonl(four_bytes << (8 * (4 - scan_data_len)));
	}
	else
	{
		four_bytes = (*(uint32_t *)scan_data);
	}

	data->payload[dir] = four_bytes;
    data->payload_len[dir] = scan_data_len;

    uint16_t  source=0;
    uint16_t  dest=0;
    lpi_plus_get_host_order_port(sess,&source ,&dest);

    data->client_port = source;
    data->server_port = dest;
    return;

}

static int lpi_plus_detect(struct lpi_plus_detect_context *ctx, struct lpi_plus_mapper *mapper, const char *payload, size_t payload_len)
{
	lpi_module_t *plpi_mod=lpi_guess_protocol(&(ctx->lpi_data));
	if(plpi_mod==NULL)return 0;

    lpi_module_t *extend_result = lpi_plus_extended_guess(ctx, plpi_mod->protocol, payload, payload_len);
    if (extend_result)plpi_mod=extend_result;

    int new_appid=lpi_plus_name2appid(mapper, plpi_mod->name, strlen(plpi_mod->name));

	//if appid won't be outer of tunnel, stop detecting
	if (new_appid>0 
		&& (plpi_mod->category != LPI_CATEGORY_TUNNELLING 
		&& plpi_mod->category != LPI_CATEGORY_NAT 
		&& plpi_mod->protocol != LPI_PROTO_RDP 
		&& plpi_mod->protocol != LPI_PROTO_UDP_RDP 
		&& plpi_mod->protocol != LPI_PROTO_UDP_RTP 
		&& plpi_mod->protocol != LPI_PROTO_UDP_RTCP))
	{
		ctx->stop_detect=1;
	}
	return new_appid;
}

static void lpi_plus_on_session(struct session *sess, enum session_state state, struct packet *pkt, void *args)
{
	if (state == SESSION_STATE_CLOSED)
	{
		assert(pkt == NULL);
		return;
	}
	struct lpi_plus *env=(struct lpi_plus *)args;
	struct lpi_plus_exdata *exdata = (struct lpi_plus_exdata *)session_get_exdata(sess, env->lpip_session_exdata_idx);
	if(exdata==NULL)
	{
		exdata= CALLOC(struct lpi_plus_exdata, 1);
		session_set_exdata(sess, env->lpip_session_exdata_idx, exdata);	
	}
	if(exdata->ctx.stop_detect==1)return;
	if(exdata->ctx.detected_pkt_cnt>=env->max_pkts)return;

	uint16_t payload_len=packet_get_payload_len(pkt);
	const char *payload=packet_get_payload_data(pkt);
	if (payload!=NULL && payload_len>0)//detect packet with payload only
	{
        lpi_plus_context_update(sess, &exdata->ctx, payload, payload_len);
		int appid=lpi_plus_detect(&exdata->ctx, env->mapper, payload, payload_len);
        if(appid>0 && lpi_plus_appid_update(exdata->appid, &(exdata->appid_num), appid))
		{
			struct appid_message *msg=lpi_plus_message_new(sess, exdata->appid, exdata->appid_num);
			if(0 > mq_runtime_publish_message(module_manager_get_mq_runtime(env->mod_mgr),
											 env->topic_appid, 
											 msg))FREE(msg);
		}
	}
	return;
}

static void lpi_plus_exdata_free(int idx __unused, void *ex_ptr, void *arg __unused)
{
	if(ex_ptr==NULL)return;
	FREE(ex_ptr);
}

void lpi_plus_exit(struct module_manager *mod_mgr, struct module *mod)
{
	if(mod_mgr==NULL)return;
	if(mod)
	{
		struct lpi_plus *env=(struct lpi_plus *)module_get_ctx(mod);
		lpi_free_library();
        lpi_plus_mapper_free(env->mapper);
		FREE(env);
		module_free(mod);
	}
}

static void appid_message_free(void *msg, void *msg_free_arg __unused)
{
	if(msg==NULL)return;
	FREE(msg);
}

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wcast-function-type"

static void lpi_plus_appid_on_msg_dispatch(int topic_id __unused, 
									 void *msg, 
									on_msg_cb_func* on_msg_cb, 
									 void *on_msg_cb_arg, 
									 void *dispatch_arg __unused)
{
	on_appid_callback *appid_cb = (on_appid_callback *)on_msg_cb;
	struct appid_message *appid_msg=(struct appid_message *)msg;
	appid_cb(appid_msg->sess, appid_msg->appid, appid_msg->appid_num, on_msg_cb_arg);
}

int lpi_plus_appid_subscribe(struct lpi_plus *lpip, on_appid_callback *cb, void *args)
{
	if(lpip==NULL)return -1;
	struct module_manager *mod_mgr=lpip->mod_mgr;
	int appid_topic_id=mq_schema_get_topic_id(module_manager_get_mq_schema(mod_mgr), LPIP_APPID_MESSAGE_TOPIC);
	if(appid_topic_id<0)
	{
		appid_topic_id=mq_schema_create_topic(module_manager_get_mq_schema(mod_mgr), LPIP_APPID_MESSAGE_TOPIC, lpi_plus_appid_on_msg_dispatch, mod_mgr, appid_message_free, NULL);
	}
	return mq_schema_subscribe(module_manager_get_mq_schema(mod_mgr), appid_topic_id, (on_msg_cb_func *)cb, args);
}

#pragma GCC diagnostic pop

int lpi_plus_create_appid_topic(struct module_manager *mod_mgr)
{
	int app_topic_id=mq_schema_get_topic_id(module_manager_get_mq_schema(mod_mgr), LPIP_APPID_MESSAGE_TOPIC);
	if(app_topic_id < 0)
	{
		app_topic_id=mq_schema_create_topic(module_manager_get_mq_schema(mod_mgr), LPIP_APPID_MESSAGE_TOPIC, lpi_plus_appid_on_msg_dispatch, NULL,appid_message_free, NULL);
	}
	return app_topic_id;
}

struct module *lpi_plus_init(struct module_manager *mod_mgr)
{
	if(mod_mgr==NULL)return NULL;

	struct lpi_plus *env=CALLOC(struct lpi_plus, 1);
	struct module *mod=module_new("LPI_PLUS", env);
	env->mod_mgr=mod_mgr;
	env->max_pkts=16;//TODO: load from toml
	struct module *sess_mgr_mod=module_manager_get_module(mod_mgr, SESSION_MANAGER_MODULE_NAME);
	struct session_manager *sess_mgr=module_to_session_manager(sess_mgr_mod);
	struct mq_schema *mq_s=module_manager_get_mq_schema(mod_mgr);

	if(sess_mgr==NULL || mq_s==NULL)
	{
		goto INIT_ERROR;
	}

	if(lpi_init_library()<0)
	{
		goto INIT_ERROR;
	}
	env->mapper=lpi_plus_mapper_new(l7_protocol_file);// TODO: load path from toml
	if(env->mapper == NULL)
	{
		goto INIT_ERROR;
	}

	session_manager_subscribe_tcp(sess_mgr,lpi_plus_on_session, env);
	session_manager_subscribe_udp(sess_mgr, lpi_plus_on_session, env);

	env->lpip_session_exdata_idx = session_manager_new_session_exdata_index(sess_mgr, "EXDATA_LPI", lpi_plus_exdata_free, NULL);	
	env->topic_appid=lpi_plus_create_appid_topic(mod_mgr);
	if(env->topic_appid<0)
	{
		goto INIT_ERROR;
	}

	return mod;

INIT_ERROR:
	lpi_plus_exit(mod_mgr, mod);
	exit(-1);
	return NULL;

}

struct lpi_plus *module_to_lpi_plus(struct module *mod)
{
	if(mod==NULL)return NULL;
	assert(strcmp(module_get_name(mod), LPI_PLUS_MODULE_NAME) == 0);  
	struct lpi_plus *lpi_p=(struct lpi_plus *)module_get_ctx(mod);
	if(lpi_p==NULL)return NULL;
	return lpi_p;
}