#include <string.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>
#include <linux/tcp.h>

#include <tfe_utils.h>
#include <tfe_pkt_util.h>

// always success
void tfe_pkt_parse_ipv4_header(const void *a_packet, struct pkt_info *pktinfo)
{
	pktinfo->addr_type = ADDR_TYPE_IPV4;

	pktinfo->iphdr.v4 = (struct iphdr *)a_packet;
	pktinfo->iphdr_len = pktinfo->iphdr.v4->ihl * 4;
	pktinfo->ip_totlen = ntohs(pktinfo->iphdr.v4->tot_len);

	pktinfo->tcphdr = (struct tcphdr *)((char *)pktinfo->iphdr.v4 + pktinfo->iphdr_len);
	pktinfo->tcphdr_len = pktinfo->tcphdr->doff * 4;

	pktinfo->data = (char *)pktinfo->tcphdr + pktinfo->tcphdr_len;
	pktinfo->data_len = pktinfo->ip_totlen - pktinfo->iphdr_len - pktinfo->tcphdr_len;

	pktinfo->parse_failed = 0;
}

// check pktinfo->parse_failed for status
void tfe_pkt_parse_ipv6_header(const void *a_packet, struct pkt_info *pktinfo)
{
	pktinfo->addr_type = ADDR_TYPE_IPV6;

	pktinfo->iphdr.v6 = (struct ip6_hdr *)a_packet;
	pktinfo->ip_totlen = ntohs(pktinfo->iphdr.v6->ip6_ctlun.ip6_un1.ip6_un1_plen) + sizeof(struct ip6_hdr);

	uint8_t next_hdr_type = pktinfo->iphdr.v6->ip6_ctlun.ip6_un1.ip6_un1_nxt;
	char *next_hdr_ptr = (char *)pktinfo->iphdr.v6 + sizeof(struct ip6_hdr);
	int skip_len = 0;
	while (1)
	{
		switch(next_hdr_type)
		{
			case IPPROTO_TCP:
				//parse tcphdr
				pktinfo->iphdr_len = next_hdr_ptr - (char *)a_packet;

				pktinfo->tcphdr = (struct tcphdr *)next_hdr_ptr;
				pktinfo->tcphdr_len = pktinfo->tcphdr->doff * 4;

				pktinfo->data = (char *)pktinfo->tcphdr + pktinfo->tcphdr_len;
				pktinfo->data_len = pktinfo->ip_totlen - pktinfo->iphdr_len - pktinfo->tcphdr_len;

				pktinfo->parse_failed = 0;
				return;
			case IPPROTO_HOPOPTS: /* fall through */
			case IPPROTO_ROUTING: /* fall through */
			case IPPROTO_AH: /* fall through */
			case IPPROTO_DSTOPTS: /* fall through */
				skip_len = (*(next_hdr_ptr + 1)) * 8 + 8;
				next_hdr_type = *next_hdr_ptr;
				next_hdr_ptr += skip_len;
				break;
			case IPPROTO_NONE: /* fall through */
			default:
				pktinfo->parse_failed = 1;
				return;
		}
	}
}

uint16_t tfe_pkt_checksum_ip(const void *buf, size_t hdr_len)
{
	unsigned long sum = 0;
	const uint16_t *ip1;
	ip1 = (const uint16_t *)buf;
	while (hdr_len > 1)
	{
		sum += *ip1++;
		if (sum & 0x80000000)
		{
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		hdr_len -= 2;
	}
	while (sum >> 16)
	{
		sum = (sum & 0xFFFF) + (sum >> 16);
	}
	return (~sum);
}

uint16_t tfe_pkt_checksum_tcp_v4(const void *_buf, size_t len, in_addr_t src_addr, in_addr_t dest_addr)
{
	const uint16_t *buf = (u_int16_t *)_buf;
	uint16_t *ip_src = (uint16_t *)&src_addr;
	uint16_t *ip_dst = (uint16_t *)&dest_addr;
	uint32_t sum;
	size_t length = len;

	// Calculate the sum
	sum = 0;
	while (len > 1)
	{
		sum += *buf++;
		if (sum & 0x80000000)
		{
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		len -= 2;
	}

	if (len & 1)
	{
		// Add the padding if the packet lenght is odd
		sum += *((uint8_t *)buf);
	}

	// Add the pseudo-header
	sum += *(ip_src++);
	sum += *ip_src;
	sum += *(ip_dst++);
	sum += *ip_dst;
	sum += htons(IPPROTO_TCP);
	sum += htons(length);

	// Add the carries
	while (sum >> 16)
	{
		sum = (sum & 0xFFFF) + (sum >> 16);
	}

	// Return the one's complement of sum
	return ((uint16_t)(~sum));
}

uint16_t tfe_pkt_checksum_tcp_v6(const void *_buf, size_t len, struct in6_addr src_addr, struct in6_addr dest_addr)
{
	const uint16_t *buf = (u_int16_t *)_buf;
	uint16_t *ip_src = (uint16_t *)&src_addr;
	uint16_t *ip_dst = (uint16_t *)&dest_addr;
	uint32_t sum;
	size_t length = len;
	int i = 0;

	// Calculate the sum
	sum = 0;
	while (len > 1)
	{
		sum += *buf++;
		if (sum & 0x80000000)
		{
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		len -= 2;
	}

	if (len & 1)
	{
		// Add the padding if the packet lenght is odd
		sum += *((uint8_t *)buf);
	}

	// Add the pseudo-header
	for (i = 0; i < 8; i++)
	{
		sum += *ip_src;
		ip_src++;
	}
	for ( i = 0; i < 8; i++)
	{
		sum += *ip_dst;
		ip_dst++;
	}
	sum += htons(IPPROTO_TCP);
	sum += htons(length);

	// Add the carries
	while (sum >> 16)
	{
		sum = (sum & 0xFFFF) + (sum >> 16);
	}

	// Return the one's complement of sum
	return ((uint16_t)(~sum));
}

/*
 * 目的：在 IP 的 Payload ${data} 中查找指定的 tcp ${option}。
 *
 * 已知：
 * 1.所有的 tcp options 所占的存储空间的长度为 ${optlen}
 * 2.${out_opt_buff} 输出缓冲区的最大长度为 ${out_opt_buff_size}
 *
 * 返回值：
 * 1.若找到指定的 tcp ${option} 则返回 1，并将该 tcp ${option} 对应的值拷贝到 ${out_opt_buff} 中，并将拷贝的值所占的存储空间记录到 ${out_optlen} 中
 * 2.若未找到指定的 tcp ${option} 则返回 0
 */
int tfe_pkt_find_tcp_option(uint8_t option, char *data, unsigned int opts_total_len, uint8_t *out_opt_len, char *out_opt_buff, unsigned int out_opt_buff_size)
{
	const uint8_t *op;
	unsigned int i;

	if (!opts_total_len)
		return 0;

	op = (uint8_t *)(data + sizeof(struct tcphdr));
	for (i = 0; i < opts_total_len;)
	{
		uint8_t __optlen;
		uint8_t __valuelen;

		if (op[i] == option)
		{
			if (op[i] < 2)
			{
				*out_opt_len = 0;
				return 1;
			}

			__optlen = op[i + 1];
			if (__optlen <= 2)
			{
				TFE_LOG_ERROR(g_default_logger, "failed at parse tcp options, tcp option length must be larger than 2, but the value is %u", __optlen);
				return 0;
			}

			__valuelen = __optlen - 2;
			if (__valuelen > out_opt_buff_size)
			{
				TFE_LOG_ERROR(g_default_logger, "failed at parse tcp options, tcp option length is larger than input buffer");
				return 0;
			}

			*out_opt_len = __valuelen;
			memcpy(out_opt_buff, &op[i + 2], __valuelen);
			return 1;
		}

		if (op[i] < 2)
			i++;
		else
			i += op[i + 1] ?: 1;
	}

	return 0;
}