#define _GNU_SOURCE
#include <arpa/inet.h>
#include <ctype.h>
#include <errno.h>
#include <getopt.h>
#include <linux/if_ether.h>
#include <linux/if_packet.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <locale.h>
#include <net/if.h>
#include <netinet/in.h>
#include <sched.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>
#include <wchar.h>

#ifdef UNICODE
#define ok(fmt, ...) wprintf(L"\e[32m\e[0m " fmt, ##__VA_ARGS__)
#define warn(fmt, ...) wprintf(L"\e[33m\e[0m " fmt, ##__VA_ARGS__)
#define incoming(fmt, ...) wprintf(L"\e[36m\e[0m " fmt, ##__VA_ARGS__)
#define error(fmt, ...) wprintf(L"\e[31m\e[0m " fmt, ##__VA_ARGS__)
#define UNKNOWN L"\e[90m�\e[0m"
#else
#define ok(fmt, ...) wprintf(L"\e[32m[Y]\e[0m " fmt, ##__VA_ARGS__)
#define warn(fmt, ...) wprintf(L"\e[33m[W]\e[0m " fmt, ##__VA_ARGS__)
#define incoming(fmt, ...) wprintf(L"\e[36m[I]\e[0m " fmt, ##__VA_ARGS__)
#define error(fmt, ...) wprintf(L"\e[31m[E]\e[0m " fmt, ##__VA_ARGS__)
#define UNKNOWN L"\e[90m?\e[0m"
#endif

#ifndef PAGE_SIZE
#define PAGE_SIZE 1UL << 12
#endif

#ifndef DUMPLEN
#define DUMPLEN (16UL)
#endif

struct args {
	const char *iface;
	int fanout_type;
	int cpu;
	int gid;
};

#define err_handler(api) \
	error(L"%s() failed - %s (%d)\n", (api), strerror(errno), errno)

#define MAC_FMT "%02x:%02x:%02x:%02x:%02x:%02x"
#define MAC_VA_LIST(_mac) \
	(_mac)[0], (_mac)[1], (_mac)[2], (_mac)[3], (_mac)[4], (_mac)[5]

volatile sig_atomic_t stop = 0;
static void handle_sigint(int sig __attribute__((unused)))
{
	stop = 1;
}

static void dump_ip_address(char *buf, __be32 addr)
{
	struct sockaddr_in sa;

	sa.sin_addr.s_addr = addr;
	inet_ntop(AF_INET, &(sa.sin_addr), buf, INET_ADDRSTRLEN);
}

static void parse_packet(const char *buf)
{
	size_t i;
	off_t off;

	const struct ethhdr *eth;
	const struct iphdr *iph;
	const struct udphdr *udph;
	const char *payload;

	char src_ip[INET_ADDRSTRLEN];
	char dst_ip[INET_ADDRSTRLEN];

	eth = (const struct ethhdr *)buf;
	off = sizeof(*eth);
	if (ntohs(eth->h_proto) != ETH_P_IP) {
		warn(L"not IPv4 packet, skip\n");
		return;
	}

	iph = (const struct iphdr *)(buf + off);
	off += iph->ihl << 2;
	if (iph->protocol != IPPROTO_UDP) {
		warn(L"not UDP packet, skip\n");
		return;
	}

	udph = (const struct udphdr *)(buf + off);
	off += sizeof(*udph);

	payload = buf + off;

	dump_ip_address(src_ip, iph->saddr);
	dump_ip_address(dst_ip, iph->daddr);

	incoming(L"D(" MAC_FMT ", %s:%hu) - S(" MAC_FMT ", %s:%hu): ",
		 MAC_VA_LIST(eth->h_dest), dst_ip, ntohs(udph->dest),
		 MAC_VA_LIST(eth->h_source), src_ip, ntohs(udph->source));

	for (i = 0; i < DUMPLEN - 1; ++i)
		wprintf(L"%02X ", payload[i] & 0xFF);
	wprintf(L"%02X (", payload[i] & 0xFF);

	for (i = 0; i < DUMPLEN - 1; ++i)
		if (isprint(payload[i]))
			wprintf(L"%c ", payload[i] & 0xFF);
		else
			wprintf(L"%ls ", UNKNOWN);

	if (isprint(payload[i]))
		wprintf(L"%lc)\n", payload[i] & 0xFF);
	else
		wprintf(L"%ls)\n", UNKNOWN);
}

static int setup_signals(void)
{
	struct sigaction sa;

	memset(&sa, 0x0, sizeof(sa));
	sigfillset(&sa.sa_mask);
	sa.sa_handler = handle_sigint;

	return sigaction(SIGINT, &sa, NULL);
}

_Noreturn static void usage(const char *prog, int exitcode)
{
	// clang-format off
	wprintf(L"Usage: %s [OPTIONS] -- <IFACE>\n", prog);
	wprintf(L"Available options\n");
	wprintf(L"  -h [--help]                    Display this message and exits\n");
	wprintf(L"  -f [--fanout-type]   <TYPE>    Which fanout type to use (default: none)\n");
	wprintf(L"  -g [--fanout-group]  <GROUP>   Which fanout group to belong in\n");
	wprintf(L"  -c [--cpu]           <CPU>     The CPU to be bound on\n");
	exit(exitcode);
	// clang-format on
}

#define PACKET_NO_FANOUT (-1)
static void parse_args(int argc, char **argv, struct args *args)
{
	int c;
	char *endptr;
	const char *shortopts = ":hf:g:c:";
	const struct option longopts[] = {
		{
			.name = "help",
			.has_arg = no_argument,
			.flag = NULL,
			.val = 'h',
		},
		{
			.name = "fanout-type",
			.has_arg = required_argument,
			.flag = NULL,
			.val = 'f',
		},
		{
			.name = "fanout-group",
			.has_arg = required_argument,
			.flag = NULL,
			.val = 'g',
		},
		{
			.name = "cpu",
			.has_arg = required_argument,
			.flag = NULL,
			.val = 'c',
		},
	};

	while ((c = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1) {
		switch (c) {
		case 'h':
			usage(argv[0], EXIT_SUCCESS);
		case 'f':
			if (!strcasecmp("hash", optarg)) {
				args->fanout_type = PACKET_FANOUT_HASH;
			} else if (!strcasecmp("lb", optarg)) {
				args->fanout_type = PACKET_FANOUT_LB;
			} else if (!strcasecmp("cpu", optarg)) {
				args->fanout_type = PACKET_FANOUT_CPU;
			} else if (!strcasecmp("none", optarg)) {
				args->fanout_type = PACKET_NO_FANOUT;
			} else {
				error(L"invalid fanout type\n");
				usage(argv[0], EXIT_FAILURE);
			}
			break;
		case 'g':
			args->gid = strtoul(optarg, &endptr, 0);
			if ('\0' != *endptr) {
				error(L"invalid GID\n");
				usage(argv[0], EXIT_FAILURE);
			}
			break;
		case 'c':
			args->cpu = strtoul(optarg, &endptr, 0);
			if ('\0' != *endptr) {
				error(L"invalid CPU\n");
				usage(argv[0], EXIT_FAILURE);
			}
			break;
		case ':':
		case '?':
			error(L"invalid option: %c\n", optopt);
			usage(argv[0], EXIT_FAILURE);
		default:
			error(L"getopt_long error\n");
			usage(argv[0], EXIT_FAILURE);
		}
	}

	if (optind != argc - 1) {
		error(L"missing interface\n");
		usage(argv[0], EXIT_FAILURE);
	}
	args->iface = argv[optind];

	return;
}

static int set_affinity(int cpu_core)
{
	int rv;
	long nproc;
	cpu_set_t mask;

	nproc = sysconf(_SC_NPROCESSORS_ONLN);
	if (nproc == -1) {
		err_handler("sysconf");
		return -1;
	}

	if (cpu_core > nproc) {
		error(L"CPU core must be smaller than %ld\n", nproc);
		return -1;
	}

	CPU_ZERO(&mask);
	CPU_SET(cpu_core, &mask);

	rv = sched_setaffinity(0, sizeof(mask), &mask);
	if (rv) {
		err_handler("sched_setaffinity");
		return -1;
	}

	return 0;
}

static const char *fanout_type_str[] = {
	"hash",
	"lb",
	"cpu",
};

static void dump_args(struct args *args)
{
	ok("settings\n");
	wprintf(L"  - fanout type..: %s (%d)\n",
		fanout_type_str[args->fanout_type], args->fanout_type);
	wprintf(L"  - fanout GID...: %d\n", args->gid);
	if (args->cpu != -1)
		wprintf(L"  pin to CPU...: %d\n", args->cpu);
	wprintf(L"  - interface....: %s\n\n", args->iface);
}

int main(int argc, char **argv)
{
	long rc;
	int sockfd, t;
	struct args args = { .fanout_type = PACKET_NO_FANOUT,
			     .iface = NULL,
			     .cpu = -1,
			     .gid = -1 };

	char buf[PAGE_SIZE];
	struct sockaddr_ll addr;

	setlocale(LC_CTYPE, "");
	parse_args(argc, argv, &args);
	dump_args(&args);

	if (args.cpu != -1) {
		if (set_affinity(args.cpu))
			return EXIT_FAILURE;

		ok(L"affinity set to CPU %d\n", args.cpu);
		if (args.fanout_type != PACKET_FANOUT_CPU)
			warn(L"affinity has been set but fanout type is not %s\n",
			     fanout_type_str[PACKET_FANOUT_CPU]);
	}

	sockfd = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
	if (sockfd == -1) {
		err_handler("socket");
		return EXIT_FAILURE;
	}

	ok(L"listening socket initialized\n");

	memset(&addr, 0x0, sizeof(addr));
	addr.sll_family = AF_PACKET;
	addr.sll_protocol = htons(ETH_P_IP);
	addr.sll_ifindex = if_nametoindex(args.iface);

	rc = bind(sockfd, (const struct sockaddr *)&addr, sizeof(addr));
	if (rc == -1) {
		err_handler("bind");
		goto out;
	}

	ok(L"listening socket binded to %s\n", args.iface);

	if (args.fanout_type == PACKET_NO_FANOUT && args.gid != -1)
		warn(L"GID %d will not be used, since fanout is 'none'\n",
		     args.gid);

	if (args.fanout_type != PACKET_NO_FANOUT) {
		t = (args.gid | (args.fanout_type << 16));
		rc = setsockopt(sockfd, SOL_PACKET, PACKET_FANOUT, &t,
				sizeof(int));
		if (rc == -1) {
			err_handler("setsockopt");
			goto out;
		}

		ok(L"listening socket put into fanout group %d (type: %s)\n",
		   args.gid, fanout_type_str[args.fanout_type]);
	}

	rc = setup_signals();
	if (rc == -1) {
		err_handler("setup_signals");
		goto out;
	}

	ok("you can now hit Ctrl-C to quit\n");

	while (!stop) {
		memset(buf, 0x0, PAGE_SIZE);
		rc = read(sockfd, buf, PAGE_SIZE);
		if (rc == -1) {
			if (errno != EINTR)
				err_handler("read");
			else
				break;
		}

		parse_packet(buf);
	}

out:
	close(sockfd);
	return rc;
}
