stellar-stellar/decoders/lpi_plus/lpip_module.c

/*
 * 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 "stellar/lpi_plus.h"

#include "lpip_extend.h"


#define	MAX_APPID_NUM	8

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 lpip_packet_exdata_idx;
	struct module_manager *mod_mgr;
	struct lpi_plus_mapper *mapper;
	struct session_manager *sess_mgr;
};

struct lpi_plus_per_session_ctx
{
	unsigned int detected_pkt_cnt;
	int stop_detect;
	struct lpi_plus_detect_context *detector_ctx;
	int appid[MAX_APPID_NUM];
	int packet_sequence[MAX_APPID_NUM];
	size_t appid_num;
};

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

	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;
}


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 packet *pkt, 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);
    (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(pkt, 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, int *stop_detect)
{
	lpi_module_t *plpi_mod=lpi_guess_protocol(&(ctx->lpi_data));
	if(plpi_mod==NULL)return 0;
	if(stop_detect!=NULL)*stop_detect=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))
	{
		if(stop_detect!=NULL)*stop_detect=1;
	}
	return new_appid;
}

void lpi_plus_on_packet(struct packet *pkt, struct module *mod)
{
	if(pkt==NULL||mod==NULL)return;

	struct lpi_plus *env=(struct lpi_plus *)module_get_ctx(mod);
	struct session *sess=packet_exdata_to_session(env->sess_mgr, pkt);
	
	if(sess==NULL)return;
	
	if (session_get_current_state(sess) == SESSION_STATE_CLOSED)
	{
		assert(pkt == NULL);
		return;
	}

	struct lpi_plus_per_session_ctx *lpip_sess_ctx = (struct lpi_plus_per_session_ctx *)session_get_exdata(sess, env->lpip_session_exdata_idx);
	if(lpip_sess_ctx==NULL)
	{
		lpip_sess_ctx= CALLOC(struct lpi_plus_per_session_ctx, 1);
		session_set_exdata(sess, env->lpip_session_exdata_idx, lpip_sess_ctx);	
	}
	if(lpip_sess_ctx->stop_detect==1)return;
	if(lpip_sess_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
	{
		if(lpip_sess_ctx->detector_ctx==NULL)lpip_sess_ctx->detector_ctx=CALLOC(struct lpi_plus_detect_context, 1);
        lpi_plus_context_update(pkt, sess, lpip_sess_ctx->detector_ctx, payload, payload_len);
		int appid=lpi_plus_detect(lpip_sess_ctx->detector_ctx, env->mapper, payload, payload_len, &lpip_sess_ctx->stop_detect);
		lpip_sess_ctx->detected_pkt_cnt+=1;
        if(appid>0 && lpi_plus_appid_update(lpip_sess_ctx->appid, &(lpip_sess_ctx->appid_num), appid))
		{
			lpip_sess_ctx->packet_sequence[lpip_sess_ctx->appid_num-1]=lpip_sess_ctx->detected_pkt_cnt;
			packet_set_exdata(pkt, env->lpip_packet_exdata_idx, lpip_sess_ctx);
		}
	}
	return;
}

int32_t *packet_exdata_to_lpip_appid(struct lpi_plus *lpip, struct packet *pkt, size_t *appid_num)
{
	if(lpip==NULL || pkt==NULL || appid_num==NULL)return NULL;

	struct lpi_plus_per_session_ctx *lpip_sess_ctx=(struct lpi_plus_per_session_ctx *)packet_get_exdata(pkt, lpip->lpip_packet_exdata_idx);
	if(lpip_sess_ctx==NULL)return NULL;

	*appid_num=lpip_sess_ctx->appid_num;
	return lpip_sess_ctx->appid;
}

static void lpi_plus_exdata_free(int idx __unused, void *ex_ptr, void *arg __unused)
{
	if(ex_ptr==NULL)return;
	struct lpi_plus_per_session_ctx *lpip_sess_ctx=(struct lpi_plus_per_session_ctx *)ex_ptr;
	if(lpip_sess_ctx->detector_ctx)FREE(lpip_sess_ctx->detector_ctx);
	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);
	}
}



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

	struct module *sess_mgr_mod=module_manager_get_module(mod_mgr, SESSION_MANAGER_MODULE_NAME);
	struct module *pkt_mgr_mod=module_manager_get_module(mod_mgr, PACKET_MANAGER_MODULE_NAME);
	if(sess_mgr_mod==NULL || pkt_mgr_mod == 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
	env->sess_mgr=module_to_session_manager(sess_mgr_mod);

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

	env->lpip_session_exdata_idx = session_manager_new_session_exdata_index(env->sess_mgr, "EXDATA_LPI", lpi_plus_exdata_free, NULL);	

	struct packet_manager *pkt_mgr = module_to_packet_manager(pkt_mgr_mod);
	env->lpip_packet_exdata_idx = packet_manager_new_packet_exdata_index(pkt_mgr, "EXDATA_LPI", NULL, NULL);

	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;
}