#include "kni_dynamic_bypass.h"

#define SSL_INFO_LEN 512

struct dynamic_bypass_ssl_feature{
	size_t vlen;
	char value[SSL_INFO_LEN];
};

extern struct kni_handle *g_kni_handle;
extern struct kni_field_stat_handle *g_kni_fs_handle;

char * stream_addr_to_str(struct pme_info *pmeinfo,struct pkt_info *pktinfo)
{
	char * __str_ret = NULL;
	if(pmeinfo==NULL || pktinfo == NULL) return NULL;
	if (pmeinfo->addr_type == ADDR_TYPE_IPV4)
	{
		char __src_addr[INET_ADDRSTRLEN];
		char __dst_addr[INET_ADDRSTRLEN];
		uint16_t __src_port;
		uint16_t __dst_port;
		if(pktinfo->iphdr.v4->saddr < pktinfo->iphdr.v4->daddr){
			__src_port = ntohs((uint16_t) pmeinfo->stream->addr.tuple4_v4->source);
			__dst_port = ntohs((uint16_t) pmeinfo->stream->addr.tuple4_v4->dest);
			inet_ntop(AF_INET, &(pmeinfo->stream->addr.tuple4_v4->saddr), __src_addr, sizeof(__src_addr));
			inet_ntop(AF_INET, &(pmeinfo->stream->addr.tuple4_v4->daddr), __dst_addr, sizeof(__dst_addr));
		}
		else
		{
			__src_port = ntohs((uint16_t) pmeinfo->stream->addr.tuple4_v4->dest);
			__dst_port = ntohs((uint16_t) pmeinfo->stream->addr.tuple4_v4->source);
			inet_ntop(AF_INET, &(pmeinfo->stream->addr.tuple4_v4->daddr), __src_addr, sizeof(__src_addr));
			inet_ntop(AF_INET, &(pmeinfo->stream->addr.tuple4_v4->saddr), __dst_addr, sizeof(__dst_addr));
		}
		asprintf(&__str_ret, "%s %u %s %u", __src_addr, __src_port, __dst_addr, __dst_port);
	}

	if (pmeinfo->addr_type == ADDR_TYPE_IPV6)
	{
		char __src_addr[INET6_ADDRSTRLEN];
		char __dst_addr[INET6_ADDRSTRLEN];
		uint16_t __src_port;
		uint16_t __dst_port;
		if(memcmp((void*)&(pktinfo->iphdr.v6->ip6_src), (void*)&(pktinfo->iphdr.v6->ip6_dst), IPV6_ADDR_LEN) < 0){
			__src_port = ntohs((uint16_t) pktinfo->tcphdr->source);
			__dst_port = ntohs((uint16_t) pktinfo->tcphdr->dest);
			memcpy(__src_addr, &(pktinfo->iphdr.v6->ip6_src), sizeof(__src_addr));
			memcpy(__dst_addr, &(pktinfo->iphdr.v6->ip6_dst), sizeof(__dst_addr));
		}
		else
		{
			__src_port = ntohs((uint16_t) pktinfo->tcphdr->dest);
			__dst_port = ntohs((uint16_t) pktinfo->tcphdr->source);
			memcpy(__src_addr, &(pktinfo->iphdr.v6->ip6_dst), sizeof(__src_addr));
			memcpy(__dst_addr, &(pktinfo->iphdr.v6->ip6_src), sizeof(__dst_addr));
		}
		asprintf(&__str_ret, "%s %u %s %u", __src_addr, __src_port, __dst_addr, __dst_port);
	}

	return __str_ret;
}

int stream_addr_str_split(char* addr_str, const char** sip, const char** sport, const char** dip, const char** dport)
{
	const char* seps=" ";
	char* saveptr=NULL, *subtoken=NULL, *str=NULL;
	int i=0;
	for (str = addr_str, i=0; ; str = NULL, i++)
	{
		subtoken = strtok_r(str, seps, &saveptr);
		if (subtoken == NULL)
			break;
		switch(i)
		{
			case 0:
				if(sip!=NULL) *sip=subtoken;
				break;
			case 1:
				if(sport!=NULL) *sport=subtoken;
				break;
			case 2:
				if(dip!=NULL) *dip=subtoken;
				break;
			case 3:
				if(dport!=NULL)	*dport=subtoken;
				break;
			default:
				return -1;
				break;
		}
	}
	return 0;

}

static long sslinfo2bypass_htable_search_cb(void *data, const uchar *key, uint size, void *user_args){
	if(data != NULL)
		return 0;
	else
		return -1;
}

static int sslinfo2bypass_htable_search(MESA_htable_handle htable,char *key, int klen,  struct pme_info *pmeinfo)
{
	void *logger = g_kni_handle->local_logger;
    long cb_ret = -1;
	void *value;
    value = MESA_htable_search_cb(htable, (const unsigned char *)key, (unsigned int )klen, sslinfo2bypass_htable_search_cb, NULL, &cb_ret);
	if(cb_ret == 0)
	{
		pmeinfo->ssl_pinningst = *((uint64_t *)value);
		KNI_LOG_DEBUG(logger, "MESA_htable: success to search, table = sslinfo2bypass_htable,key = %s, key_size = %d, ret =%d", key, klen,cb_ret);
	}
	else
	{
		KNI_LOG_DEBUG(logger, "MESA_htable:Search result value is null, table = sslinfo2bypass_htable,key = %s, key_size = %d, ret =%d", key, klen,cb_ret);
	}
	return cb_ret;
}

static int traceid2sslinfo_htable_add(char *key, uint klen, void * value,int thread_seq)
{
	void *logger = g_kni_handle->local_logger;
	int ret;
	MESA_htable_handle traceid2sslinfo_htable = g_kni_handle->threads_handle[thread_seq].traceid2sslinfo_htable;
	ret = MESA_htable_add(traceid2sslinfo_htable, (const unsigned char *)key, klen, (const void*)value);
	if(ret >= 0)
	{
		FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_ADD_SUCC], 0, FS_OP_ADD, 1);
		FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_CNT], 0, FS_OP_ADD, 1);
		KNI_LOG_DEBUG(logger, "MESA_htable:Success to add, table = traceid2sslinfo_htable,key = %s, key_size = %d, ret =%d , thread_seq=%d", key, klen,ret, thread_seq);
	}
	else
	{
		KNI_LOG_ERROR(logger, "MESA_htable:Fail to add, table = traceid2sslinfo_htable,key = %s, key_size = %d, ret =%d , thread_seq=%d", key, klen,ret, thread_seq);
	}
	return ret;
}


static int sslinfo2bypass_htable_add(char *key, uint klen, void * value, void *logger)
{
	int ret = -1;
	MESA_htable_handle sslinfo2bypass_htable = g_kni_handle->sslinfo2bypass_htable;
	ret = MESA_htable_add(sslinfo2bypass_htable, (const unsigned char *)key, klen, (const void*)value);
	if(ret < 0)
	{
		KNI_LOG_ERROR(logger, "MESA_htable: Failed at add, table = sslinfo2bypass_htable, key = %s, key_size = %d, ret = %d",key, klen, ret);
	}
	else
	{
		KNI_LOG_DEBUG(logger, "MESA_htable: Success at add, table = sslinfo2bypass_htable, key = %s, key_size = %d, ret = %d",key, klen, ret);
		FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_SSL2PASS_ADD_SUCC], 0, FS_OP_ADD, 1);
		FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_SSL2PASS_CNT], 0, FS_OP_ADD, 1);

	}
	return ret;
}

static long traceid2sslinfo_htable_search_cb(void *data, const uchar *key, uint size, void *user_args)
{
	struct dynamic_bypass_ssl_feature * ssl_feature = (struct dynamic_bypass_ssl_feature *)data;
	void *logger = g_kni_handle->local_logger;
	uint64_t * p_ssl_pinningst = (uint64_t *)user_args;
	uint64_t *value = NULL;
	long cb_ret = -1;

	if(ssl_feature != NULL)
	{
		value = ALLOC(uint64_t, 1);
		*value = *p_ssl_pinningst;
		if(sslinfo2bypass_htable_add(ssl_feature->value, ssl_feature->vlen, (void *)value, logger) < 0)
		{
			KNI_LOG_ERROR(logger, "Dynamic bypass, fail add to table = sslinfo2bypass_htable, key = %s, key_size = %d",
						ssl_feature->value, ssl_feature->vlen);
		}
		cb_ret = 0;
	} 
	return cb_ret;
}


static int traceid2sslinfo_htable_search(MESA_htable_handle htable,char *key, int klen, int thread_seq,uint64_t ssl_pinningst)
{
    long cb_ret = -1;
	void *logger = g_kni_handle->local_logger;
    MESA_htable_search_cb(htable, (const unsigned char *)key, (unsigned int )klen, traceid2sslinfo_htable_search_cb, &ssl_pinningst, &cb_ret);
	if(cb_ret >= 0)
	{
		KNI_LOG_DEBUG(logger, "MESA_htable: success to search, table = traceid2sslinfo_htable,thread_seq = %d,key = %s, key_size = %d, ret = %d", 
							thread_seq, key, klen,cb_ret);
	}
	else
	{
		KNI_LOG_ERROR(logger, "MESA_htable: Fail to search, table = traceid2sslinfo_htable,thread_seq = %d,key = %s, key_size = %d, ret = %d", 
							thread_seq, key, klen,cb_ret);
	}
	
	return cb_ret;
}

static void traceid2sslinfo_htable_data_free_cb(void *data){
	FREE(&data);
}

static void sslinfo2bypass_htable_data_free_cb(void *data){
	FREE(&data);
}

static int traceid2sslinfo_htable_expire_notify_cb(void *data, int eliminate_type)
{
	FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_DEL_SUCC], 0, FS_OP_ADD, 1);
	FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_CNT], 0, FS_OP_ADD, -1);
	return 1;
}

static int sslinfo2bypass_htable_expire_notify_cb(void *data, int eliminate_type)
{
	FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_SSL2PASS_DEL_SUCC], 0, FS_OP_ADD, 1);
	FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_SSL2PASS_CNT], 0, FS_OP_ADD, -1);
	return 1;
}



static int get_ssl_conn_info(struct pme_info *pmeinfo,struct pkt_info *pktinfo, struct dynamic_bypass_ssl_feature *ssl_feature)
{
	void *logger = g_kni_handle->local_logger;

	const char *sip=NULL, *sport=NULL, *dip=NULL, *dport=NULL;
	char *addr_str = NULL;

	if(pmeinfo->session_attribute->ja3_fingerprint == NULL)
	{
		KNI_LOG_DEBUG(logger, "Dynamic bypass:get stream label ja3_fingerprint is null, stream traceid = %s", pmeinfo->stream_traceid);
		return 1;
	}

	addr_str = stream_addr_to_str(pmeinfo,pktinfo);
	stream_addr_str_split(addr_str, &sip, &sport, &dip, &dport);


	ssl_feature->vlen = snprintf(ssl_feature->value, SSL_INFO_LEN, "%s:%s:%s",
																	pmeinfo->session_attribute->ja3_fingerprint,sip,
																	pmeinfo->domain_len > 0 ? (char*)&(pmeinfo->domain): dip);

	FREE(&addr_str);
	return 0;
}


int first_data_ssl_dynamic_bypass(struct streaminfo *stream, struct pme_info *pmeinfo, struct pkt_info *pktinfo, int thread_seq)
{
	struct dynamic_bypass_ssl_feature *ssl_feature = NULL;
	void *logger = g_kni_handle->local_logger;

	int ret = -1;
	if(pmeinfo->protocol != PROTO_SSL){
		KNI_LOG_DEBUG(logger, "Dynamic bypass: stream protocol is not ssl: stream traceid = %s, protocol = %d", pmeinfo->stream_traceid,pmeinfo->protocol);
		return ret;
	}

	ssl_feature = ALLOC(struct dynamic_bypass_ssl_feature,1);
 	if(get_ssl_conn_info(pmeinfo,pktinfo, ssl_feature) == 1){
		FREE(&ssl_feature);
	 	return ret;
	}
	
	KNI_LOG_DEBUG(logger, "Dynamic bypass: ssl_info: %s, ssl_info len:%d, stream traceid = %s", ssl_feature->value, ssl_feature->vlen, pmeinfo->stream_traceid);

	if(sslinfo2bypass_htable_search(g_kni_handle->sslinfo2bypass_htable,ssl_feature->value, ssl_feature->vlen,pmeinfo) == 0)
	{
		KNI_LOG_DEBUG(logger, "Dynamic bypass: passthrough ok, stream traceid = %s,ssl_pinningst = %d", pmeinfo->stream_traceid, pmeinfo->ssl_pinningst);
		FREE(&ssl_feature);
		ret = 0; 
	}
	else
	{
		KNI_LOG_DEBUG(logger, "Dynamic bypass: not passthrough, stream traceid = %s", pmeinfo->stream_traceid);
		if(traceid2sslinfo_htable_add(pmeinfo->stream_traceid, strlen(pmeinfo->stream_traceid),ssl_feature,thread_seq) >=0)
		{
			ret = 1;
		}
		else
		{
			FREE(&ssl_feature);
		}
	}

	return ret;
}


void next_data_ssl_dynamic_bypass(struct pkt_info *pktinfo)
{
	//FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_INTCP_BYTE], 0, FS_OP_ADD, pktinfo->ip_totlen);
	return;
}


int ssl_dynamic_bypass_htable_add(struct pme_info *pmeinfo)
{
	void *logger = g_kni_handle->local_logger;
	long cb_ret;
	int ret = -1;
	MESA_htable_handle traceid2sslinfo_htable = g_kni_handle->threads_handle[pmeinfo->thread_seq].traceid2sslinfo_htable;
	if(pmeinfo->ssl_intercept_state == 0)
	{
		cb_ret = traceid2sslinfo_htable_search(traceid2sslinfo_htable,pmeinfo->stream_traceid, strlen(pmeinfo->stream_traceid), pmeinfo->thread_seq, pmeinfo->ssl_pinningst);
		if(cb_ret >= 0)
		{
			cb_ret = MESA_htable_del(traceid2sslinfo_htable, (const unsigned char *)pmeinfo->stream_traceid, strlen(pmeinfo->stream_traceid), NULL);
			if(cb_ret < 0)
			{
				KNI_LOG_ERROR(logger, "Delete traceid2sslinfo_htable fail,key=%s",pmeinfo->stream_traceid);
			}
			else
			{
				FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_DEL_SUCC], 0, FS_OP_ADD, 1);
				FS_operate(g_kni_fs_handle->handle, g_kni_fs_handle->fields[KNI_FIELD_ID2SSL_CNT], 0, FS_OP_ADD, -1);
			}
			ret = 0;
		}
	}
	return ret;		
}


int ssl_dynamic_bypass_htable_init(const char *profile,void *local_logger)
{
	struct kni_htable_opt opt;
	MESA_htable_handle sslinfo2bypass_htable = NULL;
	int ret = 0;
	int ssl_dynamic_bypass_enable = 1;

	MESA_load_profile_int_def(profile, "ssl_dynamic_bypass", "enabled", &ssl_dynamic_bypass_enable, 1);
    KNI_LOG_ERROR(local_logger, "Dynamic bypass: MESA_prof_load, [%s]:\n enabled: %d", "ssl_dynamic_bypass",ssl_dynamic_bypass_enable);
	g_kni_handle->ssl_dynamic_bypass_enable = ssl_dynamic_bypass_enable;

	if(ssl_dynamic_bypass_enable == 1){
		//init traceid2sslinfo_htable
		memset(&opt, 0, sizeof(opt));
		kni_get_htable_opt(&opt, profile, "traceid2sslinfo_htable", (void*)traceid2sslinfo_htable_data_free_cb, (void *)traceid2sslinfo_htable_expire_notify_cb, local_logger);
		for(int i = 0; i < g_kni_handle->thread_count; i++){
			MESA_htable_handle traceid2sslinfo_htable = kni_create_htable((char*)"traceid2sslinfo_htable", &opt, local_logger);
			if(traceid2sslinfo_htable == NULL){
				KNI_LOG_ERROR(local_logger, "Failed at kni_create_htable, table = traceid2sslinfo_htable");
				ret = -1;
				return ret;
			}
			g_kni_handle->threads_handle[i].traceid2sslinfo_htable = traceid2sslinfo_htable;
		}
		
		//init sslinfo2bypass_htable
		memset(&opt, 0, sizeof(opt));
		kni_get_htable_opt(&opt, profile, "sslinfo2bypass_htable", (void*)sslinfo2bypass_htable_data_free_cb, (void *)sslinfo2bypass_htable_expire_notify_cb, local_logger);
		sslinfo2bypass_htable = kni_create_htable((char*)"sslinfo2bypass_htable", &opt, local_logger);
		if(sslinfo2bypass_htable == NULL){
			KNI_LOG_ERROR(local_logger, "Failed at create sslinfo2bypass_htable");
			ret = -2;
			return ret;
		}
		g_kni_handle->sslinfo2bypass_htable = sslinfo2bypass_htable;
	}
	return ret;

}