/*
 * 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 "appid/appid_internal.h"

#include "lpi_plus.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_env
{
	unsigned int max_pkts;
	int lpip_session_exdata_idx;
	int topic_appid;
	struct stellar_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_get0_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->origin=ORIGIN_LPI_PLUS;
	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]);
		result->surrogate_id[i]=0;
	}

	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_mapper *mapper, int appid)
{
	struct lpi_plus_appid_info *out=NULL;
	HASH_FIND(hh_appid, 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 *env=(struct lpi_plus_env *)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(stellar_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 stellar_module_manager *mod_mgr, struct stellar_module *mod)
{
	if(mod_mgr==NULL)return;
	if(mod)
	{
		struct lpi_plus_env *env=(struct lpi_plus_env *)stellar_module_get_ctx(mod);
		lpi_free_library();
        lpi_plus_mapper_free(env->mapper);
		FREE(env);
		stellar_module_free(mod);
	}
}

struct stellar_module *lpi_plus_init(struct stellar_module_manager *mod_mgr)
{
	if(mod_mgr==NULL)return NULL;

	struct lpi_plus_env *env=CALLOC(struct lpi_plus_env, 1);
	struct stellar_module *mod=stellar_module_new("LPI_PLUS", env);
	env->mod_mgr=mod_mgr;
	env->max_pkts=16;//TODO: load from toml

	struct session_manager *sess_mgr=stellar_module_get_session_manager(mod_mgr);
	struct mq_schema *mq_s=stellar_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=stellar_appid_create_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_mapper *stellar_module_get_lpip(struct stellar_module_manager *mod_mgr)
{
	if(mod_mgr==NULL)return NULL;
	struct stellar_module *mod=stellar_module_manager_get_module(mod_mgr, "LPI_PLUS");
	if(mod==NULL)return NULL;
	struct lpi_plus_env *lpi_p=(struct lpi_plus_env *)stellar_module_get_ctx(mod);
	if(lpi_p==NULL)return NULL;
	return lpi_p->mapper;
}