/*
 * stats.c - implements wanstats for hosts.
 *
 * Copyright (C) 2023 IOPSYS Software Solutions AB. All rights reserved.
 *
 * See LICENSE file for license related information.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <errno.h>

#include <libnetfilter_conntrack/libnetfilter_conntrack.h>

#include <json-c/json.h>
#include <libubox/blobmsg.h>
#include <libubox/blobmsg_json.h>
#include <libubox/uloop.h>
#include <libubox/ustream.h>
#include <libubox/utils.h>
#include <libubus.h>

#include <easy/easy.h>

#include "debug.h"
#include "util.h"
#include "timer.h"
#include "neigh.h"
#include "hostmngr.h"
#include "lookup3_hash.h"


void neigh_flowtable_flush(void *priv)
{
	struct hostmngr_private *p = (struct hostmngr_private *)priv;
	struct neigh_flow_table *ft = (struct neigh_flow_table *)&p->nftable;
	struct neigh_flow_entry *e = NULL;
	struct hlist_node *tmp;
	int i;

	for (i = 0; i < NEIGH_FLOW_ENTRIES_MAX; i++) {
		hlist_for_each_entry_safe(e, tmp, &ft->table[i], hlist) {
			hlist_del(&e->hlist, &ft->table[i]);
			free(e);
		}
	}
}

void neigh_iptable_flush(void *priv)
{
	struct hostmngr_private *p = (struct hostmngr_private *)priv;
	struct neigh_flow_table *ft = (struct neigh_flow_table *)&p->nftable;
	struct neigh_ip_entry *e = NULL;
	struct hlist_node *tmp;
	int i;

	for (i = 0; i < NEIGH_IP_ENTRIES_MAX; i++) {
		hlist_for_each_entry_safe(e, tmp, &ft->iptable[i], iphlist) {
			hlist_del(&e->iphlist, &ft->iptable[i]);
			free(e);
		}
	}
}


uint32_t ip4flow_hash(struct ip_address *sip, uint16_t sport,
		      struct ip_address *dip, uint16_t dport,
		      uint32_t protocol)
{
	uint32_t initval = 0xdeadbeef;
	uint32_t key;

	if (sip->addr.ip4.s_addr < dip->addr.ip4.s_addr) {
		uint32_t pp = (sport | (dport << 16));
		uint32_t arr[4] = {sip->addr.ip4.s_addr, dip->addr.ip4.s_addr, pp, protocol};

		key = hashword(arr, 4, initval);
	} else {
		uint32_t pp = (dport | (sport << 16));
		uint32_t arr[4] = {dip->addr.ip4.s_addr, sip->addr.ip4.s_addr, pp, protocol};

		key = hashword(arr, 4, initval);
	}

	return key & (NEIGH_FLOW_ENTRIES_MAX - 1);
}

uint32_t ip6flow_hash(struct ip_address *sip, uint16_t sport,
		      struct ip_address *dip, uint16_t dport,
		      uint32_t protocol)
{
	uint32_t initval = 0xdeadbeef;
	uint32_t pp = (sport | (dport << 16));
	uint32_t key;
	uint32_t arr[10] = {*(uint32_t *)sip->addr.ip6.s6_addr,
				*(uint32_t *)dip->addr.ip6.s6_addr, pp, protocol};

	key = hashword(arr, 10, initval);

	return key & (NEIGH_FLOW_ENTRIES_MAX - 1);
}


struct neigh_flow_entry *neigh_flow_create(void *nfq,
					   struct ip_address *sip,
					   struct ip_address *dip,
					   uint16_t sport,
					   uint16_t dport,
					   uint32_t l4proto)
{
	struct neigh_flow_table *ft = (struct neigh_flow_table *)nfq;
	struct neigh_flow_entry *e;
	struct neigh_ip_entry *n = NULL;
	uint32_t nkey;
	uint32_t fkey;


	e = calloc(1, sizeof(*e));
	if (!e) {
		err("%s: -ENOMEM\n", __func__);
		return NULL;
	}

	memcpy(&e->sip, sip, sizeof(struct ip_address));
	memcpy(&e->dip, dip, sizeof(struct ip_address));
	e->sport = sport;
	e->dport = dport;
	e->l4proto = l4proto;
	e->ipneigh = NULL;

	nkey = ipaddr_hash(sip);
	hlist_for_each_entry(n, &ft->iptable[nkey], iphlist) {
		if (ipaddr_equal(&n->ip, sip)) {
			//char sipbuf[32] = {0};

			e->ipneigh = n;
			memcpy(e->macaddr, n->macaddr, 6);
			/*
			inet_ntop(e->sip.family, &e->sip.addr, sipbuf, sizeof(sipbuf));
			dbg("Link new-flow entry %s to ----> neigh " MACFMT"\n",
			    sipbuf, MAC2STR(n->macaddr));
			*/

			if (l4proto == IPPROTO_TCP)
				e->ipneigh->num_tcp++;
			else if (l4proto == IPPROTO_UDP)
				e->ipneigh->num_udp++;
			break;
		}
	}

	if (e->sip.family == AF_INET)
		fkey = ip4flow_hash(sip, sport, dip, dport, l4proto);
	else
		fkey = ip6flow_hash(sip, sport, dip, dport, l4proto);

	hlist_add_head(&e->hlist, &ft->table[fkey]);
	ft->num++;

	return e;
}

struct neigh_flow_entry *neigh_flow_lookup(void *nfq,
					   struct ip_address *sip,
					   struct ip_address *dip,
					   uint16_t sport,
					   uint16_t dport,
					   uint32_t l4proto)
{
	struct neigh_flow_table *ft = (struct neigh_flow_table *)nfq;
	struct neigh_flow_entry *e = NULL;
	uint32_t key;


	if (sip->family == AF_INET && dip->family == AF_INET)
		key = ip4flow_hash(sip, sport, dip, dport, l4proto);
	else
		key = ip6flow_hash(sip, sport, dip, dport, l4proto);

	hlist_for_each_entry(e, &ft->table[key], hlist) {
		if (ipaddr_equal(&e->sip, sip) && ipaddr_equal(&e->dip, dip)
		    && e->sport == sport && e->dport == dport
		    && e->l4proto == l4proto) {

			return e;
		}
	}

	return NULL;
}

void neigh_flow_delete(void *nfq, struct ip_address *sip, struct ip_address *dip,
		       uint16_t sport, uint16_t dport, uint32_t l4proto,
		       const struct neigh_ip_entry *ipneigh)
{
	struct neigh_flow_table *ft = (struct neigh_flow_table *)nfq;
	struct neigh_flow_entry *e = NULL;
	struct hlist_node *tmp;
	uint32_t key;


	if (sip->family == AF_INET && dip->family == AF_INET)
		key = ip4flow_hash(sip, sport, dip, dport, l4proto);
	else
		key = ip6flow_hash(sip, sport, dip, dport, l4proto);

	hlist_for_each_entry_safe(e, tmp, &ft->table[key], hlist) {
		if (ipaddr_equal(&e->sip, sip) && ipaddr_equal(&e->dip, dip)
		    && e->sport == sport && e->dport == dport
		    && e->l4proto == l4proto) {

			hlist_del(&e->hlist, &ft->table[key]);
			ft->num--;
			if (ipneigh && e->ipneigh == ipneigh) {
				if (l4proto == IPPROTO_TCP && e->ipneigh->num_tcp > 0)
					e->ipneigh->num_tcp--;
				else if (l4proto == IPPROTO_UDP && e->ipneigh->num_udp > 0)
					e->ipneigh->num_udp--;

				e->ipneigh = NULL;
			}
			free(e);
			break;
		}
	}
}

struct neigh_flow_entry *neigh_flow_lookup2(void *nfq,
					   struct ip_address *sip,
					   struct ip_address *dip,
					   uint16_t sport,
					   uint16_t dport,
					   uint32_t l4proto)
{
	struct neigh_flow_table *ft = (struct neigh_flow_table *)nfq;
	struct neigh_flow_entry *e = NULL;
	struct neigh_ip_entry *n;
	uint32_t nkey;
	uint32_t key;


	if (sip->family == AF_INET && dip->family == AF_INET)
		key = ip4flow_hash(sip, sport, dip, dport, l4proto);
	else
		key = ip6flow_hash(sip, sport, dip, dport, l4proto);

	hlist_for_each_entry(e, &ft->table[key], hlist) {
		if (ipaddr_equal(&e->sip, sip) && ipaddr_equal(&e->dip, dip)
		    && e->sport == sport && e->dport == dport
		    && e->l4proto == l4proto) {

			return e;
		}
	}

	/* entry not found; create one */
	e = calloc(1, sizeof(*e));
	if (!e) {
		err("%s: -ENOMEM\n", __func__);
		return NULL;
	}

	memcpy(&e->sip, sip, sizeof(struct ip_address));
	memcpy(&e->dip, dip, sizeof(struct ip_address));
	e->sport = sport;
	e->dport = dport;
	e->l4proto = l4proto;
	e->ipneigh = NULL;
	nkey = ipaddr_hash(sip);
	hlist_for_each_entry(n, &ft->iptable[nkey], iphlist) {
		if (ipaddr_equal(&n->ip, sip)) {
			//char sipbuf[32] = {0};

			e->ipneigh = n;
			memcpy(e->macaddr, n->macaddr, 6);
			/*
			inet_ntop(e->sip.family, &e->sip.addr, sipbuf, sizeof(sipbuf));
			dbg("Link new-flow entry %s to ----> neigh " MACFMT"\n",
			    sipbuf, MAC2STR(n->macaddr));
			*/

			if (l4proto == IPPROTO_TCP)
				e->ipneigh->num_tcp++;
			else if (l4proto == IPPROTO_UDP)
				e->ipneigh->num_udp++;
			break;
		}
	}

	if (e->sip.family == AF_INET)
		key = ip4flow_hash(sip, sport, dip, dport, l4proto);
	else
		key = ip6flow_hash(sip, sport, dip, dport, l4proto);

	hlist_add_head(&e->hlist, &ft->table[key]);
	ft->num++;

	return e;
}

struct neigh_ip_entry *neigh_ip_entry_lookup(void *priv, struct ip_address *ip)
{
	struct hostmngr_private *p = (struct hostmngr_private *)priv;
	struct neigh_ip_entry *e = NULL;
	uint32_t key;


	if (!priv || !ip)
		return NULL;

	key = ipaddr_hash(ip);

	/* dbg("%s: ip = 0x%x, idx = %u\n", __func__, ip->addr.ip4.s_addr, key); */
	hlist_for_each_entry(e, &p->nftable.iptable[key], iphlist) {
		if (ipaddr_equal(&e->ip, ip)) {
			return e;
		}
	}

	return NULL;
}

struct neigh_ip_entry *neigh_ip_entry_lookup2(void *priv, struct ip_address *ip)
{
	struct neigh_ip_entry *e = NULL;
	struct neigh_flow_table *ft;
	uint32_t key;


	if (!priv || !ip)
		return NULL;

	ft = &((struct hostmngr_private *)priv)->nftable;
	key = ipaddr_hash(ip);
	hlist_for_each_entry(e, &ft->iptable[key], iphlist) {
		if (ipaddr_equal(&e->ip, ip)) {
			return e;
		}
	}

	/* create new neigh_ip_entry */
	e = calloc(1, sizeof(*e));
	if (!e) {
		err("%s: -ENOMEM\n", __func__);
		return NULL;
	}

	memcpy(&e->ip, ip, sizeof(struct ip_address));
	key = ipaddr_hash(ip);
	hlist_add_head(&e->iphlist, &ft->iptable[key]);

	return e;
}

struct ct_nlevent {
	struct uloop_fd uloop;
	void (*error_cb)(struct ct_nlevent *e, int error);
	void (*event_cb)(struct ct_nlevent *e);
};

struct event_socket {
	struct ct_nlevent ev;
	struct nfct_handle *h;
	int sock_bufsize;
};

static void handle_error(struct ct_nlevent *e, int error)
{
	//struct event_socket *ev_sock = container_of(e, struct event_socket, ev);

	if (error != ENOBUFS)
		goto err;

	/* TODO
	ev_sock->sock_bufsize *= 2;
	if (nl_socket_set_buffer_size(nfct_fd(ev_sock->h), ev_sock->sock_bufsize, 0))
		goto err;
	*/

	return;

err:
	e->uloop.cb = NULL;
	uloop_fd_delete(&e->uloop);
}

static void recv_nlevents(struct ct_nlevent *e)
{
	struct event_socket *evs = container_of(e, struct event_socket, ev);
	int ret;

	ret = nfct_catch(evs->h);
	if (ret == -1 && errno != EAGAIN) {
		dbg("%s: ret = %d (%s)\n", __func__, ret, strerror(errno));
	}
}

static struct event_socket ct_event = {
	.ev = {
		.uloop = {.fd = - 1, },
		.error_cb = handle_error,
		.event_cb = recv_nlevents,
	},
	.h = NULL,
	.sock_bufsize = 0x20000,
};

static int nfct_cb(enum nf_conntrack_msg_type type, struct nf_conntrack *ct,
		   void *data)
{
	struct hostmngr_private *priv = (struct hostmngr_private *)data;
	struct nfct_attr_grp_ctrs cnto = {0};
	struct nfct_attr_grp_ctrs cntr = {0};
	struct nfct_attr_grp_ipv4 ip4 = {0};
	struct nfct_attr_grp_ipv6 ip6 = {0};
	struct nfct_attr_grp_port port = {0};
	struct neigh_flow_entry *e = NULL;
	struct neigh_ip_entry *ipe;
	struct ip_address sip = {0};
	struct ip_address dip = {0};
	uint8_t l4proto = 0;
	char srcip[256] = {0};
	char dstip[256] = {0};

	if (nfct_attr_grp_is_set(ct, ATTR_GRP_ORIG_IPV4)) {
		nfct_get_attr_grp(ct, ATTR_GRP_ORIG_IPV4, &ip4);
		inet_ntop(AF_INET, &ip4.src, srcip, sizeof(srcip));
		inet_ntop(AF_INET, &ip4.dst, dstip, sizeof(dstip));

		sip.family = AF_INET;
		memcpy(&sip.addr, &ip4.src, sizeof(ip4.src));
		dip.family = AF_INET;
		memcpy(&dip.addr, &ip4.dst, sizeof(ip4.dst));

	} else if (nfct_attr_grp_is_set(ct, ATTR_GRP_ORIG_IPV6)) {
		nfct_get_attr_grp(ct, ATTR_GRP_ORIG_IPV6, &ip6);
		inet_ntop(AF_INET6, &ip6.src, srcip, sizeof(srcip));
		inet_ntop(AF_INET6, &ip6.dst, dstip, sizeof(dstip));

		sip.family = AF_INET6;
		memcpy(&sip.addr, &ip6.src, sizeof(ip6.src));
		dip.family = AF_INET6;
		memcpy(&dip.addr, &ip6.dst, sizeof(ip6.dst));
	}

	if (nfct_attr_grp_is_set(ct, ATTR_GRP_ORIG_PORT))
		nfct_get_attr_grp(ct, ATTR_GRP_ORIG_PORT, &port);

	if (nfct_attr_grp_is_set(ct, ATTR_ORIG_L4PROTO))
		l4proto = nfct_get_attr_u8(ct, ATTR_ORIG_L4PROTO);

	if (nfct_attr_grp_is_set(ct, ATTR_GRP_ORIG_COUNTERS)) {
		nfct_get_attr_grp(ct, ATTR_GRP_ORIG_COUNTERS, &cnto);
		dbg("Orig Counters: packets = %ju, bytes = %ju\n",
		    (uintmax_t)cnto.packets, (uintmax_t)cnto.bytes);
	}

	if (nfct_attr_grp_is_set(ct, ATTR_GRP_REPL_COUNTERS)) {
		nfct_get_attr_grp(ct, ATTR_GRP_REPL_COUNTERS, &cntr);
		dbg("Repl Counters: packets = %ju, bytes = %ju\n",
		    (uintmax_t)cntr.packets, (uintmax_t)cntr.bytes);
	}

	if (type != NFCT_T_NEW && type != NFCT_T_UPDATE && type != NFCT_T_DESTROY)
		return NFCT_CB_CONTINUE;

	ipe = neigh_ip_entry_lookup(priv, &sip);
	if (!ipe) {
		if (type == NFCT_T_DESTROY) {
			/* dbg("DESTROY: l4proto = %d  IP-src: %s (sport = %hu), IP-dst: %s (dport = %hu)\n",
			    l4proto, srcip, port.sport, dstip, port.dport); */

			neigh_flow_delete(&priv->nftable,
					  &sip, &dip,
					  port.sport, port.dport,
					  l4proto,
					  NULL);
		}

		return NFCT_CB_CONTINUE;
	}

	/* dbg("Type = 0x%x: l4proto = %d  IP-src: %s (sport = %hu), IP-dst: %s (dport = %hu)\n",
	    type, l4proto, srcip, port.sport, dstip, port.dport); */

	e = neigh_flow_lookup2(&priv->nftable, &sip, &dip, port.sport, port.dport, l4proto);
	if (!e)
		return NFCT_CB_CONTINUE;

	if (cnto.bytes > 0 || cntr.bytes > 0) {
		if (e->ipneigh) {
			if (e->ipneigh->neigh) {
				//dbg("[1] Neigh present now\n");
				if (cnto.bytes > 0) {
					e->ipneigh->ws.ul_packets += cnto.packets;
					e->ipneigh->ws.ul_bytes += cnto.bytes;
				}
				if (cntr.bytes > 0) {
					e->ipneigh->ws.dl_packets += cntr.packets;
					e->ipneigh->ws.dl_bytes += cntr.bytes;
				}
			}

			if (type == NFCT_T_DESTROY) {
				/* dbg("Destroy: l4proto = %d  IP-src: %s (sport = %hu), IP-dst: %s (dport = %hu)\n",
				    l4proto, srcip, port.sport, dstip, port.dport); */

				neigh_flow_delete(&priv->nftable,
						  &sip, &dip,
						  port.sport, port.dport,
						  l4proto,
						  e->ipneigh);

				return NFCT_CB_CONTINUE;
			}
		} else {
			struct neigh_ip_entry *n = NULL;
			uint32_t nkey;

			nkey = ipaddr_hash(&sip);
			hlist_for_each_entry(n, &priv->nftable.iptable[nkey], iphlist) {
				if (ipaddr_equal(&n->ip, &sip)) {
					//char sipbuf[32] = {0};

					e->ipneigh = n;
					memcpy(e->macaddr, n->macaddr, 6);

					if (e->ipneigh->neigh) {
						dbg("[2] Neigh present now\n");
						/*
						inet_ntop(e->sip.family, &e->sip.addr, sipbuf, sizeof(sipbuf));
						dbg("Link flow entry %s to ----> neigh " MACFMT"\n",
						    sipbuf, MAC2STR(n->macaddr));
						*/

						if (cnto.bytes > 0) {
							e->ipneigh->ws.ul_packets += cnto.packets;
							e->ipneigh->ws.ul_bytes += cnto.bytes;
						}
						if (cntr.bytes > 0) {
							e->ipneigh->ws.dl_packets += cntr.packets;
							e->ipneigh->ws.dl_bytes += cntr.bytes;
						}
					}
					break;
				}
			}
		}
	}

	return NFCT_CB_CONTINUE;
}

static void ct_receive_nlevents(struct uloop_fd *u, unsigned int events)
{
	struct ct_nlevent *e = container_of(u, struct ct_nlevent, uloop);

	if (u->error) {
		int ret = -1;
		socklen_t ret_len = sizeof(ret);

		u->error = false;
		if (e->error_cb &&
		    getsockopt(u->fd, SOL_SOCKET, SO_ERROR, &ret, &ret_len) == 0) {
			e->error_cb(e, ret);
		}
	}

	if (e->event_cb) {
		e->event_cb(e);
		return;
	}
}

int nfct_get_entries_nolo(struct hostmngr_private *priv)
{
	struct nfct_filter_ipv4 filter_ipv4 = {
		.addr = ntohl(inet_addr("127.0.0.1")),
		.mask = 0xffffffff,
	};

	struct nfct_filter_ipv6 filter_ipv6 = {
		.addr = { 0x0, 0x0, 0x0, 0x1 },
		.mask = { 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff },
	};
	struct nfct_filter *filter;
	struct nfct_handle *h;


	h = nfct_open(CONNTRACK, NFCT_ALL_CT_GROUPS);
	if (!h) {
		perror("nfct_open");
		return 0;
	}

	filter = nfct_filter_create();
	if (!filter) {
		perror("nfct_create_filter");
		nfct_close(h);
		return 0;
	}

	nfct_filter_set_logic(filter, NFCT_FILTER_SRC_IPV4, NFCT_FILTER_LOGIC_NEGATIVE);
	nfct_filter_add_attr(filter, NFCT_FILTER_SRC_IPV4, &filter_ipv4);

	nfct_filter_set_logic(filter, NFCT_FILTER_SRC_IPV6, NFCT_FILTER_LOGIC_NEGATIVE);
	nfct_filter_add_attr(filter, NFCT_FILTER_SRC_IPV6, &filter_ipv6);

	if (nfct_filter_attach(nfct_fd(h), filter) == -1) {
		perror("nfct_filter_attach");
		nfct_filter_destroy(filter);
		nfct_close(h);
		return 0;
	}

	nfct_filter_destroy(filter);

	nfct_callback_register(h, NFCT_T_ALL, nfct_cb, priv);

	ct_event.h = h;
	ct_event.ev.uloop.fd = nfct_fd(ct_event.h);
	ct_event.ev.uloop.cb = ct_receive_nlevents;
	uloop_fd_add(&ct_event.ev.uloop, ULOOP_READ |
		     ((ct_event.ev.error_cb) ? ULOOP_ERROR_CB : 0));

	return 0;
}

void nfct_cleanup(void)
{
	uloop_fd_delete(&ct_event.ev.uloop);
	nfct_close(ct_event.h);
}
