#include <ssl_stream.h>
#include <tfe_utils.h>
#include <MESA/Maat_rule.h>
#include <cjson/cJSON.h>
#include <openssl/ssl.h>
#include <assert.h>
#include <tfe_resource.h>
#include <tfe_cmsg.h>

struct ssl_policy_enforcer
{
	Maat_feather_t maat;
	int policy_table_id;
	int profile_table_id;
	void* logger;
};
struct intercept_param
{
	int policy_id;
	int ref_cnt;
	int keyring_for_trusted;
	int keyring_for_untrusted;
	int decryption_profile_id;
};

struct decryption_param
{
	int ref_cnt;
	int bypass_ev_cert;
	int bypass_ct_cert;
	int bypass_mutual_auth;
	int bypass_pinning;
	int bypass_uninstall_cert_traffic;
	int bypass_protocol_errors;
	int no_verify_cn;
	int no_verify_issuer;
	int no_verify_self_signed;
	int no_verify_expry_date;
	int block_fake_cert;
	int ssl_min_version;
	int ssl_max_version;
	int allow_http2;
	int mirror_client_version;
};

void intercept_param_dup_cb(int table_id, MAAT_PLUGIN_EX_DATA* to, MAAT_PLUGIN_EX_DATA* from, long argl, void* argp)
{
	struct intercept_param* param= (struct intercept_param*) *from;
	if(param)
	{
		 __sync_add_and_fetch(&(param->ref_cnt), 1);
		*to = param;
	}
	else
	{
		*to=NULL;
	}
	return;
}
void intercept_param_new_cb(int table_id, const char* key, const char* table_line, MAAT_PLUGIN_EX_DATA* ad, long argl, void* argp)
{
	int ret=0;
	size_t intercept_user_region_offset=0, len=0;
	char* json_str=NULL;
	cJSON *json=NULL, *item=NULL;
	struct intercept_param* param=NULL;

	struct ssl_policy_enforcer* enforcer=(struct ssl_policy_enforcer*)argp;
	ret=Maat_helper_read_column(table_line, 7, &intercept_user_region_offset, &len);
	if(ret<0)
	{
		TFE_LOG_ERROR(enforcer->logger, "Get intercept user region: %s", table_line);
		return;
	}
	json_str=ALLOC(char, len+1);
	memcpy(json_str, table_line+intercept_user_region_offset, len);
	json=cJSON_Parse(json_str);
	if(json==NULL)
	{
		TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: id = %s", key);
		goto error_out;
	}

	item=cJSON_GetObjectItem(json, "protocol");
	if(unlikely(!item || !cJSON_IsString(item)))
	{
		TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %s invalid protocol format", key);
        goto error_out;
	}
	if(0!=strcasecmp(item->valuestring, "SSL")&& 0!=strcasecmp(item->valuestring, "HTTP"))
	{
		goto error_out;
	}

	param=ALLOC(struct intercept_param, 1);
	param->policy_id=atoi(key);
	param->ref_cnt=1;
	/*
	param->bypass_mutual_auth=1;
	param->bypass_pinning=1;
	param->mirror_client_version=1;
	*/
	param->keyring_for_trusted=1;
	param->keyring_for_untrusted=0;
	param->decryption_profile_id=0;

	item=cJSON_GetObjectItem(json, "keyring_for_trusted");
	if(item)
	{
		if(item->type==cJSON_Number)
		{
			param->keyring_for_trusted=item->valueint;
		}
		else if(item->type==cJSON_String)
		{
			param->keyring_for_trusted=atoi(item->valuestring);
		}
		else
		{
			TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %d invalid keyring_for_trusted format", param->policy_id);
		}
	}

	item=cJSON_GetObjectItem(json, "keyring_for_untrusted");
	if(item)
	{
		if(item->type==cJSON_Number)
		{
			param->keyring_for_untrusted=item->valueint;
		}
		else if(item->type==cJSON_String)
		{
			param->keyring_for_untrusted=atoi(item->valuestring);
		}
		else
		{
			TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %d invalid keyring_for_untrusted format", param->policy_id);
		}
	}

	item=cJSON_GetObjectItem(json, "decryption");
	if(item)
	{
		if(item->type==cJSON_Number)
		{
			param->decryption_profile_id=item->valueint;
		}
		else if(item->type==cJSON_String)
		{
			param->decryption_profile_id=atoi(item->valuestring);
		}
		else
		{
			TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter:%d invalid decryption format", param->policy_id);
		}
	}
	*ad=param;
	TFE_LOG_INFO(enforcer->logger, "Add intercept policy: %d", param->policy_id);
error_out:
	cJSON_Delete(json);
	free(json_str);
	return;
}
void intercept_param_free_cb(int table_id, MAAT_PLUGIN_EX_DATA* ad, long argl, void* argp)
{
	struct ssl_policy_enforcer* enforcer=(struct ssl_policy_enforcer*)argp;
	struct intercept_param* param= (struct intercept_param*) *ad;
	if(param==NULL)
	{
		return;
	}

    if ((__sync_sub_and_fetch(&param->ref_cnt, 1) == 0))
	{
		TFE_LOG_INFO(enforcer->logger, "Del intercept policy %d", param->policy_id);\
		free(param);
		*ad=NULL;
	}
}
void intercept_param_free(struct intercept_param* param)
{
	intercept_param_free_cb(0, (void**)&param, 0, NULL);
	return;
}

void profile_param_dup_cb(int table_id, MAAT_PLUGIN_EX_DATA* to, MAAT_PLUGIN_EX_DATA* from, long argl, void* argp)
{
	struct decryption_param* param= (struct decryption_param*) *from;
	if(param)
	{
		 __sync_add_and_fetch(&(param->ref_cnt), 1);
		*to = param;
	}
	else
	{
		*to=NULL;
	}
	return;
}
void profile_param_free_cb(int table_id, MAAT_PLUGIN_EX_DATA* ad, long argl, void* argp)
{
	struct decryption_param* param= (struct decryption_param*) *ad;
	if(param==NULL)
	{
		return;
	}

    if ((__sync_sub_and_fetch(&param->ref_cnt, 1) == 0))
	{
		free(param);
		*ad=NULL;
	}
}
void profile_param_free(struct decryption_param* param)
{
	profile_param_free_cb(0, (void**)&param, 0, NULL);
	return;
}
void profile_param_new_cb(int table_id, const char* key, const char* table_line, MAAT_PLUGIN_EX_DATA* ad, long argl, void* argp)
{
	int ret=0;
	size_t offset=0, len=0;
	char* json_str=NULL;
	cJSON *json=NULL, *exclusions=NULL, *cert_verify=NULL, *approach=NULL, *ssl_ver=NULL, *item=NULL;
	struct decryption_param* param=NULL;
	struct ssl_policy_enforcer* enforcer=(struct ssl_policy_enforcer*)argp;

	ret=Maat_helper_read_column(table_line, 3, &offset, &len);
	if(ret<0)
	{
		TFE_LOG_ERROR(enforcer->logger, "Get decryption param: %s", table_line);
		goto error_out;
	}
	json_str=ALLOC(char, len+1);
	memcpy(json_str, table_line+offset, len);
	json=cJSON_Parse(json_str);
	if(json==NULL)
	{
		TFE_LOG_ERROR(enforcer->logger, "Invalid decryption parameter: %s", table_line);
		goto error_out;
	}

	param=ALLOC(struct decryption_param, 1);
	param->ref_cnt=1;
	param->bypass_mutual_auth=1;
	param->bypass_pinning=1;
	param->mirror_client_version=1;

	exclusions=cJSON_GetObjectItem(json, "dynamic_bypass");
	if(exclusions)
	{
		item=cJSON_GetObjectItem(exclusions, "ev_cert");
		if(item && item->type==cJSON_Number) param->bypass_ev_cert=item->valueint;
		item=cJSON_GetObjectItem(exclusions, "cert_transparency");
		if(item && item->type==cJSON_Number) param->bypass_ct_cert=item->valueint;
		item=cJSON_GetObjectItem(exclusions, "mutual_authentication");
		if(item && item->type==cJSON_Number) param->bypass_mutual_auth=item->valueint;
		item=cJSON_GetObjectItem(exclusions, "cert_pinning");
		if(item && item->type==cJSON_Number) param->bypass_pinning=item->valueint;
		item=cJSON_GetObjectItem(exclusions, "protocol_errors");
		if(item && item->type==cJSON_Number) param->bypass_protocol_errors=item->valueint;
		item=cJSON_GetObjectItem(exclusions, "trusted_root_cert_is_not_installed_on_client");
		if(item && item->type==cJSON_Number) param->bypass_uninstall_cert_traffic=item->valueint;
	}
	ssl_ver=cJSON_GetObjectItem(json, "protocol_version");
	if(ssl_ver)
	{
		item=cJSON_GetObjectItem(ssl_ver, "mirror_client");
		if(item && item->type==cJSON_Number) param->mirror_client_version=item->valueint;
		if(!param->mirror_client_version)
		{
			item=cJSON_GetObjectItem(ssl_ver, "min");
			if(item && item->type==cJSON_String) param->ssl_min_version=sslver_str2num(item->valuestring);
			item=cJSON_GetObjectItem(ssl_ver, "max");
			if(item && item->type==cJSON_String) param->ssl_max_version=sslver_str2num(item->valuestring);
			if(param->ssl_min_version<0||param->ssl_max_version<0)
			{
				param->mirror_client_version=1;
				TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: ssl version = %s", item->valuestring);
			}
		}
		item=cJSON_GetObjectItem(ssl_ver, "allow_http2");
		if(item && item->type==cJSON_Number) param->allow_http2=item->valueint;
	}
	cert_verify=cJSON_GetObjectItem(json, "certificate_checks");
	if(cert_verify)
	{
		approach=cJSON_GetObjectItem(cert_verify, "approach");
		if(approach)
		{
			item=cJSON_GetObjectItem(approach, "cn");
			if(item && item->type==cJSON_Number && item->valueint==0) param->no_verify_cn=1;
			item=cJSON_GetObjectItem(approach, "issuer");
			if(item && item->type==cJSON_Number && item->valueint==0) param->no_verify_issuer=1;
			item=cJSON_GetObjectItem(approach, "self-signed");
			if(item && item->type==cJSON_Number && item->valueint==0) param->no_verify_self_signed=1;
			item=cJSON_GetObjectItem(approach, "expiration");
			if(item && item->type==cJSON_Number && item->valueint==0) param->no_verify_expry_date=1;
		}
		item=cJSON_GetObjectItem(cert_verify, "fail_action");
		if(item && item->type==cJSON_String)
		{
			if(0==strcasecmp(item->valuestring, "Fail-Close"))
			{
				param->block_fake_cert=1;
			}
		}
	}

	*ad=param;
	TFE_LOG_INFO(enforcer->logger, "Add decryption profile: %s", key);
error_out:
	if (json)
		cJSON_Delete(json);
	if (json_str)
		free(json_str);
	return;
}
struct ssl_policy_enforcer* ssl_policy_enforcer_create(void* logger)
{
	struct ssl_policy_enforcer* enforcer=ALLOC(struct ssl_policy_enforcer, 1);
	enforcer->maat=(Maat_feather_t)tfe_bussiness_resouce_get(STATIC_MAAT);;
	enforcer->logger=logger;
	enforcer->policy_table_id=Maat_table_register(enforcer->maat, "TSG_SECURITY_COMPILE");
	assert(enforcer->policy_table_id >= 0);
	enforcer->profile_table_id=Maat_table_register(enforcer->maat, "TSG_PROFILE_DECRYPTION");
	assert(enforcer->profile_table_id >= 0);
	UNUSED int ret=Maat_plugin_EX_register(enforcer->maat,
							enforcer->policy_table_id,
							intercept_param_new_cb,
							intercept_param_free_cb,
							intercept_param_dup_cb,
							NULL,
							0,
							enforcer);
	assert(ret==0);
	ret=Maat_plugin_EX_register(enforcer->maat,
							enforcer->profile_table_id,
							profile_param_new_cb,
							profile_param_free_cb,
							profile_param_dup_cb,
							NULL,
							0,
							enforcer);
	assert(ret==0);
	return enforcer;
}
enum ssl_stream_action ssl_policy_enforce(struct ssl_stream *upstream, void* u_para)
{
	UNUSED struct ssl_policy_enforcer* enforcer=(struct ssl_policy_enforcer*)u_para;
	struct intercept_param *policy_param=NULL;
	struct decryption_param *profile_param=NULL;
	enum ssl_stream_action action=SSL_ACTION_PASSTHROUGH;
	UNUSED int ret=0;
	int policy_id=0;
	char policy_id_str[16]={0};
	char profile_id_str[16]={0};
	char sni[512], addr_string[512];
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_INTERCEPT_POLICY_ID, &policy_id);
	assert(ret==0);
	snprintf(policy_id_str, sizeof(policy_id_str), "%d", policy_id);
	policy_param=(struct intercept_param *)Maat_plugin_get_EX_data(enforcer->maat, enforcer->policy_table_id, policy_id_str);
	if(policy_param==NULL)
	{
		TFE_LOG_INFO(enforcer->logger, "Failed to get intercept parameter of policy %d.", policy_id);
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Invalid Intercept Param");
		return SSL_ACTION_PASSTHROUGH;
	}
	else
	{
		ssl_stream_get_string_opt(upstream, SSL_STREAM_OPT_SNI, sni, sizeof(sni));
		ssl_stream_get_string_opt(upstream, SSL_STREAM_OPT_ADDR, addr_string, sizeof(addr_string));
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy %d", addr_string, sni, policy_id);
	}

	snprintf(profile_id_str, sizeof(profile_id_str), "%u", policy_param->decryption_profile_id);
	profile_param=(struct decryption_param *)Maat_plugin_get_EX_data(enforcer->maat, enforcer->profile_table_id, profile_id_str);
	if (profile_param==NULL)
	{
		TFE_LOG_INFO(enforcer->logger, "Failed to get decryption parameter of profile %s.", profile_id_str);
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Invalid Decryption Param");
		return SSL_ACTION_PASSTHROUGH;
	}
	int pinning_staus=0, is_ev=0, is_ct=0, is_mauth=0, has_error=0, ja3_pinning_status=0;
 	if(!profile_param->mirror_client_version)
	{
		ret=ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_PROTOCOL_MIN_VERSION, profile_param->ssl_min_version);
		ret=ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_PROTOCOL_MAX_VERSION, profile_param->ssl_max_version);
	}
	if(profile_param->allow_http2)
	{
		ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_ENABLE_ALPN, 1);
	}
	ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_NO_VERIFY_COMMON_NAME, profile_param->no_verify_cn);
	ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_NO_VERIFY_ISSUER, profile_param->no_verify_issuer);
	ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_NO_VERIFY_SELF_SIGNED, profile_param->no_verify_self_signed);
	ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_NO_VERIFY_EXPIRY_DATE, profile_param->no_verify_expry_date);
	if(profile_param->block_fake_cert)
	{
		ret=ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_BLOCK_FAKE_CERT, 1);
	}
	ret=ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_KEYRING_FOR_TRUSTED, policy_param->keyring_for_trusted);
	ret=ssl_stream_set_integer_opt(upstream, SSL_STREAM_OPT_KEYRING_FOR_UNTRUSTED, policy_param->keyring_for_untrusted);

	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_PINNING_STATUS, &pinning_staus);
	assert(ret==0);
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_JA3_PINNING_STATUS, &ja3_pinning_status);
	assert(ret==0);
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_IS_EV_CERT, &is_ev);
	assert(ret==0);
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_IS_MUTUAL_AUTH, &is_mauth);
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_IS_CT_CERT, &is_ct);
	ret=ssl_stream_get_integer_opt(upstream, SSL_STREAM_OPT_HAS_PROTOCOL_ERRORS, &has_error);
	assert(ret==0);

	if (pinning_staus == 1 && ja3_pinning_status == JA3_PINNING_STATUS_NOT_PINNING && profile_param->bypass_uninstall_cert_traffic)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Certificate Not Installed");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to Certificate Not Installed", addr_string, sni, policy_param->policy_id);
	}
	else if ((pinning_staus == 1 || ja3_pinning_status == JA3_PINNING_STATUS_IS_PINNING) && ja3_pinning_status != JA3_PINNING_STATUS_NOT_PINNING && profile_param->bypass_pinning)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Certificate Pinning");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to Certificate Pinning", addr_string, sni, policy_param->policy_id);
	}
	else if (is_mauth && profile_param->bypass_mutual_auth)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Mutual Authentication");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to Mutual Authentication", addr_string, sni, policy_param->policy_id);
	}
	else if (is_ev && profile_param->bypass_ev_cert)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "EV Certificate");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to EV Certificate", addr_string, sni, policy_param->policy_id);
	}
	else if (is_ct && profile_param->bypass_ct_cert)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Certificate Transparency");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to Certificate Transparency", addr_string, sni, policy_param->policy_id);
	}
	else if (has_error && profile_param->bypass_protocol_errors)
	{
		action = SSL_ACTION_PASSTHROUGH;
		ssl_stream_set_cmsg_string(upstream, TFE_CMSG_SSL_PASSTHROUGH_REASON, "Protocol Errors");
		TFE_LOG_DEBUG(enforcer->logger, "%s %s enforce policy_id %d, action PASSTHROUGH due to Protocol Errors", addr_string, sni, policy_param->policy_id);
	}
	else
	{
		action = SSL_ACTION_INTERCEPT;
	}

	intercept_param_free(policy_param);
	profile_param_free(profile_param);
	policy_param=NULL;
	profile_param=NULL;
	return action;
}