#include <sys/types.h>
#include <sys/socket.h>
#include <sys/param.h>
#include <sys/time.h>

#if !defined(__OpenBSD__)
#define __FAVOR_BSD
#endif
#include <netinet/in.h>
#include <arpa/inet.h>
#include <net/if.h>
#include <netinet/in_systm.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <netinet/ip_icmp.h>

#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <string.h>
#include <errno.h>
#include <pcap.h>

#include "glougloud.h"
#include "external/imsgev.h"

/* XXX a user process should be able to have multiple clients, so that we
 * can send the same data to multiple machines with almost no overhead */

#define NULL_HDRLEN 4 // XXX portable ?
#define NODE_MAX_WITHOUT_TIMEOUT 1000
#define NODE_TIMEOUT 300 // XXX conf ?
#define CONN_TIMEOUT 300 // XXX conf ?
#define CONN_TIMEOUT_UDP 20 // XXX conf ?
#define CONN_TIMEOUT_ICMP 10 // XXX conf ?
#define CONN_FREEIDS_COUNT 65536 /* 2^16 as long as freeids are u_int16_t */
#define CONNTIMER 5 // XXX conf ?

struct node {
	LIST_ENTRY(node) entry;
	struct in_addr	addr;
	time_t		lastseen;
	int		used;
	short		namelen;
#define NODENAME_WAITING -1
#define NODENAME_FAILED -2
	char		*name;
};

enum connstate {
	CONNSTATE_ESTABLISHED,
	CONNSTATE_TCPFIN,
	CONNSTATE_TCPFIN2
};

struct conn {
	LIST_ENTRY(conn) entry;
	u_int		id;
	enum connstate	state;
	struct node 	*src;
	u_int		src_port;
	struct node 	*dst;
	u_int		dst_port;
	u_int		proto;
	u_int		size;
	time_t		lastseen;
};

struct user {
	LIST_ENTRY(user)	entry;
	struct sockaddr_in	addr;
};

struct capture {
	int		fd;
	struct imsgev	iev;

	pcap_t		*pcap;
	struct event	pcap_ev;
	struct timeval	pcap_tv;
	pcap_handler	pcap_handler;

	LIST_HEAD(, conn) conn_list;
	LIST_HEAD(, node) node_list;
	int		node_count;
	u_int16_t	conn_freeids[CONN_FREEIDS_COUNT];
	int		conn_freeids_ptr;
	struct event	conntimer_ev;
	struct timeval	conntimer_tv;
	time_t		time;

	int		pinvalid;
	int		ptruncated;
};

struct phandler {
	pcap_handler f;
	int type;
};

static void		 ip_handle(struct ip *, const u_char *, u_int);
static void		 sendto_all(struct packet *, int);
static struct node	*node_add(struct in_addr *);
static void		 node_del(struct node *);
static struct node	*node_find(struct in_addr *);
static void		 conn_add(struct in_addr *, int, struct in_addr *, int, int, int);
static void		 conn_data(struct conn *, int, int);
static void		 conn_del(struct conn *);
static pcap_handler	 lookup_phandler(int);
static struct user	*finduser(struct sockaddr_in *);

static void		phandler_ether(u_char *,
					const struct pcap_pkthdr *, const u_char *);
static void		phandler_loop(u_char *,
					const struct pcap_pkthdr *, const u_char *);
static void		ev_pcap(int, short, void *);
static void		ev_timer(int, short, void *);

static void		usrconn(struct imsgev *, struct imsg *);
static void		imsgev_main(struct imsgev *, int, struct imsg *);
static void   imsgev_main_needfd(struct imsgev *);

static struct phandler phandlers[] = {
        { phandler_ether,               DLT_EN10MB },
        { phandler_ether,               DLT_IEEE802 },
        { phandler_loop,                DLT_LOOP },
        { NULL,                         0 },
};
struct capture		*cap;
LIST_HEAD(, user)	usr_list;
int			usr_count = 0;

static void
sig_handler(int sig, short why, void *data)
{
	log_info("user: got signal %d", sig);
	if (sig == SIGINT || sig == SIGTERM)
		event_loopexit(NULL);
}

int
user_init(int fd[2], pcap_t *pcap)
{
	struct event ev_sigint, ev_sigterm, ev_sigchld, ev_sighup;
	int pid, i;

	pid = fork();
	if (pid < 0)
		fatal("user fork");
	if (pid > 0)
		return pid;

	setproctitle("user");
	event_init();
	close(fd[0]);

	signal_set(&ev_sigint, SIGINT, sig_handler, NULL);
	signal_set(&ev_sigterm, SIGTERM, sig_handler, NULL);
	signal_set(&ev_sigchld, SIGCHLD, sig_handler, NULL);
	signal_set(&ev_sighup, SIGHUP, sig_handler, NULL);
	signal_add(&ev_sigint, NULL);
	signal_add(&ev_sigterm, NULL);
	signal_add(&ev_sigchld, NULL);
	signal_add(&ev_sighup, NULL);
	signal(SIGPIPE, SIG_IGN);

	cap = xcalloc(1, sizeof(struct capture));
	cap->fd = fd[1];
	imsgev_init(&cap->iev, cap->fd, NULL, imsgev_main, imsgev_main_needfd);
	for (i=0; i<CONN_FREEIDS_COUNT-1; i++)
		cap->conn_freeids[i] = i;
	cap->time = time(NULL);

	cap->pcap = pcap;
	cap->pcap_handler = lookup_phandler(pcap_datalink(pcap));
	cap->pcap_tv.tv_sec = 0;
	cap->pcap_tv.tv_usec = PCAP_TO;
	event_set(&cap->pcap_ev, pcap_fileno(cap->pcap),
		EV_READ, ev_pcap, NULL);

	cap->conntimer_tv.tv_sec = CONNTIMER;
	cap->conntimer_tv.tv_usec = 0;
	evtimer_set(&cap->conntimer_ev, ev_timer, NULL);
	if (event_add(&cap->conntimer_ev, &cap->conntimer_tv) == -1)
		fatal("user: event_add conntimer failed: %s", strerror(errno));

	droppriv();

	log_info("user: entering event loop");
	event_dispatch();

	log_info("user: exiting");
	exit(0);
}

/*
 * Parse an IP packet and descide what to do with it.
 * 'ip' is a pointer the the captured IP packet
 * 'pend' is a pointer to the end of the captured IP packet
 * 'wirelen' is the size of the IP packet off the wire
 */
#define NOTCAPTURED(v) ((u_char *)v > (u_char *)pend - sizeof(*v))
#define NOTRECEIVED(v) (wirelen < sizeof(v))
static void
ip_handle(struct ip *ip, const u_char *pend, u_int wirelen)
{
	u_int		len, ip_hlen, off;
	const u_char	*cp;
	struct tcphdr	*tcph;
	struct udphdr	*udph;
	struct icmp	*icmp;
	struct in_addr	src, dst;
	u_int		src_port, dst_port;
	u_int		size, proto, close, response;
	struct conn	*c, *conn;

	if (NOTCAPTURED(ip)) {
		log_pinvalid("user: ip truncated (ip %x pend %x sizeof(ip) %d",
        ip, pend, sizeof(ip));
		cap->ptruncated++;
		return;
	}

	if (ip->ip_v != IPVERSION) {
		log_pinvalid("user: invalid ip version");
		cap->pinvalid++;
		return;
	}

	len = ntohs(ip->ip_len);
	if (wirelen < len) {
		log_pinvalid("user: ip too small");
		cap->pinvalid++;
		len = wirelen;
	}

	ip_hlen = ip->ip_hl * 4;
	if (ip_hlen < sizeof(struct ip) || ip_hlen > len) {
		log_pinvalid("user: ip_hlen invalid, %d", ip_hlen);
		cap->pinvalid++;
		return;
	}
	len -= ip_hlen;

	src.s_addr = ntohl(ip->ip_src.s_addr);
	dst.s_addr = ntohl(ip->ip_dst.s_addr);
	src_port = 0;
	dst_port = 0;
	proto = IPPROTO_IP;
	size = len;
	close = 0;

	off = ntohs(ip->ip_off);
	if ((off & IP_OFFMASK) == 0) {
		cp = (const u_char *)ip + ip_hlen;
		switch (ip->ip_p) {

		case IPPROTO_TCP:
			tcph = (struct tcphdr *)cp;
			if (NOTCAPTURED(&tcph->th_flags)) {
				log_pinvalid("user: tcp truncated");
				cap->ptruncated++;
				return;
			}
			if (NOTRECEIVED(*tcph)) {
				log_pinvalid("user: tcp too small");
				cap->pinvalid++;
				return;
			}
			src_port = ntohs(tcph->th_sport);
			dst_port = ntohs(tcph->th_dport);
			proto = IPPROTO_TCP;
			size = len - sizeof(*tcph);
			if ((tcph->th_flags & TH_FIN) &&
				(tcph->th_flags & TH_ACK))
				close = 1;
			break;

		case IPPROTO_UDP:
			udph = (struct udphdr *)cp;
			if (NOTCAPTURED(&udph->uh_dport)) {
				log_pinvalid("user: udp truncated, "
					"ip %x, udph %x, uh_port %x, pend %x, ip_hlen %d",
					ip, udph, &udph->uh_dport, pend, ip_hlen);
				cap->ptruncated++;
				return;
			}
			if (NOTRECEIVED(*udph)) {
				log_pinvalid("user: udp too small");
				cap->pinvalid++;
				return;
			}
			src_port = ntohs(udph->uh_sport);
			dst_port = ntohs(udph->uh_dport);
			proto = IPPROTO_UDP;
			size = len - sizeof(*udph);
			break;

		case IPPROTO_ICMP:
			icmp = (struct icmp *)cp;
			if (NOTRECEIVED(*icmp)) {
				log_pinvalid("user: icmp too small");
				cap->pinvalid++;
				return;
			}
			proto = IPPROTO_ICMP;
			size = len - sizeof(*icmp);
			break;

		default:
			log_warn("user: unknown ip protocol !");
			break;
		}
	} else {
		/*
		 * if this isn't the first frag, we're missing the
		 * next level protocol header.
		 */
		log_tmp("user: got a fragmented ip packet !");
	}

	conn = NULL;
	LIST_FOREACH(c, &cap->conn_list, entry) {
		if (((c->src->addr.s_addr == src.s_addr &&
				c->src_port == src_port &&
				c->dst->addr.s_addr == dst.s_addr &&
				c->dst_port == dst_port) ||
			(c->src->addr.s_addr == dst.s_addr &&
				 c->src_port == dst_port &&
				 c->dst->addr.s_addr == src.s_addr &&
				 c->dst_port == src_port)) &&
			c->proto == proto) {
			conn = c;
			if (c->src->addr.s_addr == src.s_addr)
				response = 0;
			else
				response = 1;
			break;
		}
	}

	if (conn) {
		if (!close) {
			conn_data(conn, size, response);
		} else {
			conn_del(conn);
		}
	} else {
		if (!close) {
			conn_add(&src, src_port, &dst, dst_port, proto, size);
		} else {
			log_warn("user: captured connection close w/o open !");
		}
	}
}

/* XXX all the packets must have data htos */
static void
sendto_all(struct packet *p, int size)
{
	struct user *usr;

	log_debug("SEND PACKET (size %.2d): %x %x", size, p->ver, p->type);

	LIST_FOREACH(usr, &usr_list, entry) {
		if (sendto(net_socket, p, size, 0,
			(struct sockaddr *)&usr->addr, sizeof(usr->addr)) == -1)
			log_warn("send_to failed: %s", strerror(errno));
	}
}

static struct node *
node_add(struct in_addr *addr)
{
	struct node *n;
	struct imsg_usrdns_req req;

	log_debug("user: node_add");

	n = xcalloc(1, sizeof(struct node));
	n->addr.s_addr = addr->s_addr;
	n->namelen = NODENAME_WAITING;
	n->lastseen = cap->time;
	LIST_INSERT_HEAD(&cap->node_list, n, entry);
	cap->node_count++;

	req.addr.s_addr = addr->s_addr;
	imsgev_compose(&cap->iev, IMSG_USRDNS_REQ, 0, 0, -1,
		&req, sizeof(req));

	return n;
}

static void
node_del(struct node *n)
{
	if (n->used)
		fatal("user: trying to remove a used node !");
	log_debug("user: node_del");

	LIST_REMOVE(n, entry);
	free(n->name);
	free(n);
	cap->node_count--;
}

static struct node *
node_find(struct in_addr *remote)
{
	struct node *n;

	LIST_FOREACH(n, &cap->node_list, entry)
		if (n->addr.s_addr == remote->s_addr)
			return n;
	return NULL;
}

static void
conn_add(struct in_addr *src, int src_port, struct in_addr *dst, int dst_port, int proto, int size)
{
	struct packet p;
	struct conn *c;
	struct node *srcnode;
	struct node *dstnode;
	int id;

	log_debug("user: conn_add, %x:%d->%x:%d %d [%d]",
		src->s_addr, src_port, dst->s_addr, dst_port, proto, size);
	if (cap->conn_freeids_ptr == CONN_FREEIDS_COUNT) {
		log_warn("user: out of connection identifiers !");
		return;
	}

	id = cap->conn_freeids[cap->conn_freeids_ptr];
	cap->conn_freeids_ptr++;

	srcnode = node_find(src);
	if (!srcnode)
		srcnode = node_add(src);
	srcnode->used++;
	dstnode = node_find(dst);
	if (!dstnode)
		dstnode = node_add(dst);
	dstnode->used++;

	c = xmalloc(sizeof(struct conn));
	c->id = id;
	c->state = CONNSTATE_ESTABLISHED;
	c->src = srcnode;
	c->src_port = src_port;
	c->dst = dstnode;
	c->dst_port = dst_port;
	c->proto = proto;
	c->size = size;
	c->lastseen = cap->time;
	LIST_INSERT_HEAD(&cap->conn_list, c, entry);

	bzero(&p, sizeof(p));
	p.ver = PACKET_VERSION;
	p.type = PACKET_NEWCONN;
	p.newconn_id = id;
	p.newconn_src = htonl(src->s_addr);
	p.newconn_dst = htonl(dst->s_addr);
	p.newconn_proto = proto;
	p.newconn_size = htons(size << 8);
	sendto_all(&p, PACKET_NEWCONN_SIZE);
}

static void
conn_data(struct conn *c, int size, int response)
{
	struct packet p;

	log_debug("user: conn_data");

	c->lastseen = cap->time;

	bzero(&p, sizeof(p));
	p.ver = PACKET_VERSION;
	p.type = PACKET_DATA;
	p.data_connid = c->id;
	p.data_size = htons(size << 8 | response); //XXX
	sendto_all(&p, PACKET_DATA_SIZE);
}

static void
conn_del(struct conn *c)
{
	struct packet p;

	log_debug("user: conn_del");

	if (c->proto == IPPROTO_TCP) {
		switch (c->state) {
		case CONNSTATE_ESTABLISHED:
			c->state = CONNSTATE_TCPFIN;
			return;
		case CONNSTATE_TCPFIN:
			c->state = CONNSTATE_TCPFIN2;
			return;
		case CONNSTATE_TCPFIN2:
			break;
		}
	}

	bzero(&p, sizeof(p));
	p.ver = PACKET_VERSION;
	p.type = PACKET_DELCONN;
	p.delconn_id = c->id;
	sendto_all(&p, PACKET_DELCONN_SIZE);

  if (cap->conn_freeids_ptr == 0)
    fatal("cap->conn_freeids_ptr == 0");
	cap->conn_freeids_ptr--;
	cap->conn_freeids[cap->conn_freeids_ptr] = c->id;

	c->src->used--;
	c->dst->used--;

	LIST_REMOVE(c, entry);
	free(c);
}

static pcap_handler
lookup_phandler(int type)
{
        struct phandler *p;

        for (p = phandlers; p->f; ++p) {
                if (type == p->type)
                        return p->f;
        }
        fatal("user: unknown data link type 0x%x", type);
        /* NOTREACHED */
	return NULL;
}

static struct user *
finduser(struct sockaddr_in *remote)
{
	struct user *usr;
	struct sockaddr_in *u;

	LIST_FOREACH(usr, &usr_list, entry) {
		u = &usr->addr;
		if (u->sin_addr.s_addr == remote->sin_addr.s_addr &&
			u->sin_port == remote->sin_port)
			return usr;
	}
	return NULL;
}

static void
phandler_ether(u_char *user, const struct pcap_pkthdr *h, const u_char *p)
{
	struct ether_header *ep;
	struct ip *ip;
	u_short ether_type;
	const u_char *pend;
	u_int	len;

	log_debug("user: pcap handler ethernet !");

	/* XXX here i assume that packets are alligned, which might not
	 * be the case when using dump files, says tcpdump sources */

	ep = (struct ether_header *)p;
	pend = p + h->caplen;
	len = h->len - sizeof(struct ether_header);


	ether_type = ntohs(ep->ether_type);
	if (ether_type <= ETHERMTU)
		log_tmp("llc packet !");
	else {
		switch (ether_type) {
		case ETHERTYPE_IP:
			log_tmp("ether IP");
			ip = (struct ip *)((u_char *)ep + sizeof(struct ether_header));
			ip_handle(ip, pend, len);
			break;
		default:
			log_tmp("non ip packet !");
			break;
		}
	}
}

static void
phandler_loop(u_char *user, const struct pcap_pkthdr *h, const u_char *p)
{
	struct ip *ip;
	struct ether_header *ep;
	u_short ether_type;
	u_int family;
	const u_char *pend;
	u_int	len;

	log_debug("user: pcap handler loop !");

	/* XXX here i assume that packets are alligned, which might not
	 * be the case when using dump files, says tcpdump sources */

	pend = p + h->caplen;
	len = h->len;

	memcpy((char *)&family, (char *)p, sizeof(family));
	family = ntohl(family);
	switch (family) {
	case AF_INET:
		log_tmp("loop family AF_INET");
		ip = (struct ip *)(p + NULL_HDRLEN);
		len -= NULL_HDRLEN;
		ip_handle(ip, pend, len);
		break;
#if defined(__OpenBSD__)
	case AF_LINK:
#else
	case AF_LOCAL:
#endif
		ep = (struct ether_header *)(p + NULL_HDRLEN);
		ether_type = ntohs(ep->ether_type);
		if (ether_type <= ETHERMTU)
			log_tmp("llc packet !");
		else {
			switch (ether_type) {
			case ETHERTYPE_IP:
				log_tmp("loop family AF_LINK IP");
				ip = (struct ip *)((u_char *)ep + sizeof(*ep));
				len -= NULL_HDRLEN + sizeof(*ep);
				ip_handle(ip, pend, len);
				break;
			default:
				log_tmp("loop non ip packet !");
				break;
			}
		}
	default:
		log_tmp("unknown family %x !", family);
		break;
	}
}

static void
ev_pcap(int fd, short why, void *data)
{
	log_tmp("ev_pcap");
	pcap_dispatch(cap->pcap, PCAP_COUNT, cap->pcap_handler, NULL);

	/* reschedule */
	if (event_add(&cap->pcap_ev, &cap->pcap_tv) == -1)
		fatal("user: event_add pcap failed : %s", strerror(errno));
}

static void
ev_timer(int fd, short why, void *data)
{
	struct conn *c, *ctmp;
	struct node *n, *ntmp;
	int i, to;

	log_debug("ev_timer");
	cap->time = time(NULL);

	i = 0;
	LIST_FOREACH_SAFE(c, &cap->conn_list, entry, ctmp) {
		switch (c->proto) {
		case IPPROTO_UDP:
			to = CONN_TIMEOUT_UDP;
			break;
		case IPPROTO_ICMP:
			to = CONN_TIMEOUT_ICMP;
			break;
		default:
			to = CONN_TIMEOUT;
			break;
		}
		if (cap->time > c->lastseen + to)
			conn_del(c);
		else
			i++;
	}

	if (cap->node_count > NODE_MAX_WITHOUT_TIMEOUT) {
		LIST_FOREACH_SAFE(n, &cap->node_list, entry, ntmp) {
			if (n->used == 0 &&
				cap->time > n->lastseen + NODE_TIMEOUT)
				node_del(n);
		}
	}

	log_debug("user: ev_timer leaving with %d active connections and %d active nodes", i, cap->node_count);
	if (event_add(&cap->conntimer_ev, &cap->conntimer_tv) == -1)
		fatal("user: event_add conntimer failed : %s", strerror(errno));
}

static void
usrconn(struct imsgev *iev, struct imsg *imsg)
{
	struct imsg_usrconn	*req;
	struct imsg_usrconn	res;
	struct user		*usr;

	req = imsg->data;

	if (req->deco) {
		usr = finduser(&req->addr);
		if (!usr)
			fatal("user: trying to deco an inexistant user !");
		LIST_REMOVE(usr, entry);
		free(usr);
		usr_count--;
		if (usr_count == 0)
			event_del(&cap->pcap_ev);

		return;
	}

	usr = xmalloc(sizeof(struct user));
	addrcpy(&usr->addr, &req->addr);
	LIST_INSERT_HEAD(&usr_list, usr, entry);
	if (usr_count == 0) {
		if (event_add(&cap->pcap_ev, &cap->pcap_tv) == -1)
			fatal("user: event_add failed : %s", strerror(errno));
	}
	usr_count++;

	addrcpy(&res.addr, &usr->addr);
	res.ok = 1;
	imsgev_compose(&cap->iev, IMSG_USRCONN_RES, 0, 0, -1,
		&res, sizeof(res));
}

static void
usrdns(struct imsgev *iev, struct imsg *imsg)
{
	struct imsg_usrdns_res *res;
	struct node *n;
	struct packet p;

	res = imsg->data;

	n = node_find(&res->addr);
	if (!n)
		fatal("user: received usrdns response for inexistant node !");
	if (n->namelen != NODENAME_WAITING)
		fatal("user: received usrdns response for a nonwaiting node!");

	if (res->name[0] == '\0') {
		log_debug("user: resolv for %x failed", res->addr.s_addr);
		n->namelen = NODENAME_FAILED;
		return;
	}
	n->namelen = strnlen(res->name, DNSNAME_MAX);
	n->name = strndup(res->name, DNSNAME_MAX);

	log_debug("user: sending node name of %x is %s",
		n->addr.s_addr, n->name);
	bzero(&p, sizeof(p));
	p.ver = PACKET_VERSION;
	p.type = PACKET_NAME;
	p.name_addr = htonl(n->addr.s_addr);
	p.name_len = n->namelen;
	strncpy(p.name_fqdn, n->name, sizeof(p.name_fqdn));
	sendto_all(&p, PACKET_NAME_SIZE + p.name_len);
}

static void
imsgev_main(struct imsgev *iev, int code, struct imsg *imsg)
{
	switch (code) {
	case IMSGEV_IMSG:
		log_debug("user: %s got imsg %i on fd %i",
			__func__, imsg->hdr.type, iev->ibuf.fd);
		switch (imsg->hdr.type) {
		case IMSG_USRCONN_REQ:
			usrconn(iev, imsg);
			break;
		case IMSG_USRDNS_RES:
			usrdns(iev, imsg);
			break;
		default:
			fatal("user: %s, unexpected imsg %d", 
				__func__, imsg->hdr.type);
		}
		break;
	case IMSGEV_EREAD:
	case IMSGEV_EWRITE:
	case IMSGEV_EIMSG:
		fatal("user: imsgev read/write error");
		/* NOTREACHED */
		break;
	case IMSGEV_DONE:
		event_loopexit(NULL);
		/* NOTREACHED */
		break;
	}
}

static void
imsgev_main_needfd(struct imsgev *iev)
{
    fatal("server: imsgev_main_needfd");
}