// dhcp_protect.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <errno.h>
#include <time.h>
#include <linux/netfilter.h>
#include <libnetfilter_queue/libnetfilter_queue.h>
#include <uthash.h>
#include "dhcp_protect.h"

// main function
int main(int argc, char **argv) {
	char *configfile;
	dp_conf conf;
	struct nfq_q_handle *qh;

	if ( argc == 2 ) {
		configfile = argv[1];
	}
	else {
		usage(argv[0]);
		return EXIT_FAILURE;
	}

	if ( load_config(&conf, configfile) == NULL )
		return EXIT_FAILURE;

	nfq_start(&conf);

	return 0;
}

// start netfilter queue
void nfq_start(dp_conf *conf) {
	struct nfq_handle *h;
	struct nfq_q_handle *qh;
	int fd;

	if ( ( h = nfq_open() ) == NULL ) {
		fprintf(stderr,"error during nfq_open() %s\n", strerror(errno));
		return;
	}

	if ( ( qh = nfq_create_queue(h, conf->queue, &dp_callback, conf) ) == NULL ) {
		fprintf(stderr, "error during nfq_create_queue() %s\n", strerror(errno));
		return;
	}

	if ( nfq_set_mode(qh, NFQNL_COPY_PACKET, 1500) < 0 ) {
		fprintf(stderr,"error during nfq_set_mode() %s\n", strerror(errno));
		return;
	}

	if ( nfq_set_queue_flags(qh, NFQA_CFG_F_FAIL_OPEN, NFQA_CFG_F_FAIL_OPEN) < 0 ) {
		fprintf(stderr,"error during nfq_set_queue_flags() %s\n", strerror(errno));
		return;
	}

	fd = nfq_fd(h);

	while(1) {
		static char buf[65536];
		int rv;

		if ((rv = recv(fd, buf, sizeof(buf), 0)) >= 0) {
			nfq_handle_packet(h, buf, rv); /* send packet to callback */
		}
	}
}

// display usage
void usage(char *prog) {
	fprintf(stderr,"Usage: %s <configuration file>\n",prog);
}


// load configuration file
dp_conf *load_config(dp_conf *conf, char *file) {
	FILE *fh;
	char *line = NULL;
	size_t len = 0;
	int error = 0;

	printf("Loading configuration %s\n",file);
	

	if ( ( fh = fopen(file,"r") ) == NULL ) {
		fprintf(stderr,"Failed to open configuration file '%s': %s\n", file, strerror(errno));
		return NULL;
	}

	while (getline(&line, &len, fh) != -1) {
		char *name, *value;

		name  = strtok(line, "=\r\n ");
		value = strtok(NULL,"=\r\n ");

		if ( name == NULL || value == NULL || name[0] == '#' )
			continue;

		if ( strcmp(name,"max_pkt_per_interval")==0 )
			conf->pktint = atoi(value);
		else if ( strcmp(name,"interval")==0 )
			conf->interval = atoi(value);
		else if ( strcmp(name, "debug")==0 )
			conf->debug = atoi(value) ? 1 : 0;
		else if ( strcmp(name, "blacklist_time")==0 )
			conf->bltime = atoi(value);
		else if ( strcmp(name, "queue")==0 )
			conf->queue = atoi(value);
		else
			fprintf(stderr,"unknown directive '%s', ignored\n", name);

		free(line);
	}
	fclose(fh);

	if ( conf->pktint < 1 || conf->pktint > 1000 ) {
		fprintf(stderr,"max_pkt_per_interval value invalid (min 1, max 1000)\n");
		error=1;
	}
	if ( conf->interval < 10 || conf->interval > 60 ) {
		fprintf(stderr,"interval value invalid (min 10, max 60)\n");
		error=1;
	}
	if ( conf->debug < 0 || conf->debug > 1 ) {
		fprintf(stderr,"debug value invalid (0 or 1)\n");
		error=1;
	}
	if ( conf->bltime < 10 || conf->bltime > 900 ) {
		fprintf(stderr,"blacklist_time value invalid (min 10, max 900)\n");
		error=1;
	}
	if ( conf->queue < 0 ) {
		fprintf(stderr,"queue must be a positive integer\n");
		error=1;
	}

	if ( error )
		return NULL;

	printf("Configuration:\n");
	printf("\t%-20s = %4i\n", "max_pkt_per_interval", conf->pktint);
	printf("\t%-20s = %4i\n", "interval", conf->interval);
	printf("\t%-20s = %4i\n", "debug", conf->debug);
	printf("\t%-20s = %4i\n", "blacklist_time", conf->bltime);
	printf("\t%-20s = %4i\n", "queue", conf->queue);


	return conf;
}

// decode dhcp packet
u_int32_t dhcp_check(struct nfq_data *nfa, int *verdict, dp_conf *conf) {
	unsigned char *pkt;
	int pktlen;
	int offset = 0;
	unsigned char *remoteid;
	int remoteidlen = 0;
	int found = 0;
	uint8_t ipver = 0;
	uint8_t ihl = 0;
	int i;

	pktlen = nfq_get_payload(nfa, &pkt);

	for(i=0; i<pktlen; i++) {
		if ( !(i%16) ) printf("\n");
		printf("%02x ", pkt[i]);
	}

	// can we read the IP proto and IP header length ?
	if ( pktlen > 0 ) {
		ipver = pkt[offset];
		ipver >>= 4;
		ihl   = pkt[offset];
		ihl  &= 0x0f;

		if ( ipver != 4 ) {
			if ( conf->debug ) printf("not an IPv4 packet\n");
			return NF_ACCEPT;
		}

		// jump over the IPv4 header
		offset += ihl * 4;
	}


	// jump over UDP + DHCP header
	offset += 8 + 28 + 16 + 64 + 128;

	// minimum packet size, fixed header + magic cookie (4 octets)
	if ( pktlen < offset + 4 ) {
		if ( conf->debug ) printf("packet too small\n");
		return NF_ACCEPT;
	}

	// check magic cookie
	if ( pkt[offset] != 99 || pkt[offset+1] != 130 || pkt[offset+2] != 83 || pkt[offset+3] != 99 ) {
		if ( conf->debug ) printf("invalid magic cookie %02x%02x%02x%02x\n", pkt[offset], pkt[offset+1], pkt[offset+2], pkt[offset+3]);
		return NF_ACCEPT;
	}
	offset+=4;


	// parse TLV options
	while(offset<pktlen && !found) {
		uint8_t type = pkt[offset];
		uint8_t len;

		offset++;

		printf("type : %i\n",type);

		// padding of 1 octet
		if ( type == 0 ) {
			continue;
		}

		// end of options
		if ( type == 255 )
			break;

		// make sure we can read len
		if ( offset>=pktlen )
			break;

		len = pkt[offset];
		offset++;

		printf("len  : %i\n", len);

		// can the value be read
		if ( offset+len>=pktlen )
			break;

		// option 82 parser
		if ( type == 82 ) {
			unsigned char *o82 = pkt+offset;
			int o82off = 0;

			// loop until the end, +2 to ensure we can read type and length
			while(o82off+2<len && !found) {
				uint8_t otype = o82[o82off];
				uint8_t olen  = o82[o82off+1];
				o82off+=2;

				printf("o82 type=%i len=%i\n", otype, olen);

				// remoteid
				if ( otype == 2 ) {
					// make sure we don't overflow and can read all data
					if ( o82off + olen > len ) {
						if ( conf->debug) printf("option 82.2 data too long\n");
						return NF_ACCEPT;
					}
					else {
						remoteid = o82 + o82off;
						remoteidlen = olen;
						found=1;
					}
				}

				o82off+=olen;
			}
		}

		offset+=len;
	}

	if ( found ) {
		dp_blacklist_count(conf, remoteid, remoteidlen);
		if ( dp_blacklist_check(conf, remoteid, remoteidlen) )
			return NF_ACCEPT;
		else
			return NF_DROP;
	}

	printf("got a packet, len = %i\n", pktlen);
	return NF_ACCEPT;
}

// netfilter queue callback
static int dp_callback(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg, struct nfq_data *nfa, void *data) {
	dp_conf *conf = (dp_conf*)data;
	int verdict;
	u_int32_t id = dhcp_check(nfa, &verdict, conf); /* Treat packet */
	return nfq_set_verdict(qh, id, verdict, 0, NULL); /* Verdict packet */
}

void dp_blacklist_count(dp_conf *conf, unsigned char *remoteid, int len) {
	int i;
	dp_blacklist *bl, *tmp;

	// does the element already exist
	HASH_FIND(hh, blacklists, remoteid, len, bl);

	// found it, increment the counter
	if ( bl ) {
		if ( conf->debug ) printf("BL: item found, incrementing\n");
		bl->count++;
	}
	// not found, create a new one
	else {
		bl = malloc(sizeof(dp_blacklist));

		memcpy(bl->remoteid, remoteid, len);
		bl->len = len;
		bl->count = 1;

		if ( conf->debug ) printf("BL: item not found, creating\n");

		HASH_ADD(hh, blacklists, remoteid, len, bl);
	}

	printf("count this remoteid ");
	for(i=0; i<len; i++) printf("%02x",*(uint8_t*)(remoteid+i));
	printf("\n");

	// is it time to cleanup the list?
	if ( dp_timestamp + conf->interval < time(NULL) ) {
		if ( conf->debug ) printf("cleanup interval\n");
		dp_timestamp = time(NULL);
		HASH_ITER(hh, blacklists, bl, tmp) {
			HASH_DEL(blacklists, bl);
			free(bl);
		}
	}
}

int dp_blacklist_check(dp_conf *conf, unsigned char *remoteid, int len) {
	dp_blacklist *bl;

	HASH_FIND(hh, blacklists, remoteid, len, bl);

	if ( bl ) {
		if(conf->debug) printf("found item\n");

		if ( bl->count > conf->pktint ) {
			if(conf->debug) printf("flood detected!\n");
			return NF_DROP;
		}
	}
}