#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/errno.h>
#include <sys/time.h>
#include <assert.h>
#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <arpa/inet.h>
#include <arpa/nameser.h>
#include <sys/resource.h>
#include <sys/stat.h>

#include <pcap.h>
#include "pcap_layers.h"

#define LIMIT_OPEN_FD 8192
#define LIMIT_CLT_PKTS 2048
#define LIMIT_MAXRSS (256<<20)
#define LIMIT_PKTS_IN_MEM (2<<20)
#define QUAD_A(ip) ((ntohl(ip.s_addr) >> 24) & 0xFF)
#define QUAD_B(ip) ((ntohl(ip.s_addr) >> 16) & 0xFF)
#define QUAD_C(ip) ((ntohl(ip.s_addr) >>  8) & 0xFF)
#define QUAD_D(ip) ((ntohl(ip.s_addr)      ) & 0xFF)

typedef struct _dlink_node dlink_node;
typedef struct _dlink_list dlink_list;

struct _dlink_node {
    void *data;
    dlink_node *prev;
    dlink_node *next;
};

struct _dlink_list {
    dlink_node *head;
    dlink_node *tail;
};

struct _packet {
    struct pcap_pkthdr hdr;
    void *data;
    struct _packet *next;
};

struct _client {
    struct in_addr addr;
    struct _packet *pkthead;
    struct _packet **pkttail;
    int npackets;
    pcap_dumper_t *fd;
    struct _client *next;
    dlink_node lru;
};

static struct _client *Hash[256];
static struct _dlink_list *LRU;
static unsigned int nopen = 3;
static unsigned int nclts = 0;
static unsigned int npacketsmem = 0;
static pcap_t *in = NULL;

void
dlinkAdd(void *data, dlink_node * m, dlink_list * list)
{
    m->data = data;
    m->prev = NULL;
    m->next = list->head;
    if (list->head)
	list->head->prev = m;
    list->head = m;
    if (list->tail == NULL)
	list->tail = m;
}

void
dlinkDelete(dlink_node * m, dlink_list * list)
{
    if (m->next)
	m->next->prev = m->prev;
    if (m->prev)
	m->prev->next = m->next;
    if (m == list->head)
	list->head = m->next;
    if (m == list->tail)
	list->tail = m->prev;
    m->next = m->prev = NULL;
}

void
hashDelete(struct _client *f)
{
    int i = (f->addr.s_addr >> 24) & 0xFF;
    struct _client **F;
    for (F = &Hash[i]; *F; F = &(*F)->next)
	if (f == *F) {
		*F = f->next;
		break;
	}
}

void
mksubdir(struct in_addr a)
{
    static char dir[128];
    snprintf(dir, 128, "%03d", QUAD_A(a));
    mkdir(dir, 0755);
    snprintf(dir, 128, "%03d/%03d", QUAD_A(a), QUAD_B(a));
    mkdir(dir, 0755);
    snprintf(dir, 128, "%03d/%03d/%s", QUAD_A(a), QUAD_B(a), inet_ntoa(a));
    mkdir(dir, 0755);
}

void
clt_pcap_open(struct _client *f)
{
    char file[128];

    if (NULL != f->fd)
	return;

    snprintf(file, 128, "%03d/%03d/%s/%lld.%06lld.pcap",
	QUAD_A(f->addr), QUAD_B(f->addr),
	inet_ntoa(f->addr),
	(long long int) f->pkthead->hdr.ts.tv_sec,
	(long long int) f->pkthead->hdr.ts.tv_usec);
    f->fd = pcap_dump_open(in, file);
    if (NULL == f->fd && errno == ENOENT) {
	mksubdir(f->addr);
	f->fd = pcap_dump_open(in, file);
    }
    if (NULL == f->fd) {
	perror(file);
	exit(1);
    }
    nopen++;
}


void
clt_free_packets(struct _client *f)
{
    struct _packet *p;
    struct _packet *n;
    for (p = f->pkthead; p; p = n) {
	n = p->next;
	free(p->data);
	free(p);
    }
    npacketsmem -= f->npackets;
    f->npackets = 0;
    f->pkthead = NULL;
    f->pkttail = &f->pkthead;
}

void
clt_pcap_write(struct _client *f)
{
    struct _packet *p;
    if (0 == f->npackets)
	return;
    for (p = f->pkthead; p; p = p->next)
	pcap_dump((void *)f->fd, &p->hdr, p->data);
    clt_free_packets(f);
}

void
clt_pcap_close(struct _client *f)
{
    if (NULL == f->fd)
	return;
    pcap_dump_close(f->fd);
    nopen--;
    f->fd = NULL;
}

void
clt_free(struct _client *f)
{
    if (f->npackets)
	clt_free_packets(f);
    hashDelete(f);
    dlinkDelete(&f->lru, LRU);
    free(f);
    nclts--;
}

void
close_lru(void)
{
    int nc = 0;
    dlink_node *p = LRU->tail;
    fprintf(stderr, "Closing LRU...");
    while (nopen > (LIMIT_OPEN_FD / 2) && p) {
	struct _client *f = p->data;
	p = p->prev;
	if (NULL == f->fd)
		continue;
	clt_pcap_write(f);
	clt_pcap_close(f);
	/* clt_free(f); */
	nc++;
    }
    fprintf(stderr, "%d\n", nc);
}

void
flush(struct _client *f)
{
    if (nopen >= LIMIT_OPEN_FD)
	close_lru();
    clt_pcap_open(f);
    clt_pcap_write(f);
}

long
getmaxrss(void)
{
    struct rusage ru;
    getrusage(RUSAGE_SELF, &ru);
    return ru.ru_maxrss;
}

void
flushall()
{
    int i;
    int n = 0;
    struct _client *f;
    struct _client *next;
    fprintf(stderr, "Flushing...\n");
    for (i = 0; i < 256; i++) {
	for (f = Hash[i]; f; f = next) {
	    next = f->next;
	    if (0 == f->npackets)
		continue;
	    if (f->fd == NULL)
		clt_pcap_open(f);
	    clt_pcap_write(f);
	    n++;
	    if (0 == (n % 1000))
		fprintf(stderr, "flushed %d clts, open fd: %d\n", n, nopen);
	    if (nopen >= LIMIT_OPEN_FD)
	        close_lru();
	}
    }
    fprintf(stderr, "flushed %d\n", n);
    fprintf(stderr, "open files: %d\n", nopen);
    fprintf(stderr, "max rss: %ld\n", getmaxrss());
}

void
stash2(struct _client *f, struct pcap_pkthdr *hdr, const unsigned char *data)
{
    struct _packet *p = calloc(1, sizeof(*p));
    assert(p);
    p->hdr = *hdr;
    p->data = malloc(hdr->caplen);
    assert(p->data);
    memcpy(p->data, data, hdr->caplen);
    *f->pkttail = p;
    f->pkttail = &p->next;
    f->npackets++;
    npacketsmem++;
#if 0
    if (f->npackets > LIMIT_CLT_PKTS)
	flush(f);
#endif
}

void
stash(struct in_addr a, struct pcap_pkthdr *hdr, const unsigned char *data)
{
    struct _client **F;
    struct _client *f;
    int i = (a.s_addr >> 24) & 0xFF;
    for (F = &Hash[i]; (f = *F); F = &(*F)->next) {
	if (f->addr.s_addr == a.s_addr)
	    break;
    }
    if (NULL == f) {
	nclts++;
	f = calloc(1, sizeof(*f));
	assert(f);
	f->addr = a;
	f->pkttail = &f->pkthead;
	f->next = Hash[i];
	Hash[i] = f;
    } else if (f != Hash[i]) {
	/* move to top */
	*F = f->next;
	f->next = Hash[i];
	Hash[i] = f;
    }
    stash2(f, hdr, data);
    dlinkDelete(&f->lru, LRU);
    dlinkAdd(f, &f->lru, LRU);
}

void
print_stats(struct timeval ts, int pkt_count)
{
    struct timeval now;
    gettimeofday(&now, NULL);
    fprintf(stderr, "%ld.%03ld: at %ld, %10d pkts, %9d clts, %4d files\n",
	now.tv_sec,
	now.tv_usec / 1000,
	(long)ts.tv_sec,
	pkt_count,
	nclts,
	nopen);
}

int
my_ip4_handler(const struct ip *ip4, int len, void *userdata)
{
    struct in_addr *a = userdata;
    *a = ip4->ip_src;
    return 0;
}


int
main(int argc, char *argv[])
{
    unsigned int pkt_count = 0;
    char errbuf[PCAP_ERRBUF_SIZE + 1];
    struct pcap_pkthdr hdr;
    const u_char *data;

    setbuf(stdout, NULL);
    setbuf(stderr, NULL);
    memset(Hash, '\0', sizeof(Hash));
    LRU = calloc(1, sizeof(LRU));
    assert(LRU);

    in = pcap_open_offline("-", errbuf);
    if (NULL == in) {
	fprintf(stderr, "stdin: %s", errbuf);
	exit(1);
    }
    pcap_layers_init(pcap_datalink(in));
    callback_ipv4 = my_ip4_handler;

    while ((data = pcap_next(in, &hdr))) {
	struct in_addr src;
	memset(&src, 0, sizeof(src));
	handle_pcap((u_char *) & src, &hdr, data);
	if (src.s_addr == 0)
	    continue;
	stash(src, &hdr, data);
	if (0 == (++pkt_count & 0x3FFF)) {
	    print_stats(hdr.ts, pkt_count);
	    if (npacketsmem > LIMIT_PKTS_IN_MEM) {
		flushall();
		close_lru();
	    }
	}
    }
    flushall();
    fprintf(stderr, "Max RSS  %ld KB \n", getmaxrss());
    return 0;
}