/*
 * Copyright (c) 2020 Markus Hennecke <markus-hennecke@markus-hennecke.de>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/queue.h>
#include <sys/uio.h>
#include <sys/event.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/fcntl.h>
#include <sys/param.h>
#include <sys/wait.h>
#include <arpa/inet.h>
#include <net/if.h>
#include <net/pfvar.h>
#include <err.h>
#include <errno.h>
#include <poll.h>
#include <regex.h>
#include <signal.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <syslog.h>
#include <unistd.h>
#include <imsg.h>

#include "log.h"

int	 run = 1;
char	*blacklist_table = "blacklist";

regex_t rx_failed_invalid_user;

/*
 * Simplified regex match for IPv6 addresses, the code in send_add_address()
 * will check for a valid IPv4/IPv6 address anyway so keep the regex
 * simple.
 */
#define RX_FAILED_INVALID_USER	"Failed password for " \
    "(invalid user [^ ]+|root) from "				\
    "("								\
    "([0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3})"	\
    "|"								\
    "([0-9A-Fa-f][0-9A-Fa-f:]+)"				\
    ")"								\
    " port [0-9]+ ssh"

#define AUTHLOG_FILENAME	"/var/log/authlog"
#define MAXLINELEN		1024

enum imsg_type {
	IMSG_ADD_ADDRESS,
};

typedef struct __attribute__((packed)) {
	union {
		struct in_addr	addr4;
		struct in6_addr	addr6;
	} addr;
	sa_family_t	af;
} add_address_msg_t;

static int	 add_address(int, add_address_msg_t *, const char *);
static char	*rx_error_string(int, regex_t *);
static void	 sighandler(int);
static void	 usage(void);
static void	 handle_authlog(struct imsgbuf *, FILE *);
static void	 dispatch_imsg(int dev, struct imsgbuf *);
static int	 pfdev_main(struct imsgbuf *);
int		 send_add_address(struct imsgbuf *, const char *);


__dead void
usage(void)
{
	fprintf(stderr, "usage: %s [-dv]\n", getprogname());
	exit(1);
}


char *
rx_error_string(int rc, regex_t *rx)
{
	static char error[1024];
	regerror(rc, rx, error, sizeof(error));
	return error;
}


void
sighandler(int sig)
{
	switch (sig) {
	case SIGTERM:
	case SIGINT:
		run = 0;
		break;
	case SIGCHLD:
		wait(NULL);
		run = 0;
		break;
	default:
		break;
	}
}


void
handle_authlog(struct imsgbuf *ibuf, FILE *authlog)
{
	char line[MAXLINELEN];

	while ((fgets(line, MAXLINELEN, authlog)) != NULL) {
		regmatch_t match[3];
		int rc = regexec(&rx_failed_invalid_user, line, 3, match, 0);
		switch (rc) {
		case REG_NOMATCH:
			log_debug("nomatch, line: %s", line);
			break;
		case 0: {
				const regmatch_t *m = &match[2];
				const size_t len = m->rm_eo - m->rm_so;
				char *ip = strndup(line + m->rm_so, len);
				send_add_address(ibuf, ip);
				free(ip);
			}
			break;
		default:
			log_warn("regexec: %s",
			    rx_error_string(rc, &rx_failed_invalid_user));
			break;
		}
	}
	if (feof(authlog))
		clearerr(authlog);
	else if (ferror(authlog))
		fatal(NULL);
}


int
send_add_address(struct imsgbuf *ibuf, const char *ip)
{
	add_address_msg_t	msg;
	int			n, rc;

	bzero(&msg, sizeof(msg));
	if ((rc = inet_pton(AF_INET, ip, &msg.addr)) < 1) {
		if (-1 == rc) {
			log_warn("inet_pton");
			return -1;
		}
		if ((rc = inet_pton(AF_INET6, ip, &msg.addr)) == -1) {
			log_warn("inet_pton");
			return -1;
		}
		if (0 == rc)
			return 0;
		msg.af = AF_INET6;
	} else {
		msg.af = AF_INET;
	}
	imsg_compose(ibuf, IMSG_ADD_ADDRESS, 0, 0, -1, &msg, sizeof(msg));

send:
	if ((n = msgbuf_write(&ibuf->w)) == -1 && errno != EAGAIN)
		fatal("msgbuf_write");
	if (-1 == n && errno == EAGAIN) {
		log_warnx("msgbuf_write, EAGAIN");
		goto send;
	}
	if (n == 0)
		fatal("msgbuf_write, connection closed");
	return 1;
}


int
add_address(int dev, add_address_msg_t *msg, const char *tablename)
{
	struct pfr_addr		addr;
	struct pfioc_table	io;
	struct pfr_table	table;

	bzero(&addr, sizeof(addr));
	memcpy(&addr.pfra_u, &msg->addr, sizeof(addr.pfra_u));

	addr.pfra_af = msg->af;
	addr.pfra_net = (msg->af == AF_INET) ? 32 : 128;

	bzero(&table, sizeof(struct pfr_table));
	strlcpy(table.pfrt_name, tablename, PF_TABLE_NAME_SIZE);

	bzero(&io, sizeof(struct pfioc_table));
	io.pfrio_flags = PFR_FLAG_FEEDBACK;
	io.pfrio_table = table;
	io.pfrio_buffer = &addr;
	io.pfrio_esize = sizeof(struct pfr_addr);
	io.pfrio_size = 1;
	if (ioctl(dev, DIOCRADDADDRS, &io) == -1) {
		log_warn("DIOCRADDADDRS");
		return -1;
	}

	return io.pfrio_nadd;
}


int
pfdev_main(struct imsgbuf *ibuf)
{
	int			nready;
	int			dev;
	struct pollfd		pfd[1];

	dev = open("/dev/pf", O_RDWR);
	if (-1 == dev)
		fatal("open /dev/pf failed");

	if (unveil(NULL, NULL) == -1)
		fatal("unveil");

	pfd[0].fd = ibuf->fd;
	pfd[0].events = POLLIN;

	while (run && ((nready = poll(pfd, 1, INFTIM)) != -1)) {
		if (nready == 0)
			fatal("time out");
		if ((pfd[0].revents & (POLLERR|POLLNVAL)))
			fatal("bad fd %d", pfd[0].fd);
		if ((pfd[0].revents & (POLLIN|POLLHUP)))
			dispatch_imsg(dev, ibuf);
	}
	close(dev);
	return 0;
}


void
dispatch_imsg(int dev, struct imsgbuf *ibuf)
{
	struct imsg		imsg;
	ssize_t			n, datalen;
	int			numadd;
	add_address_msg_t	msg;

again:
	if (((n = imsg_read(ibuf)) == -1) && errno != EAGAIN)
		fatal("imsg_read");
	if (n == 0 && errno == EAGAIN)
		goto again;
	if (n == 0)
		return;

	while (run) {
		if ((n = imsg_get(ibuf, &imsg)) == -1)
			fatal("imsg_get");
		if (n == 0)
			break;
		datalen = imsg.hdr.len - IMSG_HEADER_SIZE;

		switch (imsg.hdr.type) {
		case IMSG_ADD_ADDRESS:
			if (datalen < (ssize_t )sizeof(msg)) {
				fatal("corrupt imsg");
			}
			memcpy(&msg, imsg.data, sizeof(msg));
			numadd = add_address(dev, &msg, blacklist_table);
			if (numadd > 0) {
				char p[INET6_ADDRSTRLEN];
				inet_ntop(msg.af, &msg.addr, p, sizeof(p));
				log_info("Added %s to <%s>", p,
				    blacklist_table);
			}
			break;
		default:
			log_warn("Unknown imsg received");
			break;
		}
		imsg_free(&imsg);
	}
}


int
main(int argc, char **argv)
{
	int		 queue, nev;
	int		 ch, rc;
	struct kevent	 ev[2];
	int		 debug = 0, verbose = 0;
	FILE		*authlog = NULL;
	struct imsgbuf	 parent_ibuf, child_ibuf;
	int		 imsg_fds[2];
	int		 logdest = LOG_TO_STDERR;

	if (geteuid() != 0)
		errx(1, "Need root privileges to run.");

	while ((ch = getopt(argc, argv, "dv")) != -1) {
		switch (ch) {
		case 'd':
			debug = 2;
			break;
		case 'v':
			verbose++;
			break;
		default:
			usage();
			/* NOT REACHED */
		}
	}

	argc -= optind;
	argv += optind;

	if (argc != 0)
		usage();

	setprogname("ssh-sentry");
	if (!debug)
		logdest |= LOG_TO_SYSLOG;

	log_init(logdest, verbose, LOG_DAEMON);

	rc = regcomp(&rx_failed_invalid_user, RX_FAILED_INVALID_USER,
	    REG_EXTENDED);
	if (rc != 0)
		errx(1, "%s", rx_error_string(rc, &rx_failed_invalid_user));

	logdest = debug ? LOG_TO_STDERR : LOG_TO_SYSLOG;
	log_init(logdest, verbose, LOG_DAEMON);
	log_setverbose(verbose);

	if (!debug) {
		if (daemon(0, 0) == -1)
			fatal("daemon");
	}

	if ((queue = kqueue()) == -1)
		fatal("kqueue");

	if (socketpair(AF_UNIX, SOCK_STREAM, PF_UNSPEC, imsg_fds) == -1)
		fatal("socketpair");

	switch (fork()) {
	case -1:
		fatal("fork");
	case 0:
		/* child */
		signal(SIGTERM, sighandler);
		signal(SIGINT, sighandler);

		close(imsg_fds[0]);
		imsg_init(&child_ibuf, imsg_fds[1]);
		exit(pfdev_main(&child_ibuf));
	}
	close(imsg_fds[1]);
	imsg_init(&parent_ibuf, imsg_fds[0]);

	if (unveil(AUTHLOG_FILENAME, "r") == -1)
		fatal("unveil");
	if (unveil(NULL, NULL) == -1)
		fatal("unveil");

	signal(SIGTERM, sighandler);
	signal(SIGINT, sighandler);
	signal(SIGCHLD, sighandler);

	if (pledge("stdio rpath", "") == -1)
		fatal("pledge");

	if ((authlog = fopen(AUTHLOG_FILENAME, "r")) == NULL)
		fatal("%s", AUTHLOG_FILENAME);
	if (fseek(authlog, 0L, SEEK_END) < 0)
		fatal("fseek");

reopen:
	while (NULL == authlog) {
		sleep(1);
		authlog = fopen(AUTHLOG_FILENAME, "r");
	}

add_events:
	EV_SET(&ev[0], fileno(authlog), EVFILT_READ,
	    EV_ENABLE | EV_ADD | EV_CLEAR, 0, 0, NULL);
	EV_SET(&ev[1], fileno(authlog), EVFILT_VNODE,
	    EV_ENABLE | EV_ADD | EV_CLEAR,
	    NOTE_DELETE | NOTE_RENAME, 0, NULL);
	if (kevent(queue, ev, 2, NULL, 0, NULL))
		fatal(NULL);

	log_info("Starting event loop...");
	struct kevent ke;
	while (run) {
		nev = kevent(queue, NULL, 0, &ke, 1, NULL);
		if (nev == -1) {
			if (errno != EINTR) {
				log_warn("kevent");
				break;
			}
		}

		if (nev > 0) {
			if (ke.filter == EVFILT_READ) {
				handle_authlog(&parent_ibuf, authlog);
			} else if (ke.filter == EVFILT_VNODE) {
				if (ke.fflags & NOTE_RENAME) {
					log_info("%s renamed...",
					    AUTHLOG_FILENAME);
					authlog = freopen(AUTHLOG_FILENAME,
					    "r", authlog);
					if (authlog == NULL)
						goto reopen;
					goto add_events;
				} else if (ke.fflags & NOTE_DELETE) {
					log_info("%s deleted...",
					    AUTHLOG_FILENAME);
					authlog = freopen(AUTHLOG_FILENAME,
					    "r", authlog);
					if (authlog == NULL)
						goto reopen;
					goto add_events;
				}
			}
		}
	}
	close(queue);
	regfree(&rx_failed_invalid_user);

	return 0;
}