#include "decode_ipv4.h"
#include "decode_ipv6.h"
#include "decode_tcp.h"
#include "decode_udp.h"
#include "decode_gtp.h"
#include "inject_pkt.h"
#include "system.h"

#include <linux/netfilter.h> // for NF_ACCEPT
#include <libnetfilter_queue/libnetfilter_queue.h>

#ifdef Packet_Adapter_GIT_VERSION
static __attribute__((__used__)) const char *Packet_Adapter_Version = GIT_VERSION;
#else
static __attribute__((__used__)) const char *Packet_Adapter_Version = "Unknown";
#endif

typedef struct pkt_info_s
{
	uint32_t id;	   // unique ID of packet in queue
	uint16_t protocol; // hw protocol
	uint8_t hook;	   // netfilter hook
	u_int32_t mark;
	u_int32_t indev;
	u_int32_t outdev;
	u_int32_t phys_indev;
	u_int32_t phys_outdev;

	uint8_t *payload;
	uint32_t payload_len;

	char src_addr[512];
} pkt_info_t;

typedef struct union_info_s
{
	ipv4_info_t ipv4;
	ipv6_info_t ipv6;
	tcp_info_t tcp;
	udp_info_t udp;
} union_info_t;

typedef struct pkt_paser_s
{
	pkt_info_t raw;
	union_info_t external;
	gtp_info_t gtp;
	union_info_t internal;
} pkt_paser_t;

static void dump_info(pkt_paser_t *parser)
{
	uint32_t pkt_id = parser->raw.id;
	LOG_DEBUG("raw: {id: %u, protocol: %u, hook: %u, mark: %u, indev: %u, outdev: %u, phys_indev: %u, phys_outdev: %u, src_addr: %s, data_len: %u}",
			  parser->raw.id,
			  parser->raw.protocol,
			  parser->raw.hook,
			  parser->raw.mark,
			  parser->raw.indev,
			  parser->raw.outdev,
			  parser->raw.phys_indev,
			  parser->raw.phys_outdev,
			  parser->raw.src_addr,
			  parser->raw.payload_len);

	// external
	if (parser->external.ipv4.hdr)
	{
		dump_ipv4_info(pkt_id, &(parser->external.ipv4));
	}
	if (parser->external.ipv6.hdr)
	{
		dump_ipv6_info(pkt_id, &(parser->external.ipv6));
	}
	if (parser->external.udp.hdr)
	{
		dump_udp_info(pkt_id, &(parser->external.udp));
	}
	if (parser->external.tcp.hdr)
	{
		dump_tcp_info(pkt_id, &(parser->external.tcp));
	}

	// gtp
	if (parser->gtp.hdr)
	{
		dump_gtp_info(pkt_id, &(parser->gtp));
	}

	// internal
	if (parser->internal.ipv4.hdr)
	{
		dump_ipv4_info(pkt_id, &(parser->internal.ipv4));
	}
	if (parser->internal.ipv6.hdr)
	{
		dump_ipv6_info(pkt_id, &(parser->internal.ipv6));
	}
	if (parser->internal.udp.hdr)
	{
		dump_udp_info(pkt_id, &(parser->internal.udp));
	}
	if (parser->internal.tcp.hdr)
	{
		dump_tcp_info(pkt_id, &(parser->internal.tcp));
	}
}

static int decode_ip_tcp_udp(union_info_t *parser, const uint8_t *data, uint32_t len)
{
	int next_protocol = 0;
	uint8_t *payload = NULL;
	uint32_t payload_len = 0;

	if (len < IPV4_HEADER_LEN)
	{
		LOG_ERROR("Parser IP header: packet length too small %d", len);
		return -1;
	}

	if (IP_GET_RAW_VER(data) == 4)
	{
		if (decode_ipv4(&(parser->ipv4), data, len) == -1)
		{
			return -1;
		}

		payload = parser->ipv4.payload;
		payload_len = parser->ipv4.payload_len;
		next_protocol = parser->ipv4.next_protocol;
	}
	else if (IP_GET_RAW_VER(data) == 6)
	{
		if (decode_ipv6(&(parser->ipv6), data, len) == -1)
		{
			return -1;
		}
		payload = parser->ipv6.payload;
		payload_len = parser->ipv6.payload_len;
		next_protocol = parser->ipv6.next_protocol;
	}
	else
	{
		LOG_ERROR("Unknown IP version %d", IP_GET_RAW_VER(data));
		return -1;
	}

	if (next_protocol == IPPROTO_UDP)
	{
		if (decode_udp(&(parser->udp), payload, payload_len) == -1)
		{
			return -1;
		}
		return 0;
	}
	else if (next_protocol == IPPROTO_TCP)
	{
		if (decode_tcp(&(parser->tcp), payload, payload_len) == -1)
		{
			return -1;
		}
		return 0;
	}
	else
	{
		LOG_ERROR("Unknown Internal L4 next_protocol version %d", next_protocol);
		return -1;
	}
}

///////////////////////////////////////////////////////////////////////////////
// NFQ API
///////////////////////////////////////////////////////////////////////////////

static int decode_pkt(pkt_info_t *packet, struct nfgenmsg *nfmsg, struct nfq_data *nfa)
{
	struct nfqnl_msg_packet_hdr *packet_hdr = NULL;
	struct nfqnl_msg_packet_hw *packet_hw = NULL;

	packet_hdr = nfq_get_msg_packet_hdr(nfa);
	if (packet_hdr == NULL)
	{
		LOG_ERROR("Failed at nfq_get_msg_packet_hdr()");
		return 0;
	}
	packet->id = ntohl(packet_hdr->packet_id);

	packet->payload_len = nfq_get_payload(nfa, &packet->payload);
	if (packet->payload_len <= 0)
	{
		LOG_ERROR("Failed at nfq_get_payload()");
		return packet->id;
	}
	packet->protocol = ntohs(packet_hdr->hw_protocol);
	packet->hook = packet_hdr->hook;

	packet_hw = nfq_get_packet_hw(nfa);
	if (packet_hw)
	{
		int i = 0;
		int offset = 0;
		int len = sizeof(packet->src_addr);
		int hlen = ntohs(packet_hw->hw_addrlen);

		for (i = 0; i < hlen - 1; i++)
		{
			offset += snprintf(packet->src_addr + offset, len - offset, "%02x:", packet_hw->hw_addr[i]);
		}
		snprintf(packet->src_addr + offset, len - offset, "%02x", packet_hw->hw_addr[hlen - 1]);
	}

	packet->mark = nfq_get_nfmark(nfa);
	packet->indev = nfq_get_indev(nfa);
	packet->outdev = nfq_get_outdev(nfa);
	packet->phys_indev = nfq_get_physindev(nfa);
	packet->phys_outdev = nfq_get_physoutdev(nfa);

	return packet->id;
}
/*
 * nfmsg : message objetc that contains the packet
 * nfa : Netlink packet data handle
 */
static int packet_handler_cb(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg, struct nfq_data *nfa, void *data)
{
	int offest = 0;
	pkt_paser_t parser = {0};
	int packet_id = decode_pkt(&(parser.raw), nfmsg, nfa);

	// external
	if (decode_ip_tcp_udp(&(parser.external), parser.raw.payload, parser.raw.payload_len) == -1)
	{
		goto end;
	}

	if (parser.external.udp.hdr == NULL)
	{
		LOG_ERROR("External L4 protocol not UDP");
		goto end;
	}

	// decode GTP
	if (decode_gtp(&(parser.gtp), parser.external.udp.payload, parser.external.udp.payload_len) == -1)
	{
		return -1;
	}

	// internal
	if (decode_ip_tcp_udp(&(parser.internal), parser.gtp.payload, parser.gtp.payload_len) == -1)
	{
		goto end;
	}

	/*
	 * NF_DROP : discarded the packet
	 * NF_ACCEPT : the packet passes, continue iterations
	 * NF_QUEUE : inject the packet into a different queue (the target queue number is in the high 16 bits of the verdict)
	 * NF_REPEAT : iterate the same cycle once more
	 * NF_STOP : accept, but don't continue iterations
	 */
	// nfq_set_verdict()
	// nfq_set_verdict2()
	// nfq_set_verdict_batch()
	// nfq_set_verdict_batch2()
	// nfq_set_verdict_mark()

	if (parser.external.ipv4.hdr)
	{
		offest += parser.external.ipv4.hdr_len;
	}
	if (parser.external.ipv6.hdr)
	{
		offest += parser.external.ipv6.hdr_len;
	}

	offest += parser.external.udp.hdr_len;
	offest += parser.gtp.hdr_len;

	dump_info(&parser);
	LOG_DEBUG("Offset : %d", offest);

	uint8_t *inject_data = parser.raw.payload + offest;
	uint32_t inject_data_len = parser.raw.payload_len - offest;

	if (offest > 0)
	{
		if ((parser.external.ipv4.hdr && parser.internal.ipv4.hdr) || (parser.external.ipv6.hdr && parser.internal.ipv6.hdr))
		{
			return nfq_set_verdict(qh, packet_id, NF_ACCEPT, inject_data_len, inject_data);
		}

		if (parser.external.ipv4.hdr && parser.internal.ipv6.hdr)
		{
			if (inject_ipv6_pkt(parser.internal.ipv6.dst_addr, inject_data, inject_data_len) == -1)
			{
				goto end;
			}
			return nfq_set_verdict(qh, packet_id, NF_DROP, 0, NULL);
		}

		if (parser.external.ipv6.hdr && parser.internal.ipv4.hdr)
		{
			if (inject_ipv4_pkt(parser.internal.ipv4.dst_addr, inject_data, inject_data_len) == -1)
			{
				goto end;
			}
			return nfq_set_verdict(qh, packet_id, NF_DROP, 0, NULL);
		}
	}

end:
	return nfq_set_verdict(qh, packet_id, NF_ACCEPT, 0, NULL);
}

static void usage(char *cmd)
{
	fprintf(stderr, "USAGE: %s [OPTIONS]\n", cmd);
	fprintf(stderr, "    -v       -- show version\n");
	fprintf(stderr, "    -i id    -- set queue id\n");
	fprintf(stderr, "    -d       -- run daemon\n");
	fprintf(stderr, "    -h       -- show help\n");
}

/*
 * doc : http://www.netfilter.org/projects/libnetfilter_queue/doxygen/html/
 * Library setup : http://www.netfilter.org/projects/libnetfilter_queue/doxygen/html/group__LibrarySetup.html
 * Queue handling : http://www.netfilter.org/projects/libnetfilter_queue/doxygen/html/group__Queue.html
 * Message parsing : http://www.netfilter.org/projects/libnetfilter_queue/doxygen/html/group__Parsing.html
 */
int main(int argc, char **argv)
{
	int fd;
	int rv;
	int opt;
	uint16_t queue = 1;
	struct nfq_handle *handle;
	struct nfq_q_handle *q_handle;
	char buf[65535] __attribute__((aligned));

	while ((opt = getopt(argc, argv, "vi:dh")) != -1)
	{
		switch (opt)
		{
		case 'v':
			fprintf(stderr, "Packet Adapter Version: %s\n", Packet_Adapter_Version);
			return 0;
		case 'i':
			queue = atoi(optarg);
			if (queue < 0 || queue > 65535)
			{
				fprintf(stderr, "Usage: %s queueid %d out of range [0, 65535]\n", argv[0], queue);
				return 0;
			}
			break;
		case 'd':
			run_daemon();
			break;
		case 'h': /* fall through */
		default:
			usage(argv[0]);
			return 0;
		}
	}

	LOG_DEBUG("Using queue: %d", queue);

	handle = nfq_open();
	if (handle == NULL)
	{
		LOG_ERROR("Failed at nfq_open(), %d: %s", errno, strerror(errno));
		goto error;
	}

	if (nfq_unbind_pf(handle, AF_INET) < 0)
	{
		LOG_ERROR("Failed at nfq_unbind_pf(), %d: %s", errno, strerror(errno));
		goto error;
	}

	if (nfq_bind_pf(handle, AF_INET) < 0)
	{
		LOG_ERROR("Failed at nfq_bind_pf(), %d: %s", errno, strerror(errno));
		goto error;
	}

	q_handle = nfq_create_queue(handle, queue, &packet_handler_cb, NULL);
	if (q_handle == NULL)
	{
		LOG_ERROR("Failed at nfq_create_queue(), %d: %s", errno, strerror(errno));
		goto error;
	}

	/*
	 * NFQNL_COPY_NONE - noop, do not use it
	 * NFQNL_COPY_META - copy only packet metadata
	 * NFQNL_COPY_PACKET - copy entire packet
	 */
	if (nfq_set_mode(q_handle, NFQNL_COPY_PACKET, 0xffff) < 0)
	{
		LOG_ERROR("Failed at nfq_set_mode(NFQNL_COPY_PACKET), %d: %s", errno, strerror(errno));
		goto error;
	}

	if (nfq_set_queue_maxlen(q_handle, 65535) < 0)
	{
		LOG_ERROR("Failed at nfq_set_queue_maxlen(65535), %d: %s", errno, strerror(errno));
		goto error;
	}

	LOG_DEBUG("Waiting for packets...");

	fd = nfq_fd(handle);
	for (;;)
	{
		if ((rv = recv(fd, buf, sizeof(buf), 0)) >= 0)
		{
			nfq_handle_packet(handle, buf, rv);
			continue;
		}
		/* 
		 * if your application is too slow to digest the packets that
		 * are sent from kernel-space, the socket buffer that we use
		 * to enqueue packets may fill up returning ENOBUFS. Depending
		 * on your application, this error may be ignored. Please, see
		 * the doxygen documentation of this library on how to improve
		 * this situation.
		 */
		if (rv < 0 && errno == ENOBUFS)
		{
			LOG_ERROR("Losing packets !!!");
			continue;
		}

		LOG_ERROR("Failed at recv(), %d: %s", errno, strerror(errno));
	}

error:
	if (q_handle)
	{
		nfq_destroy_queue(q_handle);
	}

	if (handle)
	{
		nfq_close(handle);
	}

	return 0;
}