/*
 * (C) 2012-2013 by Pablo Neira Ayuso <pablo@netfilter.org>
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This software has been sponsored by Sophos Astaro <http://www.sophos.com>
 */

#define _GNU_SOURCE
#include "config.h"
#include <errno.h>
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <netinet/ether.h>
#include <netinet/in.h>
#include <netinet/ip6.h>
#include <net/if_arp.h>
#include <getopt.h>

#include <sys/socket.h>
#include <arpa/inet.h>

#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>

#include <libmnl/libmnl.h>
#include <libnftnl/table.h>
#include <libnftnl/trace.h>
#include <libnftnl/chain.h>
#include <libnftnl/rule.h>

#include <include/xtables.h>
#include "iptables.h" /* for xtables_globals */
#include "xtables-multi.h"
#include "nft.h"
#include "nft-arp.h"

struct cb_arg {
	uint32_t nfproto;
	bool is_event;
	struct nft_handle *h;
};

static int table_cb(const struct nlmsghdr *nlh, void *data)
{
	uint32_t type = nlh->nlmsg_type & 0xFF;
	const struct cb_arg *arg = data;
	struct nftnl_table *t;
	char buf[4096];

	t = nftnl_table_alloc();
	if (t == NULL)
		goto err;

	if (nftnl_table_nlmsg_parse(nlh, t) < 0)
		goto err_free;

	if (arg->nfproto && arg->nfproto != nftnl_table_get_u32(t, NFTNL_TABLE_FAMILY))
		goto err_free;
	nftnl_table_snprintf(buf, sizeof(buf), t, NFTNL_OUTPUT_DEFAULT, 0);
	printf(" EVENT: ");
	printf("nft: %s table: %s\n", type == NFT_MSG_NEWTABLE ? "NEW" : "DEL", buf);

err_free:
	nftnl_table_free(t);
err:
	return MNL_CB_OK;
}

static bool counters;
static bool trace;
static bool events;

static int rule_cb(const struct nlmsghdr *nlh, void *data)
{
	uint32_t type = nlh->nlmsg_type & 0xFF;
	const struct cb_arg *arg = data;
	struct nftnl_rule *r;
	uint8_t family;

	r = nftnl_rule_alloc();
	if (r == NULL)
		goto err;

	if (nftnl_rule_nlmsg_parse(nlh, r) < 0)
		goto err_free;

	family = nftnl_rule_get_u32(r, NFTNL_RULE_FAMILY);
	if (arg->nfproto && arg->nfproto != family)
		goto err_free;

	if (arg->is_event)
		printf(" EVENT: ");
	switch (family) {
	case AF_INET:
	case AF_INET6:
		printf("-%c ", family == AF_INET ? '4' : '6');
		break;
	case NFPROTO_ARP:
		printf("-0 ");
		break;
	default:
		goto err_free;
	}

	printf("-t %s ", nftnl_rule_get_str(r, NFTNL_RULE_TABLE));
	nft_rule_print_save(arg->h, r, type == NFT_MSG_NEWRULE ? NFT_RULE_APPEND :
							   NFT_RULE_DEL,
			    counters ? 0 : FMT_NOCOUNTS);
err_free:
	nftnl_rule_free(r);
err:
	return MNL_CB_OK;
}

static int chain_cb(const struct nlmsghdr *nlh, void *data)
{
	uint32_t type = nlh->nlmsg_type & 0xFF;
	const struct cb_arg *arg = data;
	struct nftnl_chain *c;
	char buf[4096];
	int family;

	c = nftnl_chain_alloc();
	if (c == NULL)
		goto err;

	if (nftnl_chain_nlmsg_parse(nlh, c) < 0)
		goto err_free;

	family = nftnl_chain_get_u32(c, NFTNL_CHAIN_FAMILY);
	if (arg->nfproto && arg->nfproto != family)
		goto err_free;

	if (nftnl_chain_is_set(c, NFTNL_CHAIN_PRIO))
		family = -1;

	printf(" EVENT: ");
	switch (family) {
	case NFPROTO_IPV4:
		family = 4;
		break;
	case NFPROTO_IPV6:
		family = 6;
		break;
	default:
		nftnl_chain_snprintf(buf, sizeof(buf), c, NFTNL_OUTPUT_DEFAULT, 0);
		printf("# nft: %s\n", buf);
		goto err_free;
	}

	printf("-%d -t %s -%c %s\n",
			family,
			nftnl_chain_get_str(c, NFTNL_CHAIN_TABLE),
			type == NFT_MSG_NEWCHAIN ? 'N' : 'X',
			nftnl_chain_get_str(c, NFTNL_CHAIN_NAME));
err_free:
	nftnl_chain_free(c);
err:
	return MNL_CB_OK;
}

static int newgen_cb(const struct nlmsghdr *nlh, void *data)
{
	uint32_t genid = 0, pid = 0;
	const struct nlattr *attr;
	const char *name = NULL;

	mnl_attr_for_each(attr, nlh, sizeof(struct nfgenmsg)) {
		switch (mnl_attr_get_type(attr)) {
		case NFTA_GEN_ID:
			if (mnl_attr_validate(attr, MNL_TYPE_U32) < 0)
				break;
		        genid = ntohl(mnl_attr_get_u32(attr));
			break;
		case NFTA_GEN_PROC_NAME:
			if (mnl_attr_validate(attr, MNL_TYPE_NUL_STRING) < 0)
				break;
			name = mnl_attr_get_str(attr);
			break;
		case NFTA_GEN_PROC_PID:
			if (mnl_attr_validate(attr, MNL_TYPE_U32) < 0)
				break;
			pid = ntohl(mnl_attr_get_u32(attr));
			break;
		}
	}

	if (name)
		printf("NEWGEN: GENID=%u PID=%u NAME=%s\n", genid, pid, name);

	return MNL_CB_OK;
}

static void trace_print_return(const struct nftnl_trace *nlt)
{
	const char *chain = NULL;

	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_JUMP_TARGET)) {
		chain = nftnl_trace_get_str(nlt, NFTNL_TRACE_JUMP_TARGET);
		printf("%s", chain);
	}
}

static void trace_print_rule(const struct nftnl_trace *nlt, struct cb_arg *args)
{
	uint64_t handle = nftnl_trace_get_u64(nlt, NFTNL_TRACE_RULE_HANDLE);
	uint32_t family = nftnl_trace_get_u32(nlt, NFTNL_TRACE_FAMILY);
	const char *table = nftnl_trace_get_str(nlt, NFTNL_TRACE_TABLE);
	const char *chain = nftnl_trace_get_str(nlt, NFTNL_TRACE_CHAIN);
        struct nftnl_rule *r;
	struct mnl_socket *nl;
	struct nlmsghdr *nlh;
	uint32_t portid;
	char buf[16536];
	int ret;

        r = nftnl_rule_alloc();
	if (r == NULL) {
		perror("OOM");
		exit(EXIT_FAILURE);
	}

	nlh = nftnl_chain_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, family, NLM_F_DUMP, 0);

        nftnl_rule_set_u32(r, NFTNL_RULE_FAMILY, family);
	nftnl_rule_set_str(r, NFTNL_RULE_CHAIN, chain);
	nftnl_rule_set_str(r, NFTNL_RULE_TABLE, table);
	nftnl_rule_set_u64(r, NFTNL_RULE_POSITION, handle);
	nftnl_rule_nlmsg_build_payload(nlh, r);
	nftnl_rule_free(r);

	nl = mnl_socket_open(NETLINK_NETFILTER);
	if (nl == NULL) {
		perror("mnl_socket_open");
		exit(EXIT_FAILURE);
	}

	if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
		perror("mnl_socket_bind");
		exit(EXIT_FAILURE);
	}

	portid = mnl_socket_get_portid(nl);
        if (mnl_socket_sendto(nl, nlh, nlh->nlmsg_len) < 0) {
                perror("mnl_socket_send");
                exit(EXIT_FAILURE);
        }

	ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
        while (ret > 0) {
		args->is_event = false;
                ret = mnl_cb_run(buf, ret, 0, portid, rule_cb, args);
                if (ret <= 0)
                        break;
                ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
        }
        if (ret == -1) {
                perror("error");
                exit(EXIT_FAILURE);
        }
        mnl_socket_close(nl);
}

static void trace_print_packet(const struct nftnl_trace *nlt, struct cb_arg *args)
{
	struct list_head stmts = LIST_HEAD_INIT(stmts);
	uint32_t nfproto, family;
	uint16_t l4proto = 0;
	uint32_t mark;
	char name[IFNAMSIZ];

	printf("PACKET: %d %08x ", args->nfproto, nftnl_trace_get_u32(nlt, NFTNL_TRACE_ID));

	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_IIF))
		printf("IN=%s ", if_indextoname(nftnl_trace_get_u32(nlt, NFTNL_TRACE_IIF), name));
	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_OIF))
		printf("OUT=%s ", if_indextoname(nftnl_trace_get_u32(nlt, NFTNL_TRACE_OIF), name));

	family = nftnl_trace_get_u32(nlt, NFTNL_TRACE_FAMILY);
	nfproto = family;
	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_NFPROTO)) {
		nfproto = nftnl_trace_get_u32(nlt, NFTNL_TRACE_NFPROTO);

		if (family != nfproto)
			printf("NFPROTO=%d ", nfproto);
	}

	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_LL_HEADER)) {
		const struct ethhdr *eh;
		const char *linklayer;
		uint32_t i, len;
		uint16_t type = nftnl_trace_get_u16(nlt, NFTNL_TRACE_IIFTYPE);

		linklayer = nftnl_trace_get_data(nlt, NFTNL_TRACE_LL_HEADER, &len);
		switch (type) {
		case ARPHRD_ETHER:
			if (len < sizeof(*eh))
			       break;
			eh = (const void *)linklayer;
			printf("MACSRC=%s ", ether_ntoa((const void *)eh->h_source));
			printf("MACDST=%s ", ether_ntoa((const void *)eh->h_dest));
			printf("MACPROTO=%04x ", ntohs(eh->h_proto));
			break;
		default:
			printf("LL=0x%x ", type);
			for (i = 0 ; i < len; i++)
				printf("%02x", linklayer[i]);
			printf(" ");
			break;
		}
	}

	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_NETWORK_HEADER)) {
		const struct ip6_hdr *ip6h;
		const struct iphdr *iph;
		uint32_t i, len;
		const char *nh;

		ip6h = nftnl_trace_get_data(nlt, NFTNL_TRACE_NETWORK_HEADER, &len);

		switch (nfproto) {
		case NFPROTO_IPV4: {
			char addrbuf[INET_ADDRSTRLEN];

			if (len < sizeof(*iph))
				break;
			iph = (const void *)ip6h;


			inet_ntop(AF_INET, &iph->saddr, addrbuf, sizeof(addrbuf));
			printf("SRC=%s ", addrbuf);
			inet_ntop(AF_INET, &iph->daddr, addrbuf, sizeof(addrbuf));
			printf("DST=%s ", addrbuf);

			printf("LEN=%d TOS=0x%x TTL=%d ID=%d", ntohs(iph->tot_len), iph->tos, iph->ttl, ntohs(iph->id));
			if (iph->frag_off & htons(0x8000))
				printf("CE ");
			if (iph->frag_off & htons(IP_DF))
				printf("DF ");
			if (iph->frag_off & htons(IP_MF))
				printf("MF ");

			if (ntohs(iph->frag_off) & 0x1fff)
				printf("FRAG:%u ", ntohs(iph->frag_off) & 0x1fff);

			l4proto = iph->protocol;
			if (iph->ihl * 4 > sizeof(*iph)) {
				unsigned int optsize;
				const char *op;

				optsize = iph->ihl * 4 - sizeof(*iph);
				op = (const char *)iph;
				op += sizeof(*iph);

				printf("OPT (");
				for (i = 0; i < optsize; i++)
					printf("%02X", op[i]);
				printf(")");
			}
			break;
		}
		case NFPROTO_IPV6: {
			uint32_t flowlabel = ntohl(*(uint32_t *)ip6h);
			char addrbuf[INET6_ADDRSTRLEN];

			if (len < sizeof(*ip6h))
				break;

			inet_ntop(AF_INET6, &ip6h->ip6_src, addrbuf, sizeof(addrbuf));
			printf("SRC=%s ", addrbuf);
			inet_ntop(AF_INET6, &ip6h->ip6_dst, addrbuf, sizeof(addrbuf));
			printf("DST=%s ", addrbuf);

			printf("LEN=%zu TC=%u HOPLIMIT=%u FLOWLBL=%u ",
				ntohs(ip6h->ip6_plen) + sizeof(*iph),
				(flowlabel & 0x0ff00000) >> 20,
				ip6h->ip6_hops,
				flowlabel & 0x000fffff);

			l4proto = ip6h->ip6_nxt;
			break;
		}
		default:
			nh = (const char *)ip6h;
			printf("NH=");
			for (i = 0 ; i < len; i++)
				printf("%02x", nh[i]);
			printf(" ");
		}
	}

	if (nftnl_trace_is_set(nlt, NFTNL_TRACE_TRANSPORT_HEADER)) {
		const struct tcphdr *tcph;
		uint32_t len;

		tcph = nftnl_trace_get_data(nlt, NFTNL_TRACE_TRANSPORT_HEADER, &len);

		switch (l4proto) {
		case IPPROTO_DCCP:
		case IPPROTO_SCTP:
		case IPPROTO_UDPLITE:
		case IPPROTO_UDP:
			if (len < 4)
				break;
			printf("SPORT=%d DPORT=%d ", ntohs(tcph->source), ntohs(tcph->dest));
			break;
		case IPPROTO_TCP:
			if (len < sizeof(*tcph))
				break;
			printf("SPORT=%d DPORT=%d ", ntohs(tcph->source), ntohs(tcph->dest));
			if (tcph->syn)
				printf("SYN ");
			if (tcph->ack)
				printf("ACK ");
			if (tcph->fin)
				printf("FIN ");
			if (tcph->rst)
				printf("RST ");
			if (tcph->psh)
				printf("PSH ");
			if (tcph->urg)
				printf("URG ");
			break;
		default:
			break;
		}
	}

	mark = nftnl_trace_get_u32(nlt, NFTNL_TRACE_MARK);
	if (mark)
		printf("MARK=0x%x ", mark);
}

static void print_verdict(struct nftnl_trace *nlt, uint32_t verdict)
{
	const char *chain;

	switch (verdict) {
	case NF_ACCEPT:
		printf("ACCEPT");
		break;
	case NF_DROP:
		printf("DROP");
		break;
	case NF_QUEUE:
		printf("QUEUE");
		break;
	case NF_STOLEN:
		printf("STOLEN");
		break;
	case NFT_BREAK:
		printf("BREAK");
		break;
	case NFT_CONTINUE:
		printf("CONTINUE");
		break;
	case NFT_GOTO:
		printf("GOTO");
		if (nftnl_trace_is_set(nlt, NFTNL_TRACE_JUMP_TARGET)) {
			chain = nftnl_trace_get_str(nlt, NFTNL_TRACE_JUMP_TARGET);
			printf(":%s", chain);
		}
		break;
	case NFT_JUMP:
		printf("JUMP");
		if (nftnl_trace_is_set(nlt, NFTNL_TRACE_JUMP_TARGET)) {
			chain = nftnl_trace_get_str(nlt, NFTNL_TRACE_JUMP_TARGET);
			printf(":%s", chain);
		}
		break;
	default:
		printf("0x%x", verdict);
		break;
	}

	printf(" ");
}

static int trace_cb(const struct nlmsghdr *nlh, struct cb_arg *arg)
{
	struct nftnl_trace *nlt;
	uint32_t verdict;

	nlt = nftnl_trace_alloc();
	if (nlt == NULL)
		goto err;

	if (nftnl_trace_nlmsg_parse(nlh, nlt) < 0)
		goto err_free;

	if (arg->nfproto &&
	    arg->nfproto != nftnl_trace_get_u32(nlt, NFTNL_TABLE_FAMILY))
		goto err_free;

	printf(" TRACE: %d %08x %s:%s", nftnl_trace_get_u32(nlt, NFTNL_TABLE_FAMILY),
					nftnl_trace_get_u32(nlt, NFTNL_TRACE_ID),
					nftnl_trace_get_str(nlt, NFTNL_TRACE_TABLE),
					nftnl_trace_get_str(nlt, NFTNL_TRACE_CHAIN));

	switch (nftnl_trace_get_u32(nlt, NFTNL_TRACE_TYPE)) {
	case NFT_TRACETYPE_RULE:
		verdict = nftnl_trace_get_u32(nlt, NFTNL_TRACE_VERDICT);
		printf(":rule:0x%llx:", (unsigned long long)nftnl_trace_get_u64(nlt, NFTNL_TRACE_RULE_HANDLE));
		print_verdict(nlt, verdict);

		if (nftnl_trace_is_set(nlt, NFTNL_TRACE_RULE_HANDLE))
			trace_print_rule(nlt, arg);
		if (nftnl_trace_is_set(nlt, NFTNL_TRACE_LL_HEADER) ||
		    nftnl_trace_is_set(nlt, NFTNL_TRACE_NETWORK_HEADER))
			trace_print_packet(nlt, arg);
		break;
	case NFT_TRACETYPE_POLICY:
		printf(":policy:");
		verdict = nftnl_trace_get_u32(nlt, NFTNL_TRACE_POLICY);

		print_verdict(nlt, verdict);
		break;
	case NFT_TRACETYPE_RETURN:
		printf(":return:");
		trace_print_return(nlt);
		break;
	}
	puts("");
err_free:
	nftnl_trace_free(nlt);
err:
	return MNL_CB_OK;
}

static int monitor_cb(const struct nlmsghdr *nlh, void *data)
{
	uint32_t type = nlh->nlmsg_type & 0xFF;
	struct cb_arg *arg = data;
	int ret = MNL_CB_OK;

	switch(type) {
	case NFT_MSG_NEWTABLE:
	case NFT_MSG_DELTABLE:
		ret = table_cb(nlh, data);
		break;
	case NFT_MSG_NEWCHAIN:
	case NFT_MSG_DELCHAIN:
		ret = chain_cb(nlh, data);
		break;
	case NFT_MSG_NEWRULE:
	case NFT_MSG_DELRULE:
		arg->is_event = true;
		ret = rule_cb(nlh, data);
		break;
	case NFT_MSG_NEWGEN:
		ret = newgen_cb(nlh, data);
		break;
	case NFT_MSG_TRACE:
		ret = trace_cb(nlh, data);
		break;
	}

	return ret;
}

static const struct option options[] = {
	{.name = "counters", .has_arg = false, .val = 'c'},
	{.name = "trace", .has_arg = false, .val = 't'},
	{.name = "event", .has_arg = false, .val = 'e'},
	{.name = "ipv4", .has_arg = false, .val = '4'},
	{.name = "ipv6", .has_arg = false, .val = '6'},
	{.name = "version", .has_arg = false, .val = 'V'},
	{.name = "help", .has_arg = false, .val = 'h'},
	{NULL},
};

static void print_usage(void)
{
	printf("%s %s\n", xtables_globals.program_name,
			  xtables_globals.program_version);
	printf("Usage: %s [ -t | -e ]\n"
	       "        --trace    -t    trace ruleset traversal of packets tagged via -j TRACE rule\n"
	       "        --event    -e    show events that modify the ruleset\n"
	       "Optional arguments:\n"
	       "        --ipv4     -4    only monitor IPv4\n"
	       "        --ipv6     -6    only monitor IPv6\n"
	       "	--counters -c    show counters in rules\n"

	       , xtables_globals.program_name);
	exit(EXIT_FAILURE);
}

int xtables_monitor_main(int argc, char *argv[])
{
	struct mnl_socket *nl;
	char buf[MNL_SOCKET_BUFFER_SIZE];
	uint32_t nfgroup = 0;
	struct nft_handle h = {};
	struct cb_arg cb_arg = {
		.h = &h,
	};
	int ret, c;

	xtables_globals.program_name = "xtables-monitor";
	/* XXX xtables_init_all does several things we don't want */
	c = xtables_init_all(&xtables_globals, NFPROTO_IPV4);
	if (c < 0) {
		fprintf(stderr, "%s/%s Failed to initialize xtables\n",
				xtables_globals.program_name,
				xtables_globals.program_version);
		exit(1);
	}
#if defined(ALL_INCLUSIVE) || defined(NO_SHARED_LIBS)
	init_extensions();
	init_extensions4();
#endif

	if (nft_init(&h, xtables_ipv4)) {
		fprintf(stderr, "%s/%s Failed to initialize nft: %s\n",
			xtables_globals.program_name,
			xtables_globals.program_version,
			strerror(errno));
		exit(EXIT_FAILURE);
	}

	opterr = 0;
	while ((c = getopt_long(argc, argv, "ceht46V", options, NULL)) != -1) {
		switch (c) {
	        case 'c':
			counters = true;
			break;
	        case 't':
			trace = true;
			break;
	        case 'e':
			events = true;
			break;
	        case 'h':
			print_usage();
			exit(0);
		case '4':
			cb_arg.nfproto = NFPROTO_IPV4;
			break;
		case '6':
			cb_arg.nfproto = NFPROTO_IPV6;
			break;
		case 'V':
			printf("xtables-monitor %s\n", PACKAGE_VERSION);
			exit(0);
		default:
			fprintf(stderr, "xtables-monitor %s: Bad argument.\n", PACKAGE_VERSION);
			fprintf(stderr, "Try `xtables-monitor -h' for more information.\n");
			exit(PARAMETER_PROBLEM);
		}
	}

	if (trace)
		nfgroup |= 1 << (NFNLGRP_NFTRACE - 1);
	if (events)
		nfgroup |= 1 << (NFNLGRP_NFTABLES - 1);

	if (nfgroup == 0) {
		print_usage();
		exit(EXIT_FAILURE);
	}

	nl = mnl_socket_open(NETLINK_NETFILTER);
	if (nl == NULL) {
		perror("cannot open nfnetlink socket");
		exit(EXIT_FAILURE);
	}

	if (mnl_socket_bind(nl, nfgroup, MNL_SOCKET_AUTOPID) < 0) {
		perror("cannot bind to nfnetlink socket");
		exit(EXIT_FAILURE);
	}

	ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
	while (ret > 0) {
		ret = mnl_cb_run(buf, ret, 0, 0, monitor_cb, &cb_arg);
		if (ret <= 0)
			break;
		ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
	}
	if (ret == -1) {
		perror("cannot receive from nfnetlink socket");
		exit(EXIT_FAILURE);
	}
	mnl_socket_close(nl);

	return EXIT_SUCCESS;
}