#include "tfe_types.h"
#include "tfe_utils.h"
#include <string.h>
#include <MESA/Maat_rule.h>

const char * tfe_stream_conn_dir_to_str(enum tfe_conn_dir dir)
{
	return (dir == CONN_DIR_DOWNSTREAM) ? "downstream" : "upstream";
}
int tfe_stream_addr_str_split(char* addr_str, const char** sip, const char** sport, const char** dip, const char** dport)
{
	const char* seps=" ";
	char* saveptr=NULL, *subtoken=NULL, *str=NULL;
	int i=0;
	for (str = addr_str, i=0; ; str = NULL, i++)
	{
		subtoken = strtok_r(str, seps, &saveptr);
		if (subtoken == NULL)
			break;
		switch(i)
		{
			case 0:
				if(sip!=NULL) *sip=subtoken;
				break;
			case 1:
				if(sport!=NULL) *sport=subtoken;
				break;
			case 2:
				if(dip!=NULL) *dip=subtoken;
				break;
			case 3:
				if(dport!=NULL)	*dport=subtoken;
				break;
			default:
				return -1;
				break;
		}
	}
	return 0;

}
char * tfe_stream_addr_to_str(const struct tfe_stream_addr * addr)
{
	char * __str_ret = NULL;
	if(addr==NULL) return NULL;
	if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
	{
		const struct tfe_stream_addr_tuple4_v4 * tuple4_v4 = addr->tuple4_v4;
		char __src_addr[INET_ADDRSTRLEN];
		char __dst_addr[INET_ADDRSTRLEN];
		uint16_t __src_port = ntohs((uint16_t) tuple4_v4->source);
		uint16_t __dst_port = ntohs((uint16_t) tuple4_v4->dest);

		inet_ntop(AF_INET, &tuple4_v4->saddr, __src_addr, sizeof(__src_addr));
		inet_ntop(AF_INET, &tuple4_v4->daddr, __dst_addr, sizeof(__dst_addr));
		asprintf(&__str_ret, "%s %u %s %u", __src_addr, __src_port, __dst_addr, __dst_port);
	}

	if (addr->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
	{
		const struct tfe_stream_addr_tuple4_v6 * tuple4_v6 = addr->tuple4_v6;
		char __src_addr[INET6_ADDRSTRLEN];
		char __dst_addr[INET6_ADDRSTRLEN];
		uint16_t __src_port = ntohs((uint16_t) tuple4_v6->source);
		uint16_t __dst_port = ntohs((uint16_t) tuple4_v6->dest);

		inet_ntop(AF_INET6, &tuple4_v6->saddr, __src_addr, sizeof(__src_addr));
		inet_ntop(AF_INET6, &tuple4_v6->daddr, __dst_addr, sizeof(__dst_addr));
		asprintf(&__str_ret, "%s %u %s %u", __src_addr, __src_port, __dst_addr, __dst_port);
	}

	return __str_ret;
}
void tfe_stream_addr_free(struct tfe_stream_addr * addr)
{
	free(addr);
}

struct tfe_stream_addr * tfe_stream_addr_reverse(struct tfe_stream_addr * orin)
{
    struct tfe_stream_addr * reverse_addr = NULL;
    if(orin->addrtype == TFE_ADDR_STREAM_TUPLE4_V4)
    {
        reverse_addr = (struct tfe_stream_addr *) malloc(sizeof(struct tfe_stream_addr) +
            sizeof(struct tfe_stream_addr_tuple4_v4));

        reverse_addr->addrtype = TFE_ADDR_STREAM_TUPLE4_V4;
        reverse_addr->addrlen = sizeof(struct tfe_stream_addr_tuple4_v4);
        reverse_addr->tuple4_v4->saddr = orin->tuple4_v4->daddr;
        reverse_addr->tuple4_v4->daddr = orin->tuple4_v4->saddr;
        reverse_addr->tuple4_v4->source = orin->tuple4_v4->dest;
        reverse_addr->tuple4_v4->dest = orin->tuple4_v4->source;
    }
    else if (orin->addrtype == TFE_ADDR_STREAM_TUPLE4_V6)
    {
        reverse_addr = (struct tfe_stream_addr *) malloc(sizeof(struct tfe_stream_addr) +
            sizeof(struct tfe_stream_addr_tuple4_v6));

        reverse_addr->addrtype = TFE_ADDR_STREAM_TUPLE4_V6;
        reverse_addr->addrlen = sizeof(struct tfe_stream_addr_tuple4_v6);
        reverse_addr->tuple4_v6->saddr = orin->tuple4_v6->daddr;
        reverse_addr->tuple4_v6->daddr = orin->tuple4_v6->saddr;
        reverse_addr->tuple4_v6->source = orin->tuple4_v6->dest;
        reverse_addr->tuple4_v6->dest = orin->tuple4_v6->source;

        return reverse_addr;
    }

    assert(reverse_addr != NULL);
    return reverse_addr;
}

struct tfe_stream_addr * tfe_stream_addr_create_by_fd(int fd, enum tfe_conn_dir dir)
{
	struct tfe_stream_addr * __stream_addr = NULL;

	struct sockaddr_storage sk_src_storage{};
	struct sockaddr * sk_src_ptr = (struct sockaddr *) &sk_src_storage;
	socklen_t sk_src_len = sizeof(sk_src_storage);

	struct sockaddr_storage sk_dst_storage{};
	struct sockaddr * sk_dst_ptr = (struct sockaddr *) &sk_dst_storage;
	socklen_t sk_dst_len = sizeof(sk_dst_storage);

	switch(dir)
	{
		case CONN_DIR_UPSTREAM:
		{
			int ret = getsockname(fd, sk_src_ptr, &sk_src_len);
			if (ret < 0) goto __errout;

			ret = getpeername(fd, sk_dst_ptr, &sk_dst_len);
			if (ret < 0) goto __errout;
			break;
		}
		case CONN_DIR_DOWNSTREAM:
		{
			int ret = getsockname(fd, sk_dst_ptr, &sk_dst_len);
			if (ret < 0) goto __errout;

			ret = getpeername(fd, sk_src_ptr, &sk_src_len);
			if (ret < 0) goto __errout;
			break;
		}
		default: { assert(0); goto __errout;}
	}

	assert(sk_src_ptr->sa_family == sk_dst_ptr->sa_family);
	if (sk_src_ptr->sa_family == AF_INET)
	{
		__stream_addr = (struct tfe_stream_addr *) malloc(sizeof(struct tfe_stream_addr) +
			sizeof(struct tfe_stream_addr_tuple4_v4));

		struct tfe_stream_addr_ipv4 * st_addr_v4 = __stream_addr->ipv4;
		struct sockaddr_in * sk_v4_src_ptr = (struct sockaddr_in *) sk_src_ptr;
		struct sockaddr_in * sk_v4_dst_ptr = (struct sockaddr_in *) sk_dst_ptr;

		__stream_addr->addrtype = TFE_ADDR_STREAM_TUPLE4_V4;
		__stream_addr->addrlen = sizeof(struct tfe_stream_addr_tuple4_v4);

		st_addr_v4->saddr.s_addr = sk_v4_src_ptr->sin_addr.s_addr;
		st_addr_v4->source = sk_v4_src_ptr->sin_port;
		st_addr_v4->daddr.s_addr = sk_v4_dst_ptr->sin_addr.s_addr;
		st_addr_v4->dest = sk_v4_dst_ptr->sin_port;
	}
	else if (sk_src_ptr->sa_family == AF_INET6)
	{
		__stream_addr = (struct tfe_stream_addr *) malloc(sizeof(struct tfe_stream_addr) +
			sizeof(struct tfe_stream_addr_tuple4_v6));

		struct tfe_stream_addr_ipv6 * st_addr_v6 = __stream_addr->ipv6;
		struct sockaddr_in6 * sk_v6_src_ptr = (struct sockaddr_in6 *) sk_src_ptr;
		struct sockaddr_in6 * sk_v6_dst_ptr = (struct sockaddr_in6 *) sk_dst_ptr;

		__stream_addr->addrtype = TFE_ADDR_STREAM_TUPLE4_V6;
		__stream_addr->addrlen = sizeof(struct tfe_stream_addr_tuple4_v4);

		st_addr_v6->saddr = sk_v6_src_ptr->sin6_addr;
		st_addr_v6->source = sk_v6_src_ptr->sin6_port;
		st_addr_v6->daddr = sk_v6_dst_ptr->sin6_addr;
		st_addr_v6->dest = sk_v6_dst_ptr->sin6_port;
	}
	else { assert(0); goto __errout; }
	return __stream_addr;

__errout:
	if (__stream_addr != NULL) free(__stream_addr);
	return NULL;
}

char * tfe_string_addr_create_by_fd(int fd, enum tfe_conn_dir dir)
{
	char * addr_str = NULL;
	struct tfe_stream_addr * stream_addr = tfe_stream_addr_create_by_fd(fd, dir);
	if (stream_addr)
	{
		addr_str = tfe_stream_addr_to_str(stream_addr);
	}
	else
	{
		addr_str = tfe_strdup("null");
	}
	tfe_stream_addr_free(stream_addr);
	return addr_str;
}