// https://www.kernel.org/doc/Documentation/connector/connector.txt
// TODO CLI options to toggle sniffer, proc_connector, fanotify, ...
// TODO map uids and gids to real names
  // implement getgrgid and getpwuid myself because cant link these static
// TODO permissions, attributes, owner, groupships of files
// TODO limits? /proc/X/limits
// TODO atexit() handler to log when this dies
// TODO emit a json log startup event
// TODO emit heartbeat json logs
// TODO multiple interfaces for sniffing

#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdarg.h>
#include <stdbool.h>
#include <string.h>
#include <limits.h>
#include <syslog.h>
#include <signal.h>
#include <netdb.h>
#include <fcntl.h>

#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <sys/stat.h>

#include "net.h"
#include "error.h"
#include "output.h"
#include "sniffer.h"
#include "fanotify.h"
#include "proc_ledger.h"
#include "hash_ledger.h"
#include "proc_connector.h"
#include "time_common.h"
#include "agent_context.h"
#include "av_rules.h"
#include "aho-corasick.h"

/* dirty hack to make this work on ancient systems */
#ifndef PROC_EVENT_PTRACE
#define PROC_EVENT_PTRACE 0x00000100
#endif

#ifndef PROC_EVENT_COMM
#define PROC_EVENT_COMM 0x00000200
#endif

#ifndef PROC_EVENT_COREDUMP
#define PROC_EVENT_COREDUMP 0x40000000
#endif

/*
 * Globals
 */
sock_t                        sock;
bool                          quiet           = false;
bool                          use_syslog      = true;
bool                          log_to_file     = false;
bool                          remote_logging  = true;
char                         *log_server      = "127.0.0.1";
port_t                        log_server_port = 55555;
char                         *outfile         = "/var/log/noawareness.json.log";
FILE                         *outfilep;

bool                          promisc         = false;

unsigned long                 maxsize         = 1024 * 1024 * 100; // 100Mb
unsigned long                 boot_time;
long                          ticks;

bool                          daemonize       = false;
char                         *pidfile         = "/var/run/noawareness.pid";
static volatile sig_atomic_t  reload_config   = false;

char                         *rules_file      = "rules";

// Write pid file to disk
void write_pid_file(const char *path, pid_t pid) {
	FILE        *pidfile;

	pidfile = fopen(path, "w");
	if (pidfile == NULL) {
		error_fatal("Unable to open PID file %s: %s", path, strerror(errno));
	}

	fprintf(pidfile, "%d", pid);
	fclose(pidfile);
}

FILE *open_log_file(const char *outfile) {
	FILE *fp;

	fp = fopen(outfile, "a+");
	if (fp == NULL) {
		error_fatal("Unable to open log file: %s", strerror(errno));
	}

	return fp;
}

// Sets up a UDP socket to a remote host for sending log events
// TODO strategies to do TCP, encryption, etc. the remote logging
//      setup is pretty dirty
static void setup_remote_logging() {
	struct addrinfo hints, *res, *rp;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;

	char pstr[6];
	snprintf(pstr, sizeof(pstr), "%d", log_server_port);

	if (getaddrinfo(log_server, pstr, &hints, &res) != 0) {
		error_fatal("getaddrinfo(): %s", strerror(errno));
	}

	for (rp = res; rp != NULL; rp = rp->ai_next) {
		sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
		if (sock == -1) {
			continue;
		}

		if (connect(sock, rp->ai_addr, rp->ai_addrlen) != -1) {
			break;
		}

		close(sock);
	}

	if (rp == NULL) {
		error_fatal("Unable to connect to %s: %s\n", log_server, strerror(errno));
	}

	char addr_s[INET6_ADDRSTRLEN];
	void *addr_ptr = NULL;

	if (rp->ai_family == AF_INET) {
		struct sockaddr_in *ipv4 = (struct sockaddr_in *)rp->ai_addr;
		addr_ptr = &(ipv4->sin_addr);
	} else if (rp->ai_family == AF_INET6) {
		struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)rp->ai_addr;
		addr_ptr = &(ipv6->sin6_addr);
	}

	if (addr_ptr &&
		inet_ntop(rp->ai_family, addr_ptr, addr_s, sizeof(addr_s))) {
		msg("Connected to logserver: %s:%d", addr_s, log_server_port);
	}

	freeaddrinfo(res);
}

// SIGHUP handler
static void handle_sighup(int sig, siginfo_t *siginfo, void *context) {
	reload_config = true;
}

// install SIGHUP handler
static void install_sighup_handler() {
	struct sigaction act = {0};

	act.sa_sigaction = handle_sighup;
	act.sa_flags = 0;

	if (sigaction(SIGHUP, &act, NULL) < 0) {
		error_fatal("sigaction(): %s", strerror(errno));
	}
}

// Variadic function that returns the maximum value input
static int max_fd(int count, ...) {
	va_list args;

	va_start(args, count);

	int max = va_arg(args, int);

	for (int i = 1; i < count; i++) {
		int val = va_arg(args, int);
		if (val > max) {
			max = val;
		}
	}

	va_end(args);

	return max;
}

ac_context_t *build_matcher_from_rules(rule_set_t *rules) {
	ac_context_t *ctx = ac_new();
	if (!ctx) {
		return NULL;
	}

	for (size_t i = 0; i < rules->patterns.count; i++) {
		pattern_t *p = &rules->patterns.patterns[i];
		if (ac_add_pattern(ctx, p->id, p->bytes, p->len) != 0) {
			error("Failed to add pattern %s to Aho-Corasick", p->id);
			ac_free(ctx);
			return NULL;
		}
	}

	if (ac_build(ctx) != 0) {
		error("Failed to build Aho-Corasick automaton");
		ac_free(ctx);
		return NULL;
	}

	return ctx;
}

// Display help/usage menu
static void usage(const char *progname) {
	fprintf(stderr, "usage: %s [-h?]\n\n", progname);
	fprintf(stderr, "    -h/-?      - Print this menu and exit.\n");
	fprintf(stderr, "    -d         - Daemonize. Default: %s\n",
			daemonize ? "yes" : "no");
	fprintf(stderr, "    -i <iface> - Interface to sniff\n");
	fprintf(stderr, "    -x         - Toggle promiscuous mode. Default: %s\n",
			promisc ? "true" : "false");
	fprintf(stderr, "    -m <bytes> - Max size of file to hash. Default: %ld\n",
			maxsize);
	fprintf(stderr, "    -o <file>  - Outfile for JSON output, Default: %s\n",
			outfile);
	fprintf(stderr, "    -O         - Toggle local JSON logging. Default: %s\n",
			log_to_file ? "true" : "false");
	fprintf(stderr, "    -P <path>  - Path to PID file. Default: %s\n", pidfile);
	fprintf(stderr, "    -r         - Toggle remote logging. Default: %s\n",
			remote_logging ? "true" : "false");
	fprintf(stderr, "    -R <path>  - Path to AV rule definitions. Default: %s\n", rules_file);
	fprintf(stderr, "    -s <IP>    - Remote log server. Default: %s\n",
			log_server);
	fprintf(stderr, "    -S         - Toggle syslog. Default: %s\n",
			use_syslog ? "true" : "false");
	fprintf(stderr, "    -p <port>  - Port of remote server. Default: %d\n",
			log_server_port);
	fprintf(stderr, "    -q         - Toggle quiet mode. Default: %s\n",
			quiet ? "true" : "false");

  exit(EXIT_FAILURE);
}

int main(int argc, char *argv[]) {
	int               opt;
	pid_t             pid;
	sock_t            proc_connector;
	sock_t            sniffer; // socket for sniffing to create network events
	char             *interface = NULL;
	int               fan_fd;  // fanotify fd
	fd_set            fdset;
	agent_context_t  *agent_ctx;

	/* Parse CLI options */
	// TODO flag to enable/disable fanotify blocking
	while((opt = getopt(argc, argv, "qrRdm:s:So:Op:rP:xi:h?")) != -1) {
		switch (opt) {
		case 'd': /* Daemonize */
			daemonize = daemonize ? false : true;
			break;

		case 'i': /* Interface for SOCK_RAW sniffing */
			// TODO if optarg == "list", list interfaces
			interface = optarg;
			break;

		case 'm': /* Maximum filesize to hash */
			maxsize = atol(optarg);
			break;

		case 'o': /* Path to outfile */
			// TODO check if writeable
			outfile = optarg;
			log_to_file = true;
			break;

		case 'O': /* Toggle logging to a file */
			log_to_file = log_to_file ? false : true;
			break;

		case 'p': /* Remote server port */
			log_server_port = atoi(optarg);
			break;

		case 'P': /* PID file location */
			pidfile = optarg;
			break;

		case 'q': /* Toggle quiet mode */
			quiet = quiet ? false : true;
			break;

		case 'r': /* Toggle remote logging */
			remote_logging = remote_logging ? false : true;
			break;

		case 'R': /* AV rules file */
			rules_file = optarg;
			break;

		case 's': /* Remote server */
			log_server = optarg;
			if (!validate_ip(log_server)) {
				error_fatal("Invalid IP address: %s", log_server);
			}
			break;

		case 'S': /* Toggle syslog */
			use_syslog = use_syslog ? false : true;
			break;

		case 'x': /* Toggle promiscuous mode */
			promisc = promisc ? false : true;
			break;

		/* All of these effectively call usage(), so roll them over */
		case 'h':
		case '?':
		default:
			usage(argv[0]);
		}
	}

	/* Set up syslog() */
	if (use_syslog) {
		openlog("noawareness", LOG_PID, LOG_USER);
	}

	/* SIGHUP handler. */
	install_sighup_handler();

	/* Open log file. */
	if (log_to_file) {
		outfilep = open_log_file(outfile);
	}

	/* Allocate agent_ctx */
	agent_ctx = calloc(1, sizeof(*agent_ctx));
	if (!agent_ctx) {
		error_fatal("calloc: %s", strerror(errno));
	}
	/* Get our hostname for reporting purposes. */
	if (gethostname(agent_ctx->hostname, sizeof(agent_ctx->hostname)) == -1) {
		error_fatal("gethostname(): %s", strerror(errno));
	}

	/* Get system boot time */
	ticks = sysconf(_SC_CLK_TCK);
	boot_time = get_boot_time();

	/* Set up remote logging */
	if (remote_logging) {
		setup_remote_logging();
	}

	/* Create hash ledger */
	agent_ctx->hash_ledger = hash_ledger_init(4096);
	if (!agent_ctx->hash_ledger) {
		fprintf(stderr, "FATAL: Error creating hash ledger: %s\n", strerror(errno));
		return EXIT_FAILURE;
	}

	/* Parse AV rules and create Aho-Corasick context */
	//agent_ctx->rules = {0};
	if (load_rules(rules_file, &agent_ctx->rules) != 0) {
		error_fatal("FATAL: Failed to load AV signatures from file %s: %s", rules_file, strerror(errno));
	}

	agent_ctx->ac = build_matcher_from_rules(&agent_ctx->rules);
	if (!agent_ctx->ac) {
		error_fatal("FATAL: Failed to create Aho-Corasick context");
	}


	/* Create process ledger--this needs to be done before the
	   fanotify handler, otherwise processes will be blocked until the
	   process ledger is hydrated */
	agent_ctx->proc_ledger = proc_ledger_init(4096);
	proc_ledger_hydrate(agent_ctx);

	/* Create proc_connector socket. */
	proc_connector = setup_proc_connector();
	if (proc_connector == -1) {
		error_fatal("error setting up proc_connector: %s", strerror(errno));
	}

	/* Create fanotify descriptor */
	fan_fd = setup_fanotify();
	if (fan_fd == -1) {
		error("Error setting up fanotify: %s", strerror(errno));
	}

	/* Create sniffer raw socket */
	// TODO flag to toggle sniffer
	if (!interface) {
		interface = get_default_iface();
		if (interface == NULL) {
			error_fatal("FATAL: Unable to determine network interface. Please specify with -i");
		}
	}

	sniffer = sniffer_init_interface(interface, promisc);
	if (sniffer == -1) {
		error("Error setting up SOCK_RAW socket: %s", strerror(errno));
	}

	/* Daemonize the process if desired. */
	if (daemonize) {
		pid = fork();
		if (pid < 0) {
			error_fatal("FATAL: fork(): %s", strerror(errno));
		}

		else if (pid > 0) {
			write_pid_file(pidfile, pid);
			exit(EXIT_SUCCESS);
		}
	}

	/* Set up select loop. */
	int setsize = max_fd(3, proc_connector, fan_fd, sniffer) + 1;

	for(;;) {
		FD_ZERO(&fdset);
		if (proc_connector != -1) FD_SET(proc_connector, &fdset);
		if (fan_fd != -1)         FD_SET(fan_fd, &fdset);
		if (sniffer != -1)        FD_SET(sniffer, &fdset);

		struct timeval select_timeout;
		select_timeout.tv_sec = 1;
		select_timeout.tv_usec = 0;

		int activity = select(setsize, &fdset, NULL, NULL, &select_timeout);
		if (activity < 0) {
			if (errno != EINTR) {
				error_fatal("select(): %s", strerror(errno));
			}
		}

		if (activity == 0) {
			//printf("no activity...\n");
			continue;
		}

		if (FD_ISSET(proc_connector, &fdset)) {
			select_proc_connector(proc_connector, agent_ctx);
		}

		if (FD_ISSET(fan_fd, &fdset)) {
			select_fanotify(fan_fd, agent_ctx);
		}

		if (FD_ISSET(sniffer, &fdset)) {
			sniffer_handle_packet(sniffer, agent_ctx);
		}

		if (reload_config) { // TODO break this into a reload function
			reload_config = false;
			msg("Caught SIGHUP.");

			if (log_to_file) {
				msg("Reloading JSON log file");
				fclose(outfilep);
				outfilep = open_log_file(outfile);
			}

			msg("Flushing hash ledger");
			hash_ledger_destroy(agent_ctx->hash_ledger);
			agent_ctx->hash_ledger = hash_ledger_init(4096);
			if (!agent_ctx->hash_ledger) {
				error_fatal("FATAL: Error creating hash ledger: %s\n", strerror(errno));
			}

			msg("Reloading AV signatures");
			pattern_table_free(&agent_ctx->rules.patterns);
			free_rules(&agent_ctx->rules);
			ac_free(agent_ctx->ac);
			if (load_rules(rules_file, &agent_ctx->rules) != 0) {
				error_fatal("FATAL: Failed to load AV signatures from file %s: %s", rules_file, strerror(errno));
			}
			agent_ctx->ac = build_matcher_from_rules(&agent_ctx->rules);
			if (!agent_ctx->ac) {
				error_fatal("FATAL: Failed to create Aho-Corasick context");
			}

			msg("Rehydrating process list");
			proc_ledger_destroy(agent_ctx->proc_ledger);
			agent_ctx->proc_ledger = proc_ledger_init(4096);
			proc_ledger_hydrate(agent_ctx);
		}
	}

	/* Shouldn't ever get here */
	proc_ledger_destroy(agent_ctx->proc_ledger);
	hash_ledger_destroy(agent_ctx->hash_ledger);
	ac_free(agent_ctx->ac);
	pattern_table_free(&agent_ctx->rules.patterns);
	free_rules(&agent_ctx->rules);
	close(proc_connector);
	close(sniffer);
	close(sock);
	closelog();

	return EXIT_SUCCESS;
}