#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <sqlite3.h>

#define DATABASE	"urlfilter.db"
#define BUF_SIZE	1024

struct callback_args {
	char *urlstr;
	int  urlstr_len;
	int flags;
};

static int daemon_callback(void *data, int cols, char **values, char **colnames)
{
	struct callback_args *args = data;

	char *dburl = values[0];
	int dburl_len = strlen(dburl);

	if (args->urlstr_len >= dburl_len) {
		char *p = args->urlstr;
		p += (args->urlstr_len - dburl_len);

		if (!strncmp(p , dburl, dburl_len))
			args->flags |= 1;
	}

	return 0;
}

static int daemon_loop(sqlite3 *db)
{
	char *query = malloc(BUF_SIZE);
	if (query == NULL) {
		perror("malloc():");
		return -1;
	}

	char *line = malloc(BUF_SIZE);
	if (line == NULL) {
		perror("malloc():");
		return -1;
	}

	setvbuf(stdout, NULL, _IOLBF, 0);

	while (fgets(line, BUF_SIZE, stdin)) {
		int ret;

		char *cp = strchr(line, '\n');
		if (cp != NULL)
			*cp= '\0';

		char *tmp = strtok(line, " \t");
		char *ipstr = strtok(NULL, " \t");
		char *urlstr = strtok(NULL, " \t");
		if (tmp == NULL || ipstr == NULL || urlstr == NULL) {
			puts("ERR");
			continue;
		}

		struct in_addr addr;
		ret = inet_aton(ipstr, &addr);
		if (ret == 0) {
			puts("ERR");
			continue;
		}

		ret = snprintf(query, BUF_SIZE, "SELECT url FROM urlfilter WHERE (ip & mask) == (%u & mask);", addr.s_addr);
		if (ret < 0 || ret >= BUF_SIZE) {
			puts("ERR");
			continue;
		}

		struct callback_args args;
		args.urlstr = urlstr;
		args.urlstr_len = strlen(args.urlstr);
		args.flags = 0;

		ret = sqlite3_exec(db, query, daemon_callback, &args, NULL);
		if (ret != SQLITE_OK) {
			puts("ERR");
			continue;
		}

		if (args.flags)
			puts("ERR");
		else
			puts("OK");
	}
	return 0;
}

static int show_callback(void *data, int cols, char **values, char **colnames)
{
	struct in_addr ip;
	struct in_addr mask;
	ip.s_addr = atoi(values[0]);
	mask.s_addr = atoi(values[1]);

	printf("%15s/", inet_ntoa(ip));
	printf("%-15s ", inet_ntoa(mask));
	printf("%s (%s)\n", values[2], values[3]);
	return 0;
}

static void show_db(sqlite3 *db)
{
	char *errstr;

	int ret = sqlite3_exec(db, "SELECT ip, mask, url, info FROM urlfilter", show_callback, NULL, &errstr);
	if (ret != SQLITE_OK) {
		fprintf(stderr, "SQL error: %s\n", errstr);
		sqlite3_free(errstr);
	}
}

static void add_db(sqlite3 *db, char *ipstr, char *maskstr, char *url, char *info)
{
	int ret;

	char *query = malloc(BUF_SIZE);
	if (query == NULL) {
		perror("malloc():");
		return;
	}

	struct in_addr ip;
	ret = inet_aton(ipstr, &ip);
	if (ret == 0) {
		fprintf(stderr, "invalid IP");
		return;
	}

	struct in_addr mask;
	ret = inet_aton(maskstr, &mask);
	if (ret == 0) {
		fprintf(stderr, "invalid Mask");
		return;
	}

	ret = snprintf(query, BUF_SIZE, "INSERT INTO urlfilter VALUES(NULL,'%u','%u','%s','%s');", ip.s_addr, mask.s_addr, url, info);
	if (ret < 0 || ret >= BUF_SIZE) {
		puts("ERR\n");
		return;
	}

	char *errstr;

	ret = sqlite3_exec(db, query, NULL, NULL, &errstr);
	if (ret != SQLITE_OK) {
		fprintf(stderr, "SQL error: %s\n", errstr);
		sqlite3_free(errstr);
	}

}

int main(int argc, char *argv[])
{
	sqlite3 *db;
	int rc = sqlite3_open(DATABASE, &db);
	if (rc) {
		fprintf(stderr, "Can't open database: %s\n", sqlite3_errmsg(db));
		sqlite3_close(db);
		exit(-1);
	}

	if (argc == 1) {
		fprintf(stderr, "USAGE:\nurlfilter daemon\nurlfilter show\nurlfilter add <ip> <netmask> <url> <info>\n\n");
		exit(-1);

	} else if (argc == 2) {
		if (!strcmp(argv[1], "daemon"))
			daemon_loop(db);

		else if (!strcmp(argv[1], "show"))
			show_db(db);

	} else if ((argc = 6) && !strcmp(argv[1], "add")) {
		add_db(db, argv[2], argv[3], argv[4], argv[5]);
	}

	sqlite3_close(db);
	return 0;
}