/*-
 * Copyright (c) 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.
 */

/*
 * Implementation of a flow tracking system.  Flow classes are registered by
 * the consumer via a call to flow_init(), and are described by struct
 * flow_class.  The flow code maintains its own data structure,
 * flow_class_data, hung off of the public flow class structure.
 *
 * TODO:
 *
 * - Would be useful to have hashed flow lists, using a fc_hash() method
 *   implemented by the consumer, who is aware of the packet format and
 *   addressing scheme.
 * - Make thread-safe so that multi-threaded consumers don't have to
 *   provide there own locking.
 */

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

#include <err.h>
#include <limits.h>
#include <pcap.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

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

/*
 * Private per-flow information.
 */
struct flow_class_data {
	TAILQ_HEAD(, flow)	fcd_flow_list;
};

/*
 * Initialize a flow class.
 */
int
flow_init(struct flow_class *fc)
{
	struct flow_class_data *fcd;

	fcd = malloc(sizeof(*fcd));
	if (fcd == NULL)
		return (-1);
	TAILQ_INIT(&fcd->fcd_flow_list);
	fc->fc_data = fcd;
	return (0);
}

/*
 * Tear down a flow class.
 */
void
flow_destroy(struct flow_class *fc)
{

	free(fc->fc_data);
}

/*
 * Find an existing flow using the consumer-provided matching function.  If
 * we did hash chains, we would first ask the consumer to provide a hash key.
 */
struct flow *
flow_find(struct flow_class *fc, struct packet *p, void *args)
{
	struct flow_class_data *fcd;
	struct flow *flow;

	fcd = fc->fc_data;
	TAILQ_FOREACH(flow, &fcd->fcd_flow_list, flow_list) {
		if (fc->fc_flow_matcher(flow, p, args))
			return (flow);
	}
	return (NULL);
}

/*
 * Allocate a new flow, let the consumer initialize its fields.
 */
struct flow *
flow_new(struct flow_class *fc, void *args)
{
	struct flow_class_data *fcd;
	struct flow *flow;

	fcd = fc->fc_data;
	flow = malloc(sizeof(*flow));
	if (flow == NULL)
		return (NULL);
	bzero(flow, sizeof(*flow));
	if (fc->fc_flow_alloc(flow, args) == -1) {
		free(flow);
		return (NULL);
	}
	TAILQ_INIT(&flow->flow_packets);
	TAILQ_INSERT_HEAD(&fcd->fcd_flow_list, flow, flow_list);
	return (flow);
}

/*
 * Update a flow for the arrival of a new packet; if necessary, allocate a
 * new flow.
 */
void
flow_new_packet(struct flow_class *fc, struct packet *p, void *args)
{
	struct flow *flow;

	flow = flow_find(fc, p, args);
	if (flow == NULL) {
		flow = flow_new(fc, args);
		if (flow == NULL) {
			warnx("malloc");
			return;
		}
	}

	flow->flow_timestamp = time(NULL);
	fc->fc_flow_update(flow, p, args);
	TAILQ_INSERT_TAIL(&flow->flow_packets, p, p_list);
}

/*
 * Free a flow -- free all packets, and allow the consumer to clear up its
 * state.
 */
void
flow_drop(struct flow_class *fc, struct flow *flow)
{
	struct flow_class_data *fcd;
	struct packet *p;

	fcd = fc->fc_data;
	TAILQ_REMOVE(&fcd->fcd_flow_list, flow, flow_list);
	while ((p = TAILQ_FIRST(&flow->flow_packets)) != NULL) {
		TAILQ_REMOVE(&flow->flow_packets, p, p_list);
		packet_free(p);
	}
	fc->fc_flow_free(flow);
	free(flow);
}

/*
 * Print out a summary of all flows, using consumer flow print routine.
 */
void
flow_summary(struct flow_class *fc)
{
	struct flow_class_data *fcd;
	struct flow *flow;

	fcd = fc->fc_data;
	TAILQ_FOREACH(flow, &fcd->fcd_flow_list, flow_list)
		fc->fc_flow_print(flow);
}

/*
 * Dump information on and prune flows with no activity to save files  Note
 * that libpcap doesn't take appending to current files well, so we might
 * want to change the logic here so that we open exlusively and choose a
 * different name if there's a collision on flow file naming.  We including
 * information on the flow flags in the filename, but it might be better to
 * write this to a different file.
 */
#define IPMAX 15
void
flow_dump(struct flow_class *fc, pcap_t *pcap, int flags)
{
	struct flow *flow, *flow_temp;
	struct flow_class_data *fcd;
	pcap_dumper_t *outdumper;
	char filename[PATH_MAX];
	struct packet *p;
	time_t now;

	fcd = fc->fc_data;
	(void)mkdir("data", 0755);

	now = time(NULL);
	TAILQ_FOREACH_SAFE(flow, &fcd->fcd_flow_list, flow_list, flow_temp) {
		if (!(flags & FLOW_DUMP_NOW)) {
			if (!fc->fc_flow_expire(flow, now))
				continue;
		}

		fc->fc_flow_name(flow, filename);
		if (flags & FLOW_DUMP_PRINT)
			printf("%s\n", filename);

		outdumper = pcap_dump_open(pcap, filename);
		if (outdumper == NULL) {
			warnx("pcap_dump_open: %s", pcap_geterr(pcap));
			continue;
		}

		TAILQ_FOREACH(p, &flow->flow_packets, p_list) {
			pcap_dump((u_char *)outdumper ,&p->p_pkthdr,
			    p->p_data);
		}

		pcap_dump_close(outdumper);
		flow_drop(fc, flow);
	}
}
