#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <linux/types.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <getopt.h>
#include <netdb.h>

#include "ct_flush.h"

/* name of device (major 10, minor 243) */
#define	CT_FLUSH_DEVICE	"/dev/ct_flush"

/* default pattern: match no IP, all ports */
#define PATTERN_INIT { {0, 0, 0, 0xffff}, {0, 0, 0, 0xffff}, 0, 0 }

int exec_flush(struct search_pattern *hostpat) {
	struct search_pattern pat;
	int dev, ret = 0;

	if ((dev = open(CT_FLUSH_DEVICE, O_RDWR)) == -1 ) {
		perror("open");
		return -1;
	}

	pat.count = 0;
	pat.proto = hostpat->proto;
	pat.src.ip = htonl(hostpat->src.ip);
	pat.src.mask = htonl(hostpat->src.mask);
	pat.src.portlo = htons(hostpat->src.portlo);
	pat.src.porthi = htons(hostpat->src.porthi);
	pat.dst.ip = htonl(hostpat->dst.ip);
	pat.dst.mask = htonl(hostpat->dst.mask);
	pat.dst.portlo = htons(hostpat->dst.portlo);
	pat.dst.porthi = htons(hostpat->dst.porthi);

	if ((ret = ioctl(dev, IOCTL_CT_FLUSH, &pat)) == -1)
		perror ("ioctl");

	hostpat->count = pat.count;

	close(dev);

	return ret;
}

int parse_proto(char *arg, u_int8_t *val) {
	struct protoent *p;

	p = getprotobyname(arg);
	if (p != NULL) {
		*val = (u_int8_t)(p->p_proto);

	} else if (!strcmp(arg, "all")) {
		*val = 0;

	} else {
		long num = strtol(arg, NULL, 0);
		if (num <= 0 || num > 255)
			return 1;

		*val = (u_int8_t)num;
	}
	return 0;
}

void print_proto(char *pre, u_int8_t val) {
	struct protoent *p;

	if (val) {
		if ((p = getprotobynumber(val))) {
			printf("%s%s ", pre, p->p_name);
		} else {
			printf("%s%d ", pre, p->p_proto);
		}
	}
}

int parse_ipmask(char *arg, u_int32_t *ip, u_int32_t *mask) {
	struct in_addr tmp;
	char *p;

	/* netmask */
	if ((p = strchr(arg, '/')) != NULL) {
		if (inet_pton(AF_INET, p+1, &tmp) > 0) {
			*mask = (u_int32_t)(tmp.s_addr);

		} else {
			long num = strtol(p+1, NULL, 0);
			if (num <= 0 || num > 32)
				return 1;

			*mask = (u_int32_t) htonl(0xFFFFFFFF<<(32-num));
		}
		*p = 0;

	} else {
		*mask = 0xFFFFFFFF;
	}

	/* ip address */
	if (inet_pton(AF_INET, arg, &tmp) <= 0)
		return 1;

	*ip = (u_int32_t)(tmp.s_addr);

	return 0;
}

void print_ipmask(char *pre, u_int32_t ip, u_int32_t mask) {
	if (mask != 0 && ip != 0) {
		int i = 32;
		char ip_str[INET_ADDRSTRLEN];
		u_int32_t bits = 0xFFFFFFFEL;
		u_int32_t maskaddr = ntohl(mask);

		inet_ntop(AF_INET, &ip, ip_str, sizeof(ip_str));

		while (--i >= 0 && maskaddr != bits)
			bits <<= 1;

		if (bits == 0) {
			printf("%s%s ", pre, ip_str);

		} else if (i >= 0) {
			printf("%s%s/%d ", pre, ip_str, i);

		} else {
			char mask_str[INET_ADDRSTRLEN];
			inet_ntop(AF_INET, &mask, mask_str, sizeof(ip_str));
			printf("%s%s/%s ", pre, ip_str, mask_str);
		}
	}
}

int parse_portrange(char *arg, u_int16_t *portlo, u_int16_t *porthi) {
	long num;
	char *p;

	/* check for portrange */
	if ((p = strchr(arg, ':')) != NULL || (p = strchr(arg, '-')) != NULL) {
		/* try parse port */
		num = strtol(p+1, NULL, 0);
		if (num < 0 || num > 65535)
			return 1;

		*porthi = (u_int16_t)num;
		*p = 0;
	}

	/* try parse port */
	num = strtol(arg, NULL, 0);
	if (num < 0 || num > 65535)
		return 1;

	*portlo = (u_int16_t)num;

	if (p == NULL)
		*porthi = *portlo;

	if (portlo > porthi)
		return 1;

	return 0;
}

void print_portrange(char *pre, u_int16_t portlo, u_int16_t porthi) {
	if (portlo == porthi)
		printf("%s%d ", pre, portlo);

	else if (portlo != 0 || porthi != 65535)
		printf("%s%d:%d ", pre, portlo, porthi);
}

static struct option opts[] = {
	{"protocol",	1, 0, 'p'},
	{"source",	1, 0, 's'},
	{"src",		1, 0, 's'},
	{"sport",	1, 0, 'a'},
	{"destination",	1, 0, 'd'},
	{"dst",		1, 0, 'd'},
	{"dport",	1, 0, 'b'},
	{"help",	0, 0, 'h'},
	{"verbose",	0, 0, 'v'},
	{0, 0, 0, 0}
};

int main(int argc, char *argv[]) {
	struct search_pattern pat = PATTERN_INIT;
	int arg = 0, code = 0, verbose = 0;

	if (argc == 1) {
		fprintf(stderr, "No argument given. try --help\n");
		exit(-1);
	}

	while (code != -1) {
		code = getopt_long(argc, argv, "p:s:d:hv", opts, &arg);

		switch (code) {
		case 'p':	/* proto */
				if (parse_proto(optarg, &pat.proto)) {
					fprintf(stderr, "Error parsing value: %s\n", optarg);
					exit(-1);
				}
				break;

		case 's':	/* src */
				if (parse_ipmask(optarg, &pat.src.ip, &pat.src.mask)) {
					fprintf(stderr, "Error parsing value: %s\n", optarg);
					exit(-1);
				}
				break;

		case 'a':	/* sport */
				if (pat.proto != IPPROTO_TCP && pat.proto != IPPROTO_UDP) {
					fprintf(stderr, "--sport only allowed with tcp/udp protocol\n");
					exit(-1);
				}
				if (parse_portrange(optarg, &pat.src.portlo, &pat.src.porthi)) {
					fprintf(stderr, "Error parsing value: %s\n", optarg);
					exit(-1);
				}
				break;

		case 'd':	/* dst */
				if (parse_ipmask(optarg, &pat.dst.ip, &pat.dst.mask)) {
					fprintf(stderr, "Error parsing value: %s\n", optarg);
					exit(-1);
				}
				break;

		case 'b':	/* dport */
				if (pat.proto != IPPROTO_TCP && pat.proto != IPPROTO_UDP) {
					fprintf(stderr, "--dport only allowed with tcp/udp protocol\n");
					exit(-1);
				}
				if (parse_portrange(optarg, &pat.dst.portlo, &pat.dst.porthi)) {
					fprintf(stderr, "Error parsing value: %s\n", optarg);
					exit(-1);
				}
				break;

		case 'h':	/* help */
				printf("Usage: ct_flush [options]\n"
					"Options: \n"
					"  --protocol     -p  protocol        ('tcp', '6')\n"
					"  --source       -s  address[/mask]\n"
					"  --destination  -d  address[/mask]\n"
					"  --sport            port[-port]     (only with proto tcp/udp)\n"
					"  --dport            port[-port]     (only with proto tcp/udp)\n"
					"  --verbose      -v                  be verbose\n"
					"  --help         -h                  this help\n"
					"\n");
				exit(0);
				break;

		case 'v':	/* verbose */
				verbose = 1;
				break;

		case '?':	/* error */
				exit(-1);
				break;

		default:	/* unknown / all options parsed */
				break;
		}
	}

	if (verbose) {
		printf("flushing: ");
		print_proto("proto=", pat.proto);
		print_ipmask("src=", pat.src.ip, pat.src.mask);
		print_portrange("sport=", pat.src.portlo, pat.src.porthi);
		print_ipmask("dst=", pat.dst.ip, pat.dst.mask);
		print_portrange("dport=", pat.dst.portlo, pat.dst.porthi);
		printf("\n");
	}

	exec_flush(&pat);

	if (verbose)
		printf("%d flushed\n", pat.count);

	return 0;
}