// dhcp_protect.c
//   Copyright 2019 Pascal Gloor
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <errno.h>
#include <time.h>
#include <stdarg.h>
#include <syslog.h>
#include <arpa/inet.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;

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

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

	openlog("dhcp_protect", LOG_PID, LOG_DAEMON);

	nfq_start(&conf);

	return 0;
}

// syslog function
void dp_log(unsigned char *remoteid, int remoteidlen, char *fmt, ...) {
	va_list argList;
	char buf[1000];
	int offset = remoteidlen*2;
	int i;

	for(i=0; i<remoteidlen; i++) {
		sprintf(buf+(i*2), "%02x", remoteid[i]);
	}
	buf[offset]=':'; offset++;
	buf[offset]=' '; offset++;

	va_start(argList, fmt);
	vsnprintf(buf+offset, sizeof(buf)-offset-1, fmt, argList);
	va_end(argList);

	syslog(LOG_DAEMON|LOG_INFO, "%s", buf);
}

// 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 if ( strcmp(name, "dryrun")==0 )
			conf->dryrun = atoi(value) ? 1 : 0;
		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 < 5 || conf->interval > 900 ) {
		fprintf(stderr,"interval value invalid (min 5, max 900)\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 ( conf->dryrun < 0 || conf->dryrun > 1 ) {
		fprintf(stderr, "dryrun value invalid (0 or 1)\n");
		error=1;
	}

	if ( error )
		return NULL;

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


	return conf;
}

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

	pktlen = nfq_get_payload(nfa, &pkt);

	if ( conf->debug ) printf("got a packet, len = %i\n", pktlen);

	// 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 ) {
			// jump to DHCPv4
			offset += ( ihl * 4 ) + 8;
			dhcpv4_check(conf, pkt, pktlen, offset, &remoteid, &remoteidlen);
		}
		else if ( ipver == 6 ) {
			// jump to DHCPv6
			offset += 48 + 8;
			dhcpv6_check(conf, pkt, pktlen, offset, &remoteid, &remoteidlen);
		}
		else {		
			if ( conf->debug ) printf("not an IPv4 packet\n");
			goto end;
		}
	}

	if ( remoteidlen>0 ) {
		// count the packet, even when blacklisted.
		dp_accounting_add(conf, remoteid, remoteidlen);

		// check if already in the blacklist
		if ( dp_blacklist_check(conf, remoteid, remoteidlen) == NF_DROP )
			rv = NF_DROP;

		// check if it must be added to the blacklist
		else if ( dp_accounting_check(conf, remoteid, remoteidlen) == NF_DROP ) {
			dp_blacklist_add(conf, remoteid, remoteidlen);
			rv = NF_DROP;
		}
	}

	end:

	dp_hash_cleanup(conf);

	return rv;
}

void dhcpv6_check(dp_conf *conf, unsigned char *pkt, int pktlen, int offset, unsigned char *remoteid, int *remoteidlen) {

	while(offset<pktlen) {
		uint8_t msgtype = (uint8_t)pkt[offset];

		switch(msgtype) {
			case 11: // RELAY-FORW
			case 12: // RELAY-REPL
				offset += 2 + 16 + 16; // msg-type, hop-count, link-addr, peer-addr
				break;
			default: // all other msgtypes
				offset += 2; // msg-type, hop-count
		}

		while(offset+1<pktlen) {
			uint8_t code = (uint8_t)pkt[offset];
			uint8_t len  = (uint8_t)pkt[offset+1];

			offset+2;

			if ( code == 1 ) { // Client identifier / DUID
				remoteid = pkt+offset;
				*remoteidlen = len;
				break;
			}
			offset+=len;
		}

		if ( *remoteidlen>0 )
			break;
	}
}


void dhcpv4_check(dp_conf *conf, unsigned char *pkt, int pktlen, int offset, unsigned char *remoteid, int *remoteidlen) {
	int hwaddrpos = offset + 28; // remember where the hw addr is if we need to fallback to it
	int hwlenpos = offset + 2;   // remember where the hw addr len is if we need to fallback to it

	// jump DHCP header
	offset += 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;
	}

	// 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;
	}
	offset+=4;


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

		offset++;

		// 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++;

		// 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");
						break;
					}
					else {
						remoteid = o82 + o82off;
						*remoteidlen = olen;
					}
				}

				o82off+=olen;
			}
		}

		offset+=len;
	}

	// if we didn't find opt82.2, we fallback to the hw addr
	if ( *remoteidlen == 0 ) {
		// nope, we won't overflow
		if ( (uint8_t)pkt[hwlenpos] <= 16 ) {
			remoteid = pkt+hwaddrpos;
			remoteidlen = (uint8_t)pkt[hwlenpos];
		}
	}
}

// netfilter queue callback
static int dp_callback(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg, struct nfq_data *nfa, void *data) {
	struct nfqnl_msg_packet_hdr *ph = nfq_get_msg_packet_hdr(nfa);
	dp_conf *conf = (dp_conf*)data;
	int id = -1;
	int verdict;

	if ( ph ) {
		id = ntohl (ph->packet_id);
		if ( conf->debug ) printf ("received packet with id %d\n", id);
	}

	verdict = dhcp_check(nfa, conf); /* Treat packet */

	// override decision for dryrun
	if ( conf->dryrun )
		verdict = NF_ACCEPT;

	return nfq_set_verdict(qh, id, verdict, 0, NULL); /* Verdict packet */
}

void dp_accounting_add(dp_conf *conf, unsigned char *remoteid, int len) {
	//int i;
	dp_accounting *ac;

	// does the element already exist
	HASH_FIND(hh, accountings, remoteid, len, ac);

	if ( conf->debug ) printf("AC: add item\n");

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

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

		HASH_ADD(hh, accountings, remoteid, len, ac);
	}
}

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

	// alrady exists?
	HASH_FIND(hh, blacklists, remoteid, len, bl);

	if ( conf->debug ) printf("BL: add item\n");

	// found an entry, push the expiration further
	if ( bl ) {
		if ( conf->debug ) printf("BL: item found -> pushing further\n");
		bl->expire = time(NULL) + conf->bltime;
	}
	// not found, create a new one
	else {
		if ( conf->debug ) printf("BL: item not found, new entry in BL\n");
		dp_log(remoteid, len, "blacklisting started");
		bl = malloc(sizeof(dp_blacklist));
		memcpy(bl->remoteid, remoteid, len);
		bl->len = len;
		bl->expire = time(NULL) + conf->bltime;
		HASH_ADD(hh, blacklists, remoteid, len, bl);
	}
}

void dp_hash_cleanup(dp_conf *conf) {
	dp_accounting *ac, *actmp;
	dp_blacklist *bl, *bltmp;

	// is it time to cleanup the list?
	// cleanup every conf->interval seconds 
	if ( dp_accountingtime + conf->interval < time(NULL) ) {
		if ( conf->debug ) printf("cleanup interval\n");
		dp_accountingtime = time(NULL);
		HASH_ITER(hh, accountings, ac, actmp) {
			HASH_DEL(accountings, ac);
			free(ac);
		}
	}

	// blacklist cleanup check every 1 sec
	if ( dp_cleanuptime < time(NULL) ) {
		if ( conf->debug ) printf("cleanup BL\n");
		dp_cleanuptime = time(NULL);
		HASH_ITER(hh, blacklists, bl, bltmp) {
			if ( bl->expire < time(NULL) ) {
				dp_log(
					bl->remoteid, bl->len, 
					"blacklisting ended");
				HASH_DEL(blacklists, bl);
				free(bl);
			}
		}
	}
}

int dp_accounting_check(dp_conf *conf, unsigned char *remoteid, int len) {
	dp_accounting *ac;

	HASH_FIND(hh, accountings, remoteid, len, ac);

	if ( conf->debug ) printf("AC Check\n");

	if ( ac ) {
		if(conf->debug) printf("AC Check: found item %i > %i ?\n", ac->count, conf->pktint);

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

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

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

	if ( conf->debug ) printf("BL Check\n");

	if ( bl ) {

		if ( bl->expire > time(NULL) ) {
			if ( conf->debug ) printf("blacklisted!\n");
			return NF_DROP;
		}
	}
	return NF_ACCEPT;
}