#include "kni_utils.h"
#include "kni_cmsg.h"

const char * tfe_cmsg_tlv_type_to_string[KNI_CMSG_TLV_NR_MAX];

struct kni_cmsg_tlv
{
	uint16_t type;
	uint16_t length;
	unsigned char value_as_string[0];
} __attribute__((packed));

struct kni_cmsg
{
	uint16_t nr_tlvs;
	struct kni_cmsg_tlv* tlvs[KNI_CMSG_TLV_NR_MAX];
	uint16_t size;
} __attribute__((packed));

struct kni_cmsg_serialize_header
{
	uint8_t __magic__[2]; /* Must be 0x4d, 0x5a */
	uint16_t nr_tlvs;
	struct kni_cmsg_tlv tlvs[0];
} __attribute__((packed));

struct kni_cmsg* kni_cmsg_init()
{
	struct kni_cmsg *cmsg = ALLOC(struct kni_cmsg, 1);
	cmsg->size = sizeof(struct kni_cmsg_serialize_header);
	return cmsg;
}

void kni_cmsg_destroy(struct kni_cmsg *cmsg)
{
	if(cmsg != NULL)
	{
		for(int i = 0; i < KNI_CMSG_TLV_NR_MAX; i++)
		{
			FREE(&(cmsg->tlvs[i]));
		}
		FREE(&cmsg);
	}
}

int kni_cmsg_set(struct kni_cmsg *cmsg, uint16_t type, const unsigned char *value, uint16_t size)
{
	if(type >= KNI_CMSG_TLV_NR_MAX)
	{
		return KNI_CMSG_INVALID_TYPE;
	}
	struct kni_cmsg_tlv *tlv = cmsg->tlvs[type];
	uint16_t length = sizeof(struct kni_cmsg_tlv) + size;
	if(tlv == NULL)
	{
		tlv = (struct kni_cmsg_tlv*)ALLOC(char, length);
		cmsg->nr_tlvs++;
		cmsg->size += length;
	}
	tlv->type = type;
	tlv->length = length;
	memcpy(tlv->value_as_string, value, size);
	cmsg->tlvs[type] = tlv;
	return 0;
}

int kni_cmsg_get(struct kni_cmsg *cmsg, uint16_t type, uint16_t *size, unsigned char **pvalue)
{
	struct kni_cmsg_tlv *tlv = NULL;
	if(type >= KNI_CMSG_TLV_NR_MAX)
	{
		*size = 0;
		return KNI_CMSG_INVALID_TYPE;
	}
	if((tlv = cmsg->tlvs[type]) == NULL){
		*size = 0;
		return KNI_CMSG_TYPE_UNSET;
	}
	*size = tlv->length - sizeof(struct kni_cmsg_tlv);
	*pvalue = tlv->value_as_string;
	return 0;
}

uint16_t kni_cmsg_serialize_size_get(struct kni_cmsg *cmsg)
{
	return cmsg->size;
}

int kni_cmsg_serialize(struct kni_cmsg *cmsg, unsigned char *buff, uint16_t bufflen, uint16_t *serialize_len)
{
	uint16_t size = cmsg->size;
	if(bufflen < size)
	{
		return KNI_CMSG_BUFF_NOT_ENOUGH;
	}
	if(size < sizeof(struct kni_cmsg_serialize_header))
	{
		return KNI_CMSG_INVALID_FORMAT;
	}
	struct kni_cmsg_serialize_header *header = (struct kni_cmsg_serialize_header*)buff;
	header->__magic__[0] = 0x4d;
	header->__magic__[1] = 0x5a;
	header->nr_tlvs = htons(cmsg->nr_tlvs);
	uint16_t offset = sizeof(struct kni_cmsg_serialize_header);
	int count = 0;
	for(int i = 0; i < KNI_CMSG_TLV_NR_MAX; i++){
		if(cmsg->tlvs[i] != NULL)
		{
			count++;
		}
	}
	if(count != cmsg->nr_tlvs)
	{
		return KNI_CMSG_INVALID_FORMAT;
	}
	for(int i = 0; i < KNI_CMSG_TLV_NR_MAX; i++)
	{
		struct kni_cmsg_tlv *tlv = cmsg->tlvs[i];
		if(tlv == NULL)
		{
			continue;
		}
		if(i != tlv->type)
		{
			return KNI_CMSG_INVALID_FORMAT;
		}
		uint16_t length = tlv->length;
		if(length < sizeof(struct kni_cmsg_tlv) || offset + length > size)
		{
			return KNI_CMSG_INVALID_FORMAT;
		}
		memcpy((char*)header + offset, (void*)tlv, length);
		struct kni_cmsg_tlv *tlv1 = (struct kni_cmsg_tlv*)((char*)header + offset);
		tlv1->type = htons(tlv->type);
		tlv1->length = htons(tlv->length);
		offset += length;
	}
	if(offset != size)
	{
		return KNI_CMSG_INVALID_FORMAT;
	}
	*serialize_len = size;
	return 0;
}

int kni_cmsg_deserialize(const unsigned char *data, uint16_t len, struct kni_cmsg** pcmsg)
{
	struct kni_cmsg *cmsg = NULL;
	struct kni_cmsg_serialize_header *header = (struct kni_cmsg_serialize_header*)data;
	int offset = 0, nr_tlvs = -1;
	if(len < sizeof(struct kni_cmsg_serialize_header))
	{
		goto error_out;
	}
	if(header->__magic__[0] != 0x4d || header->__magic__[1] != 0x5a)
	{
		goto error_out;
	}
	cmsg = ALLOC(struct kni_cmsg, 1);
	offset = sizeof(struct kni_cmsg_serialize_header);
	nr_tlvs = ntohs(header->nr_tlvs);
	for(int i = 0; i < nr_tlvs; i++)
	{
		struct kni_cmsg_tlv *tlv = (struct kni_cmsg_tlv*)(data + offset);
		if(offset + sizeof(struct kni_cmsg_tlv) > len)
		{
			goto error_out;
		}
		uint16_t type = ntohs(tlv->type);
		uint16_t length = ntohs(tlv->length);
		if(length < sizeof(struct kni_cmsg_tlv) || offset + length > len)
		{
			goto error_out;
		}
		int ret = kni_cmsg_set(cmsg, type, data + offset + sizeof(struct kni_cmsg_tlv), length - sizeof(struct kni_cmsg_tlv));
		if(ret < 0)
		{
			goto error_out;
		}
		offset += length;
	}
	cmsg->size = offset;
	*pcmsg = cmsg;
	return 0;

error_out:
	kni_cmsg_destroy(cmsg);
	return KNI_CMSG_INVALID_FORMAT;
}



void tfe_cmsg_enum_to_string()
{
	memset(tfe_cmsg_tlv_type_to_string, 0 ,sizeof(tfe_cmsg_tlv_type_to_string));
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_SEQ] = "TFE_CMSG_TCP_RESTORE_SEQ";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_ACK] = "TFE_CMSG_TCP_RESTORE_ACK";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_MSS_CLIENT] = "TFE_CMSG_TCP_RESTORE_MSS_CLIENT";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_MSS_SERVER] = "TFE_CMSG_TCP_RESTORE_MSS_SERVER";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_WSACLE_CLIENT] = "TFE_CMSG_TCP_RESTORE_WSACLE_CLIENT";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_WSACLE_SERVER] = "TFE_CMSG_TCP_RESTORE_WSACLE_SERVER";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_SACK_CLIENT] = "TFE_CMSG_TCP_RESTORE_SACK_CLIENT";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_SACK_SERVER] = "TFE_CMSG_TCP_RESTORE_SACK_SERVER";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_TS_CLIENT] = "TFE_CMSG_TCP_RESTORE_TS_CLIENT";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_TS_SERVER] = "TFE_CMSG_TCP_RESTORE_TS_SERVER";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_PROTOCOL] = "TFE_CMSG_TCP_RESTORE_PROTOCOL";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_WINDOW_CLIENT] = "TFE_CMSG_TCP_RESTORE_WINDOW_CLIENT";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_WINDOW_SERVER] = "TFE_CMSG_TCP_RESTORE_WINDOW_SERVER";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_INFO_PACKET_CUR_DIR] = "TFE_CMSG_TCP_RESTORE_INFO_PACKET_CUR_DIR";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_TS_CLIENT_VAL] = "TFE_CMSG_TCP_RESTORE_TS_CLIENT_VAL";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_RESTORE_TS_SERVER_VAL] = "TFE_CMSG_TCP_RESTORE_TS_SERVER_VAL";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_POLICY_ID] = "TFE_CMSG_POLICY_ID";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_STREAM_TRACE_ID] = "TFE_CMSG_STREAM_TRACE_ID";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_INTERCEPT_STATE] = "TFE_CMSG_SSL_INTERCEPT_STATE";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_UPSTREAM_LATENCY] = "TFE_CMSG_SSL_UPSTREAM_LATENCY";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_DOWNSTREAM_LATENCY] = "TFE_CMSG_SSL_DOWNSTREAM_LATENCY";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_UPSTREAM_VERSION] = "TFE_CMSG_SSL_UPSTREAM_VERSION";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_DOWNSTREAM_VERSION] = "TFE_CMSG_SSL_DOWNSTREAM_VERSION";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_PINNING_STATE] = "TFE_CMSG_SSL_PINNING_STATE";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_CERT_VERIFY] = "TFE_CMSG_SSL_CERT_VERIFY";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_ERROR] = "TFE_CMSG_SSL_ERROR";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_MAC] = "TFE_CMSG_SRC_MAC";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_MAC] = "TFE_CMSG_DST_MAC";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_NODELAY] = "TFE_CMSG_DOWNSTREAM_TCP_NODELAY";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_TTL] = "TFE_CMSG_DOWNSTREAM_TCP_TTL";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_KEEPALIVE] = "TFE_CMSG_DOWNSTREAM_TCP_KEEPALIVE";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_KEEPCNT] = "TFE_CMSG_DOWNSTREAM_TCP_KEEPCNT";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_KEEPIDLE] = "TFE_CMSG_DOWNSTREAM_TCP_KEEPIDLE";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_KEEPINTVL] = "TFE_CMSG_DOWNSTREAM_TCP_KEEPINTVL";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DOWNSTREAM_TCP_USER_TIMEOUT] = "TFE_CMSG_DOWNSTREAM_TCP_USER_TIMEOUT";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_NODELAY] = "TFE_CMSG_UPSTREAM_TCP_NODELAY";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_TTL] = "TFE_CMSG_UPSTREAM_TCP_TTL";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_KEEPALIVE] = "TFE_CMSG_UPSTREAM_TCP_KEEPALIVE";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_KEEPCNT] = "TFE_CMSG_UPSTREAM_TCP_KEEPCNT";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_KEEPIDLE] = "TFE_CMSG_UPSTREAM_TCP_KEEPIDLE";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_KEEPINTVL] = "TFE_CMSG_UPSTREAM_TCP_KEEPINTVL";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_UPSTREAM_TCP_USER_TIMEOUT] = "TFE_CMSG_UPSTREAM_TCP_USER_TIMEOUT";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_TCP_PASSTHROUGH] = "TFE_CMSG_TCP_PASSTHROUGH";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_SUB_ID] = "TFE_CMSG_SRC_SUB_ID";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_SUB_ID] = "TFE_CMSG_DST_SUB_ID";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_ASN] = "TFE_CMSG_SRC_ASN";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_ASN] = "TFE_CMSG_DST_ASN";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_ORGANIZATION] = "TFE_CMSG_SRC_ORGANIZATION";
    tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_ORGANIZATION] = "TFE_CMSG_DST_ORGANIZATION";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_IP_LOCATION_COUNTRY] = "TFE_CMSG_SRC_IP_LOCATION_COUNTRY";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_IP_LOCATION_COUNTRY] = "TFE_CMSG_DST_IP_LOCATION_COUNTRY";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_IP_LOCATION_PROVINE] = "TFE_CMSG_SRC_IP_LOCATION_PROVINE";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_IP_LOCATION_PROVINE] = "TFE_CMSG_DST_IP_LOCATION_PROVINE";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_IP_LOCATION_CITY] = "TFE_CMSG_SRC_IP_LOCATION_CITY";
    // tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_IP_LOCATION_CITY] = "TFE_CMSG_DST_IP_LOCATION_CITY";
	// tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_IP_LOCATION_SUBDIVISION] = "TFE_CMSG_SRC_IP_LOCATION_SUBDIVISION";
	// tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_IP_LOCATION_SUBDIVISION] = "TFE_CMSG_DST_IP_LOCATION_SUBDIVISION";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_SRC_IP_LOCATION] = "TFE_CMSG_SRC_IP_LOCATION";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_DST_IP_LOCATION] = "TFE_CMSG_DST_IP_LOCATION";

    tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_CLIENT_JA3_FINGERPRINT] = "TFE_CMSG_SSL_CLIENT_JA3_FINGERPRINT";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_FQDN_CAT_ID_NUM] = "TFE_CMSG_FQDN_CAT_ID_NUM";
	tfe_cmsg_tlv_type_to_string[TFE_CMSG_FQDN_CAT_ID_VAL] = "TFE_CMSG_FQDN_CAT_ID_VAL";

	tfe_cmsg_tlv_type_to_string[TFE_CMSG_COMMON_DIRECTION] = "TFE_CMSG_COMMON_DIRECTION";

	tfe_cmsg_tlv_type_to_string[TFE_CMSG_SSL_PASSTHROUGH_REASON] = "TFE_CMSG_SSL_PASSTHROUGH_REASON";
}