/*-
 * Copyright (c) 2005-2006 Robert N. M. Watson
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

/*
 * Simple tool for dumping TCP flows to libpcap save files.  Given an
 * interface and optional libpcap filter, dump flow information to data/
 * with useful flow names.  SIGQUIT, SIGHUP, and SIGTERM cause exit with a
 * dump of the current flow table; SIGINFO causes immediate dump, but should
 * be used with caution as it causes all current flows to be pruned.
 */

#include <sys/types.h>
#include <sys/signal.h>
#include <sys/stat.h>

#include <net/ethernet.h>

#define	__FAVOR_BSD						/* Linux. */

#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>

#include <arpa/inet.h>

#include <err.h>
#include <errno.h>
#include <limits.h>
#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>

#include "queue.h"
#include "packet.h"
#include "ethernet.h"
#include "tcpip.h"
#include "flow.h"

static pcap_t *pcap;
static int dumpnow, quit;

/*
 * TCP-specific flow flags.
 */
#define	FLOW_TCP_TTLCHANGE	0x00000002	/* TTL has changed. */
#define	FLOW_TCP_TTL2		0x00000004	/* TTL from ip2 valid. */
#define	FLOW_TCP_FIN1		0x00000008	/* FIN from ip1. */
#define	FLOW_TCP_FIN2		0x00000010	/* FIN from ip2. */
#define	FLOW_TCP_RST1		0x00000020	/* RST from ip1. */
#define	FLOW_TCP_RST2		0x00000040	/* RST from ip2. */

/*
 * Information passed down from the TCP parsing code in the pcap handler to
 * our implementations of flow methods via the flow opaque pointer 'args'.
 */
struct tcp_flow_args {
	struct ip	*tfa_ip;
	struct tcphdr	*tfa_tcp;
};

/*
 * Per-flow TCP-specific data.
 */
struct tcp_flow_data {
	struct in_addr	tfd_ip1;
	struct in_addr	tfd_ip2;
	u_short		tfd_port1;
	u_short		tfd_port2;
	u_char		tfd_ttl1;
	u_char		tfd_ttl2;
	u_int		tfd_flags;
};

/*
 * TCP implementations of flow class methods.
 */
static int	tcp_flow_alloc(struct flow *flow, void *args);
static void	tcp_flow_free(struct flow *flow);
static int	tcp_flow_matcher(struct flow *flow, struct packet *p,
		    void *args);
static void	tcp_flow_update(struct flow *flow, struct packet *p,
		    void *args);
static void	tcp_flow_name(struct flow *flow, char *path);
static int	tcp_flow_expire(struct flow *flow, time_t now);
static void	tcp_flow_print(struct flow *flow);

struct flow_class	tcp_flow_class = {
	.fc_flow_alloc = tcp_flow_alloc,
	.fc_flow_free = tcp_flow_free,
	.fc_flow_matcher = tcp_flow_matcher,
	.fc_flow_update = tcp_flow_update,
	.fc_flow_name = tcp_flow_name,
	.fc_flow_expire = tcp_flow_expire,
	.fc_flow_print = tcp_flow_print,
};

static void
usage(void)
{

	errx(EX_USAGE, "tcpflowdump interface [filter]");
}

static void
sighandler_quit(int sig)
{

	quit = 1;
}

#ifdef SIGINFO
static void
sighandler_dump(int sig)
{

	dumpnow = 1;
}
#endif

/*
 * Allocate TCP-specific per-flow state.  Initialize using TCP/IP headers.
 */
static int
tcp_flow_alloc(struct flow *flow, void *args)
{
	struct tcp_flow_data *tfd;
	struct tcp_flow_args *tfa;

	tfa = args;
	tfd = malloc(sizeof(*tfd));
	if (tfd == NULL)
		return (-1);
	bzero(tfd, sizeof(*tfd));
	tfd->tfd_ip1 = tfa->tfa_ip->ip_src;
	tfd->tfd_port1 = tfa->tfa_tcp->th_sport;
	tfd->tfd_ttl1 = tfa->tfa_ip->ip_ttl;
	tfd->tfd_ip2 = tfa->tfa_ip->ip_dst;
	tfd->tfd_port2 = tfa->tfa_tcp->th_dport;
	flow->flow_data = tfd;
	return (0);
}

/*
 * Free TCP-specific per-flow data.
 */
static void
tcp_flow_free(struct flow *flow)
{

	free(flow->flow_data);
}

/*
 * Flow matching routine for TCP: flows consist of two endpoints, each an IP
 * and port number.  Packets match flows if both endpoints match endpoints in
 * the TCP/IP headers.
 */
static int
tcp_flow_matcher(struct flow *flow, struct packet *p, void *args)
{
	struct tcp_flow_args *tfa;
	struct tcp_flow_data *tfd;

	tfa = args;
	tfd = flow->flow_data;
	if (tfd->tfd_ip1.s_addr == tfa->tfa_ip->ip_src.s_addr &&
	    tfd->tfd_port1 == tfa->tfa_tcp->th_sport &&
	    tfd->tfd_ip2.s_addr == tfa->tfa_ip->ip_dst.s_addr &&
	    tfd->tfd_port2 == tfa->tfa_tcp->th_dport)
		return (1);
	if (tfd->tfd_ip1.s_addr == tfa->tfa_ip->ip_dst.s_addr &&
	    tfd->tfd_port1 == tfa->tfa_tcp->th_dport &&
	    tfd->tfd_ip2.s_addr == tfa->tfa_ip->ip_src.s_addr &&
	    tfd->tfd_port2 == tfa->tfa_tcp->th_sport)
		return (1);
	return (0);
}

/*
 * Update a flow record based on the current TCP/IP header -- track whether
 * the flow has been reset, and whether the ttl has changed in either
 * direction.
 */
static void
tcp_flow_update(struct flow *flow, struct packet *p, void *args)
{
	struct tcp_flow_args *tfa;
	struct tcp_flow_data *tfd;
	int host;

	tfa = args;
	tfd = flow->flow_data;
	/*
	 * Decide which endpoint sent this packet.
	 */
	if (tfa->tfa_ip->ip_src.s_addr == tfd->tfd_ip1.s_addr &&
	    tfa->tfa_tcp->th_sport == tfd->tfd_port1)
		host = 1;
	else
		host = 2;

	/*
	 * If the TTL has changed, flag this flow.  Must figure out first
	 * which end point the TTL is from so we can update the right field.
	 */
	switch (host) {
	case 1:
		if (tfd->tfd_ttl1 != tfa->tfa_ip->ip_ttl) {
			tfd->tfd_ttl1 = tfa->tfa_ip->ip_ttl;
			tfd->tfd_flags |= FLOW_TCP_TTLCHANGE;
		}
		break;

	case 2:
		if (!(tfd->tfd_flags & FLOW_TCP_TTL2)) {
			tfd->tfd_ttl2 = tfa->tfa_ip->ip_ttl;
			tfd->tfd_flags |= FLOW_TCP_TTL2;
		} else if (tfd->tfd_ttl2 != tfa->tfa_ip->ip_ttl) {
			tfd->tfd_ttl2 = tfa->tfa_ip->ip_ttl;
			tfd->tfd_flags |= FLOW_TCP_TTLCHANGE;
		}
		break;

	default:
		errx(-1, "tcp_flow_update: host error");
	}

	/*
	 * Flag if we've seen FIN's and RST's in both directions.
	 */
	switch (host) {
	case 1:
		if (tfa->tfa_tcp->th_flags & TH_FIN)
			tfd->tfd_flags |= FLOW_TCP_FIN1;
		if (tfa->tfa_tcp->th_flags & TH_RST)
			tfd->tfd_flags |= FLOW_TCP_RST1;
		break;

	case 2:
		if (tfa->tfa_tcp->th_flags & TH_FIN)
			tfd->tfd_flags |= FLOW_TCP_FIN2;
		if (tfa->tfa_tcp->th_flags & TH_RST)
			tfd->tfd_flags |= FLOW_TCP_RST2;
		break;

	default:
		errx(-1, "tcp_flow_update: host error");
	}
}

/*
 * Select a name for the current flow for the purposes of creating a save
 * file.  Try to make it canonical by sorting the "lower" IP first.
 *
 * XXXRW: This API should be able to return an error, and have a way to
 * detect/handle conflicts.  Maybe a try counter?
 */
#define	IPMAX	15
static void
tcp_flow_name(struct flow *flow, char *path)
{
	char ip1[IPMAX], ip2[IPMAX];
	char trystring[PATH_MAX];
	struct tcp_flow_data *tfd;
	u_short port1, port2;
	struct stat sb;
	int ret, try;

	for (try = 0; try < 1000; try++) {
		tfd = flow->flow_data;
		if (ntohl(tfd->tfd_ip1.s_addr) < ntohl(tfd->tfd_ip2.s_addr)) {
			snprintf(ip1, IPMAX, "%s", inet_ntoa(tfd->tfd_ip1));
			snprintf(ip2, IPMAX, "%s", inet_ntoa(tfd->tfd_ip2));
			port1 = ntohs(tfd->tfd_port1);
			port2 = ntohs(tfd->tfd_port2);
		} else {
			snprintf(ip1, IPMAX, "%s", inet_ntoa(tfd->tfd_ip2));
			snprintf(ip2, IPMAX, "%s", inet_ntoa(tfd->tfd_ip1));
			port1 = ntohs(tfd->tfd_port2);
			port2 = ntohs(tfd->tfd_port1);
		}
		if (try != 0)
			snprintf(trystring, PATH_MAX, "-%d", try);
		else
			trystring[0] = '\0';
		snprintf(path, PATH_MAX, "data/%s:%d-%s:%d%s%s%s",
		    ip1, port1, ip2, port2, trystring,
		    tfd->tfd_flags & (FLOW_TCP_RST1 | FLOW_TCP_RST2) ?
		    "-rst" : "",
		    tfd->tfd_flags & FLOW_TCP_TTLCHANGE ? "-ttl" : "");

		ret = stat(path, &sb);
		if (ret < 0 && errno == ENOENT)
			break;
	}
	if (try == 1000)
		err(-1, "tcp_flow_name: too many collisions");
}

/*
 * Decide if a flow is ready to expire from the flow list.  Expire quicker if
 * we think both ends are done, but not if one end thinks it is done.
 */
static int
tcp_flow_expire(struct flow *flow, time_t now)
{
	struct tcp_flow_data *tfd;

	tfd = flow->flow_data;

	/* Expire faster if mutual fin. */
	if ((tfd->tfd_flags & FLOW_TCP_FIN1) &&
	    (tfd->tfd_flags & FLOW_TCP_FIN2))
		return (!(flow->flow_timestamp + 30 > now));

	/* Expire faster if mutual reset. */
	if ((tfd->tfd_flags & FLOW_TCP_RST1) &&
	    (tfd->tfd_flags & FLOW_TCP_RST2))
		return (!(flow->flow_timestamp + 30 > now));

	return (!(flow->flow_timestamp + 180 > now));
}

/*
 * Simple printer for TCP flows.
 */
static void
tcp_flow_print(struct flow *flow)
{
	struct tcp_flow_data *tfd;

	tfd = flow->flow_data;

	printf("%s:%d -> ", inet_ntoa(tfd->tfd_ip1), ntohs(tfd->tfd_port1));
	printf("%s:%d\n", inet_ntoa(tfd->tfd_ip2), ntohs(tfd->tfd_port2));
}

/*
 * Given a packet from pcap, find IP and TCP headers, feed to flow code.
 */
static void
handle_packet(const struct pcap_pkthdr *h, const u_char *bytes)
{
	struct tcp_flow_args tfa;
	struct ether_header eth;
	struct ip ip;
	struct tcphdr tcp;
	struct packet *p;

	/*
	 * Truncated by libpcap/bpf.
	 */
	if (h->len > h->caplen) {
		warnx("truncated packet");
		return;
	}

	p = packet_new_pkthdr(h->len, bytes, h);
	if (p == NULL) {
		warnx("packet_new");
		return;
	}

	if (packet_get_etherheader(p, &eth) < 0) {
		packet_free(p);
		warnx("packet_get_etherheader");
		return;
	}

	if (ntohs(eth.ether_type) != ETHERTYPE_IP) {
		packet_free(p);
		return;
	}

	if (packet_get_ip(p, &ip) < 0) {
		packet_free(p);
		warnx("packet_get_ip");
		return;
	}

	if (ip.ip_p != IPPROTO_TCP) {
		packet_free(p);
		return;
	}

	if (packet_get_tcp(p, &tcp) < 0) {
		packet_free(p);
		warnx("packet_get_tcp");
		return;
	}

	tfa.tfa_ip = &ip;
	tfa.tfa_tcp = &tcp;
	flow_new_packet(&tcp_flow_class, p, &tfa);
}

int
main(int argc, char *argv[])
{
	struct bpf_program bpf_program;
	char errbuf[PCAP_ERRBUF_SIZE];
	struct pcap_pkthdr pkthdr;
	const u_char *packet;

	if (argc < 2 || argc > 3)
		usage();

	if (geteuid() != 0)
		warnx("warning: euid not 0");
	if (getuid() != geteuid())
		warnx("warning: installed setuid");

	if (signal(SIGINT, sighandler_quit) == SIG_ERR)
		err(-1, "signal(SIGINT)");
	if (signal(SIGQUIT, sighandler_quit) == SIG_ERR)
		err(-1, "signal(SIGQUIT)");
	if (signal(SIGTERM, sighandler_quit) == SIG_ERR)
		err(-1, "signal(SIGTERM)");
#ifdef SIGINFO
	if (signal(SIGINFO, sighandler_dump) == SIG_ERR)
		err(-1, "signal(SIGINFO)");
#endif

	pcap = pcap_open_live(argv[1], 65535, 1, 100, errbuf);
	if (pcap == NULL)
		errx(EX_UNAVAILABLE, "pcap_open_live: %s", errbuf);

	if (pcap_set_datalink(pcap, DLT_EN10MB) < 0) {
		pcap_close(pcap);
		errx(EX_UNAVAILABLE, "pcap_set_datalink(DLT_EN10MB)");
	}

	if (argc == 3) {
		if (pcap_compile(pcap, &bpf_program, argv[2], 1, 0) < 0)
			errx(EX_USAGE, "pcap_compile");
		if (pcap_setfilter(pcap, &bpf_program) != 0)
			errx(EX_UNAVAILABLE, "pcap_setfilter");
	}

	flow_init(&tcp_flow_class);
	do {
		packet = pcap_next(pcap, &pkthdr);
		if (packet != NULL)
			handle_packet(&pkthdr, packet);
		if (dumpnow) {
			flow_dump(&tcp_flow_class, pcap,
			    FLOW_DUMP_PRINT | FLOW_DUMP_NOW);
			dumpnow = 0;
		} else {
			flow_dump(&tcp_flow_class, pcap, FLOW_DUMP_PRINT);
		}
	} while (!quit);
	flow_dump(&tcp_flow_class, pcap, FLOW_DUMP_PRINT | FLOW_DUMP_NOW);
	flow_destroy(&tcp_flow_class);

	return (0);
}
