#include <stdint.h>
#include <linux/ip.h>
#include <netinet/ip6.h>
#include <linux/netfilter.h>
#include <string.h>
#include <arpa/inet.h>
#include <stdlib.h>

#include <libnetfilter_queue/libnetfilter_queue.h>

#include "mxallowd.h"
#include "log.h"
#include "whitelist.h"

#define NFQ_PACKET_BUFFER_SIZE 4096

/*
 * Helper function to decide what to do with packets sent by iptables
 *
 */
static int handlePacket(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg,
			 struct nfq_data *nfdata, void *packet_buffer) {
	(void)nfmsg;
	(void)packet_buffer;

	struct nfqnl_msg_packet_hdr *header;
	char *payload = NULL;
	int payload_length;

	if (!(header = nfq_get_msg_packet_hdr(nfdata)))
		/* Skip the packet if the header could not be obtained */
		return -1;

	if ((payload_length = nfq_get_payload(nfdata, (char**)&payload)) == -1)
		/* Skip the packet if it has no payload */
		return -1;

	/* OK, we'll shortly need to decide, so let's clean up the whitelists here */
	cleanup_whitelist();

	char dest_address[INET6_ADDRSTRLEN+1],
	     source_address[INET6_ADDRSTRLEN+1];
	memset(dest_address, '\0', INET6_ADDRSTRLEN+1);
	memset(source_address, '\0', INET6_ADDRSTRLEN+1);

	struct iphdr* ip_header = (struct iphdr*)payload;
	if (ip_header->version == IPVERSION) {
		if (inet_ntop(AF_INET, &(ip_header->daddr), dest_address, sizeof(dest_address)) == NULL ||
		    inet_ntop(AF_INET, &(ip_header->saddr), source_address, sizeof(source_address)) == NULL)
			return -1;
	} else {
#ifdef IPV6
		struct ip6_hdr *ip6_header = (struct ip6_hdr*)payload;
		if (inet_ntop(AF_INET6, &(ip6_header->ip6_dst), dest_address, sizeof(dest_address)) == NULL ||
		    inet_ntop(AF_INET6, &(ip6_header->ip6_src), source_address, sizeof(source_address)) == NULL)
			/* Address could not be read, packet has to be malformed */
			return -1;
#endif
	}

	/* Let's see if the packet was sent to MX1 */
	if (is_included(fake_mailservers, dest_address)) {
		/* This packet was sent to MX1, whitelist the sender for MX2 */
		slog("Successful connection from %s to a fake mailserver, adding to whitelist\n", source_address);
		add_to_whitelist(source_address, NULL, false, 0, NULL);
	} else if (is_included(real_mailservers, dest_address)) {
		/* This packet was sent to MX2, let's see what to do */
		if (!is_whitelisted(source_address, true)) {
			/* The sender is not whitelisted, so we drop the packet */
			slog("Dropping connection attempt from %s to a real mailserver\n", source_address);
			blocked_attempts++;
			return nfq_set_verdict(qh, ntohl(header->packet_id), NF_DROP, 0, NULL);
		}
		slog("Successful connection from %s to a real mailserver\n", source_address);
	}
	return nfq_set_verdict(qh, ntohl(header->packet_id), NF_ACCEPT, 0, NULL);
}

void nfq_init() {
	/* Create a libnetfilterqueue-handle */
	struct nfq_handle *nfq = nfq_open();
	if (nfq == NULL)
		diem("Error creating an nfq-handle");

	/* Unbind to clean up previous instances */
#ifdef BROKEN_UNBIND
	(void)nfq_unbind_pf(nfq, AF_INET);
#else
	if (nfq_unbind_pf(nfq, AF_INET) != 0)
		diem("Error unbinding AF_INET");
#endif
	/* Bind to IPv4 */
	if (nfq_bind_pf(nfq, AF_INET) != 0)
		diem("Error binding to AF_INET");

#ifdef IPV6
	/* Unbind to clean up previous instances */
#ifdef BROKEN_UNBIND
	(void)nfq_unbind_pf(nfq, AF_INET6);
#else
	if (nfq_unbind_pf(nfq, AF_INET6) != 0)
		diem("Error unbinding AF_INET6");
#endif
	/* Bind to IPv6 */
	if (nfq_bind_pf(nfq, AF_INET6) != 0)
		diem("Error binding to AF_INET6");
#endif

	/* Create queue */
	struct nfq_q_handle *qh;
	if (!(qh = nfq_create_queue(nfq, queue_num, &handlePacket, NULL)))
		dief("Error creating queue %d\n", queue_num);

	/* We need a copy of the packet */
	if (nfq_set_mode(qh, NFQNL_COPY_PACKET, NFQ_PACKET_BUFFER_SIZE) == -1)
		diem("Cannot set mode to NFQNL_COPY_PACKET");

	char *packet_buffer = malloc(sizeof(unsigned char) * (NFQ_PACKET_BUFFER_SIZE + 1));
	if (packet_buffer == NULL)
		diep("malloc()");

	int fd = nfq_fd(nfq), rv;
	/* Read a packet in blocking mode */
	while ((rv = recv(fd, packet_buffer, NFQ_PACKET_BUFFER_SIZE, 0)) >= 0)
		nfq_handle_packet(nfq, packet_buffer, rv);

	nfq_destroy_queue(qh);
	nfq_close(nfq);
	free(packet_buffer);
}
