#include <ssl_service_cache.h>
#include <MESA/maat.h>
#include <ssl_stream.h>
#include <cjson/cJSON.h>

struct ssl_svc_addr
{
	const char *sip;
	const char *sport;
	const char *dip;
	const char *dport;
};

struct ssl_svc_client_st
{
	time_t last_update_time;
	unsigned int suspect_pinning_count;	
	unsigned int protocol_error_count;
	char is_mutual_auth;
	struct ssl_service_cache* ref_svc_cache;
};
struct ssl_svc_server_st
{
	char is_ev;
	char is_ct;
	long long ev_st_switched;
	long long ct_st_switched;
	struct ssl_service_cache* ref_svc_cache;
};
struct ssl_svc_app_st
{
	unsigned int down_ssl_success_cnt;
	struct ssl_service_cache *ref_svc_cache;
};
struct ssl_service_write_args
{
	struct ssl_service_cache* cache;
	const struct ssl_service_status* status;
};
static void ssl_svc_free_client_st(void * data)
{
	struct ssl_svc_client_st* p = (struct ssl_svc_client_st *) data;
	struct ssl_service_cache* svc_cache=p->ref_svc_cache;
	if(p->is_mutual_auth)
	{
		svc_cache->stat.mutual_auth_cli_cnt--;
	}
	if(p->suspect_pinning_count>=svc_cache->fail_as_cli_pinning_count)
	{
		svc_cache->stat.pinning_cli_cnt--;
	}

	if (p->protocol_error_count >= svc_cache->fail_as_proto_err_count)
	{
		svc_cache->stat.proto_err_cli_cnt--;
	}

	free(p);
	return;
}
static void ssl_svc_free_server_st(void * data)
{
	struct ssl_svc_server_st* p = (struct ssl_svc_server_st *) data;
	struct ssl_service_cache* svc_cache=p->ref_svc_cache;
	if(p->is_ct)
	{
		svc_cache->stat.ct_srv_cnt--;
	}
	if(p->is_ev)
	{
		svc_cache->stat.ev_srv_cnt--;
	}
	free(p);
	return;
}
static void ssl_svc_free_app_st(void *data)
{
	struct ssl_svc_app_st *p = (struct ssl_svc_app_st *)data;
	struct ssl_service_cache *svc_cache = p->ref_svc_cache;
	if (p->down_ssl_success_cnt > svc_cache->succ_as_app_not_pinning_count)
	{
		svc_cache->stat.app_not_pinning_cnt--;
	}
	free(p);
}
static size_t ssl_svc_server_st_mk_key(struct ssl_svc_addr *addr_info, const struct ssl_chello* chello, char* key_buff, size_t sz)
{
	size_t key_len=0;
	key_len=snprintf(key_buff, sz, "%s:%s:%s:", addr_info->dip, addr_info->dport, chello->sni ? chello->sni : "null");
	return key_len;
}
static size_t ssl_svc_app_st_mk_key(struct ssl_svc_addr *addr_info, const struct ssl_chello *chello, char *key_buff, size_t sz)
{
    // splicing ssl attributes to generate hashkey
    size_t key_len=snprintf(key_buff, sz, "%d:%d:%s:%s", chello->min_version.ossl_format, 
												chello->max_version.ossl_format,
												chello->sni?chello->sni: addr_info->dip ,
												chello->alpn?chello->alpn:"null");
	if(chello->cipher_suites && sz-key_len>chello->cipher_suites_len)
	{
		memcpy(key_buff+key_len, chello->cipher_suites, chello->cipher_suites_len);
		key_len+=chello->cipher_suites_len;
	}
	if(chello->sign_algos && sz-key_len > chello->sign_algos_len)
	{
		memcpy(key_buff+key_len, chello->sign_algos, chello->sign_algos_len);
		key_len+=chello->sign_algos_len;
	}
	if(chello->supported_groups && sz-key_len > chello->supported_groups_len)
	{
		memcpy(key_buff+key_len, chello->supported_groups, chello->supported_groups_len);
		key_len+=chello->supported_groups_len;
	}
	return key_len;
}
static size_t ssl_svc_client_st_mk_key(struct ssl_svc_addr *addr_info, const struct ssl_chello *chello, char *key_buff, size_t sz)
{
	size_t key_len=0;
	char chello_id_buff[sz];
	size_t chello_id_len=0;

	key_len=snprintf(key_buff, sz, "%s:", addr_info->sip);
    chello_id_len = ssl_svc_app_st_mk_key(addr_info, chello, chello_id_buff, sizeof(chello_id_buff));
    memcpy(key_buff+key_len, chello_id_buff, MIN(chello_id_len, sz-key_len));
	key_len += MIN(chello_id_len, sz-key_len);

	return key_len;
}
static long cli_st_read_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct ssl_svc_client_st* cli_st=(struct ssl_svc_client_st*)data;
	struct ssl_service_status* result=(struct ssl_service_status*)user_arg;
	
	if (cli_st == NULL)
	{
		return 0;
	}
	struct ssl_service_cache* svc_cache=cli_st->ref_svc_cache;
	if(cli_st->suspect_pinning_count==0)
	{		
		result->cli_pinning_status=PINNING_ST_NOT_PINNING;
	}
	else if(cli_st->suspect_pinning_count<svc_cache->fail_as_cli_pinning_count)
	{
		result->cli_pinning_status=PINNING_ST_MAYBE_PINNING;
	}
	else
	{
		result->cli_pinning_status=PINNING_ST_PINNING;
	}
	if(cli_st->protocol_error_count>=svc_cache->fail_as_proto_err_count)
	{
		result->has_protocol_errors=1;
	}
	result->is_mutual_auth=cli_st->is_mutual_auth;
	return 1;
}
static long cli_st_write_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct ssl_svc_client_st* cli_st=(struct ssl_svc_client_st*)data;
	struct ssl_service_write_args* args=(struct ssl_service_write_args*)user_arg;
	const struct ssl_service_status* status=args->status;
	struct ssl_service_cache* cache=args->cache;
	UNUSED int ret = 0;
	time_t now=time(NULL);
	if(cli_st==NULL)
	{
		cli_st=ALLOC(struct ssl_svc_client_st, 1);
		cli_st->ref_svc_cache=cache;
		cli_st->last_update_time = now;
		ret = MESA_htable_add(cache->cli_st_hash, key, size, cli_st);
		assert(ret >= 0);
	}
	if (now - cli_st->last_update_time > cache->fail_time_window)
	{
		if(cli_st->suspect_pinning_count<cache->fail_as_cli_pinning_count)	cli_st->suspect_pinning_count=0;
		if(cli_st->protocol_error_count<cache->fail_as_proto_err_count)	cli_st->protocol_error_count=0;
		cli_st->last_update_time = now;
	}
	if(status->cli_pinning_status!=PINNING_ST_NOT_PINNING && cli_st->suspect_pinning_count<cache->fail_as_cli_pinning_count)
	{
		if(status->cli_pinning_status==PINNING_ST_PINNING)
		{
			cli_st->suspect_pinning_count=cache->fail_as_cli_pinning_count;
		}
		else
		{
			cli_st->suspect_pinning_count++;
		}
		if(cli_st->suspect_pinning_count==cache->fail_as_cli_pinning_count)
		{
			cache->stat.pinning_cli_cnt++;
		}
	}
	else if(status->cli_pinning_status==PINNING_ST_PINNING)
	{
		cli_st->suspect_pinning_count=cache->fail_as_cli_pinning_count;
	}
	if(status->has_protocol_errors)
	{
		cli_st->protocol_error_count++;
		if(cli_st->protocol_error_count==cache->fail_as_proto_err_count)
		{
			cache->stat.proto_err_cli_cnt++;
		}
	}
	if(status->is_mutual_auth==1&&cli_st->is_mutual_auth==0)
	{
		cache->stat.mutual_auth_cli_cnt++;	
		cli_st->is_mutual_auth=1;
	}
	cli_st->last_update_time=now;	
	return 1;
}

static long srv_st_read_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct ssl_svc_server_st* srv_st=(struct ssl_svc_server_st*)data;
	struct ssl_service_status* result=(struct ssl_service_status*)user_arg;
	if (srv_st == NULL)
	{
		return 0;
	}
	result->is_ct=srv_st->is_ct;
	result->is_ev=srv_st->is_ev;
	return 1;
}

static long srv_st_write_cb(void * data, const uchar * key, uint size, void * user_arg)
{
	struct ssl_svc_server_st* srv_st=(struct ssl_svc_server_st*)data;
	struct ssl_service_write_args* args=(struct ssl_service_write_args*)user_arg;
	const struct ssl_service_status* status=args->status;
	struct ssl_service_cache* cache=args->cache;
	UNUSED int ret = 0;
	if(srv_st==NULL)
	{
		srv_st=ALLOC(struct ssl_svc_server_st, 1);
		srv_st->ref_svc_cache=cache;
		ret = MESA_htable_add(cache->srv_st_hash, key, size, srv_st);
		assert(ret >= 0);
	}
	if(status->is_ev==1&&srv_st->is_ev==0)
	{
		srv_st->is_ev=1;
		cache->stat.ev_srv_cnt++;
	}
	if(status->is_ev!=srv_st->is_ev)
	{
		srv_st->ev_st_switched++;
	}
	if(status->is_ct==1&&srv_st->is_ct==0)
	{
		srv_st->is_ct=1;
		cache->stat.ct_srv_cnt++;
	}
	if(status->is_ct!=srv_st->is_ct)
	{
		srv_st->ct_st_switched++;
	}
//	assert(srv_st->ev_st_switched<2&&srv_st->ct_st_switched<2);
	return 1;
}

static long app_st_read_cb(void *data, const uchar *key, uint size, void *user_arg)
{
	struct ssl_svc_app_st *app_st = (struct ssl_svc_app_st *)data;
	struct ssl_service_status *result = (struct ssl_service_status *)user_arg;
	if (app_st == NULL)
	{
		return 0;
	}
	if (app_st->down_ssl_success_cnt > app_st->ref_svc_cache->succ_as_app_not_pinning_count)
	{
		result->is_app_not_pinning = 1;
	}
	return 1;
}

static long app_st_write_cb(void *data, const uchar *key, uint size, void *user_arg)
{
	struct ssl_svc_app_st *app_st = (struct ssl_svc_app_st *)data;
	struct ssl_service_write_args *args = (struct ssl_service_write_args *)user_arg;
	const struct ssl_service_status *status = args->status;
	struct ssl_service_cache *cache = args->cache;
	UNUSED int ret = 0;
	if (app_st == NULL)
	{
		app_st = ALLOC(struct ssl_svc_app_st, 1);
		app_st->ref_svc_cache = cache;
		ret = MESA_htable_add(cache->app_st_hash, key, size, app_st);
		assert(ret >= 0);
	}
	if (status->is_app_not_pinning)
	{
		app_st->down_ssl_success_cnt++;
	}
	if (app_st->down_ssl_success_cnt > cache->succ_as_app_not_pinning_count)
	{
		cache->stat.app_not_pinning_cnt++;
	}
	return 1;
}

int ssl_service_cache_read(struct ssl_service_cache *svc_cache, const struct ssl_chello *chello, const struct tfe_stream *tcp_stream, struct ssl_service_status *result)
{
	long cli_st_cb_ret=0, svr_st_cb_ret=0, app_st_cb_ret=0;
	char temp_key[2048];
	unsigned char hash_key[4096]; // Size must be twice the size of temp_key
	size_t temp_key_sz=0;
	size_t hash_key_sz=0;
	if(chello->sni==NULL)
	{
		return 0;
	}

	struct ssl_svc_addr addr_info = { NULL, NULL, NULL, NULL};
 	char *addr_str = tfe_strdup(tcp_stream->str_stream_info);
    tfe_stream_addr_str_split(addr_str, &(addr_info.sip), &(addr_info.sport), &(addr_info.dip), &(addr_info.dport));

	memset(hash_key, 0, sizeof(hash_key));
    temp_key_sz = ssl_svc_client_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
    hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
	MESA_htable_search_cb(svc_cache->cli_st_hash, hash_key, (unsigned int) hash_key_sz, cli_st_read_cb, result, &cli_st_cb_ret);
	TFE_LOG_DEBUG(g_default_logger, "client table, hash:%s, found:%ld, sni:%s, addr:%s, mutual:%d, pinning:%d, err:%d",
				hash_key, cli_st_cb_ret, chello->sni, tcp_stream->str_stream_info, result->is_mutual_auth, result->cli_pinning_status, result->has_protocol_errors);

	memset(hash_key, 0, sizeof(hash_key));
    temp_key_sz = ssl_svc_server_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
    hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
	MESA_htable_search_cb(svc_cache->srv_st_hash, hash_key, (unsigned int) hash_key_sz, srv_st_read_cb, result, &svr_st_cb_ret);
	TFE_LOG_DEBUG(g_default_logger, "server table, hash:%s, found:%ld, sni:%s, addr:%s, ct:%d, ev:%d",
				hash_key, svr_st_cb_ret, chello->sni, tcp_stream->str_stream_info, result->is_ct, result->is_ev);

	memset(hash_key, 0, sizeof(hash_key));
	temp_key_sz = ssl_svc_app_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
	hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
	MESA_htable_search_cb(svc_cache->app_st_hash, hash_key, (unsigned int)hash_key_sz, app_st_read_cb, result, &app_st_cb_ret);
	TFE_LOG_DEBUG(g_default_logger, "app table, hash:%s, found:%ld, sni:%s, addr:%s, app_not_pinning:%d",
				  hash_key, app_st_cb_ret, chello->sni, addr_str, result->is_app_not_pinning);

	free(addr_str);

    if(cli_st_cb_ret||svr_st_cb_ret||app_st_cb_ret)
	{
		return 1;
	}
	else
	{
		return 0;
	}
}

void ssl_service_cache_write(struct ssl_service_cache *svc_cache, const struct ssl_chello *chello, const struct tfe_stream *tcp_stream, const struct ssl_service_status *status)
{
	long cli_st_cb_ret=0, svr_st_cb_ret=0;
	char temp_key[2048];
	unsigned char hash_key[4096]; // Size must be twice the size of temp_key
	size_t temp_key_sz=0;
	size_t hash_key_sz=0;

	if(chello == NULL || chello->sni==NULL)
	{
		return;
	}

	struct ssl_svc_addr addr_info = { NULL, NULL, NULL, NULL};
 	char *addr_str = tfe_strdup(tcp_stream->str_stream_info);
    tfe_stream_addr_str_split(addr_str, &(addr_info.sip), &(addr_info.sport), &(addr_info.dip), &(addr_info.dport));

    struct ssl_service_write_args write_args={svc_cache, status};
	if(status->is_mutual_auth||status->cli_pinning_status!=PINNING_ST_NOT_PINNING||status->has_protocol_errors)
	{
		memset(hash_key, 0, sizeof(hash_key));
        temp_key_sz = ssl_svc_client_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
        hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
        TFE_LOG_DEBUG(g_default_logger, "client table, hash:%s, sni:%s, addr:%s, mutual:%d, pinning:%d, err:%d",
                      hash_key, chello->sni, tcp_stream->str_stream_info, status->is_mutual_auth, status->cli_pinning_status, status->has_protocol_errors);
        MESA_htable_search_cb(svc_cache->cli_st_hash, hash_key, (unsigned int) hash_key_sz, cli_st_write_cb, &write_args, &cli_st_cb_ret);	
	}
	if(status->is_ct||status->is_ev)
	{	
		memset(hash_key, 0, sizeof(hash_key));
        temp_key_sz = ssl_svc_server_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
        hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
        TFE_LOG_DEBUG(g_default_logger, "server table, hash:%s, sni:%s, addr:%s, ct:%d, ev:%d",
                      hash_key, chello->sni, tcp_stream->str_stream_info, status->is_ct, status->is_ev);
        MESA_htable_search_cb(svc_cache->srv_st_hash, hash_key, (unsigned int) hash_key_sz, srv_st_write_cb, &write_args, &svr_st_cb_ret);
	}
	if (status->is_app_not_pinning)
	{
		memset(hash_key, 0, sizeof(hash_key));
		temp_key_sz = ssl_svc_app_st_mk_key(&addr_info, chello, temp_key, sizeof(temp_key));
		hash_key_sz = tfe_hexdump(hash_key, (unsigned char *)temp_key, temp_key_sz) - hash_key;
		TFE_LOG_DEBUG(g_default_logger, "app table, hash:%s, sni:%s, addr:%s, app_not_pinning:%d",
					  hash_key, chello->sni, addr_str, status->is_app_not_pinning);
		MESA_htable_search_cb(svc_cache->app_st_hash, hash_key, (unsigned int)hash_key_sz, app_st_write_cb, &write_args, &svr_st_cb_ret);
	}

	free(addr_str);
}
struct ssl_service_cache *ssl_service_cache_create(unsigned int slot_size, unsigned int expire_seconds, int fail_as_pinning_cnt, int fail_as_proto_err_cnt, int fail_time_win, int succ_as_app_not_pinning_cnt)
{
	struct ssl_service_cache * cache = ALLOC(struct ssl_service_cache, 1);
	unsigned max_num = slot_size * 4;
	UNUSED int ret = 0;
	MESA_htable_handle htable=NULL, saved[3];
	int i=0, opt_val=0;
	cache->fail_as_cli_pinning_count=fail_as_pinning_cnt;
	cache->fail_as_proto_err_count=fail_as_proto_err_cnt;
	cache->fail_time_window=fail_time_win;
	cache->succ_as_app_not_pinning_count = succ_as_app_not_pinning_cnt;
	void (*free_func[])(void *)={ssl_svc_free_client_st, ssl_svc_free_server_st, ssl_svc_free_app_st};
	for(i=0; i<3; i++)
	{
		htable = MESA_htable_born();
		opt_val=0;
		ret = MESA_htable_set_opt(htable, MHO_SCREEN_PRINT_CTRL, &opt_val, sizeof(opt_val));
		opt_val=1;
		ret = MESA_htable_set_opt(htable, MHO_THREAD_SAFE, &opt_val, sizeof(opt_val));
		opt_val=16;
		ret = MESA_htable_set_opt(htable, MHO_MUTEX_NUM, &opt_val, sizeof(opt_val));
		ret = MESA_htable_set_opt(htable, MHO_HASH_SLOT_SIZE, &slot_size, sizeof(slot_size));
		ret = MESA_htable_set_opt(htable, MHO_HASH_MAX_ELEMENT_NUM, &max_num, sizeof(max_num));
		// for app table
		if (i == 2)
		{
			unsigned int app_expire_seconds = 2 * expire_seconds;
			ret = MESA_htable_set_opt(htable, MHO_EXPIRE_TIME, &app_expire_seconds, sizeof(app_expire_seconds));
			opt_val = HASH_ELIMINATE_ALGO_LRU;
			ret = MESA_htable_set_opt(htable, MHO_ELIMIMINATE_TYPE, &opt_val, sizeof(int));
		}
		// for client table and server table
		else
		{
			ret = MESA_htable_set_opt(htable, MHO_EXPIRE_TIME, &expire_seconds, sizeof(expire_seconds));
			opt_val = HASH_ELIMINATE_ALGO_FIFO;
			ret = MESA_htable_set_opt(htable, MHO_ELIMIMINATE_TYPE, &opt_val, sizeof(int));
		}

		ret = MESA_htable_set_opt(htable, MHO_CBFUN_DATA_FREE,
			(void*)free_func[i], sizeof(free_func[i]));
		
		ret = MESA_htable_mature(htable);
		assert(ret == 0);
		saved[i]=htable;
	}
	cache->cli_st_hash=saved[0];
	cache->srv_st_hash=saved[1];
	cache->app_st_hash=saved[2];

	return cache;
}
void ssl_service_cache_destroy(struct ssl_service_cache* cache)
{
	MESA_htable_destroy(cache->cli_st_hash, NULL);
	cache->cli_st_hash=NULL;
	MESA_htable_destroy(cache->srv_st_hash, NULL);
	cache->srv_st_hash=NULL;
	MESA_htable_destroy(cache->app_st_hash, NULL);
	cache->app_st_hash=NULL;
	free(cache);
	return;
}
void ssl_service_cache_stat(struct ssl_service_cache* svc_cache, struct ssl_service_cache_statistics* result)
{
	*result=svc_cache->stat;
	return;
}