/*
 * Copy me if you can.
 * by 20h
 */

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <poll.h>
#include <ctype.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <libgen.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <linux/types.h>
#include <linux/netlink.h>

#include "arg.h"
#include "log.h"

int listfd = -1;
int dofork = 0;

struct handler {
	int fd;
	int argc;
	char **argv;
};

void
spawn_handler(struct handler *child)
{
	int pipefd[2];

	if (pipe(pipefd) == -1)
		edie("pipe");

	pid_t pid = fork();

	if (pid < 0)
		edie("fork");

	if (pid == 0) {
		close(pipefd[1]);
		dup2(pipefd[0], 0);
		execv(child->argv[0], child->argv);
		edie("execl");
	}
	close(pipefd[0]);
	child->fd = pipefd[1];
}

int
write_to_handler(struct handler *child, const char *buf, size_t count)
{
	if (child->fd == -1)
		spawn_handler(child);
	return write(child->fd, buf, count);
}

void
sighandler(int sig)
{
	switch(sig) {
	case SIGHUP:
	case SIGINT:
	case SIGQUIT:
	case SIGABRT:
	case SIGTERM:
		exit(0);
		break;
	default:
		break;
	}
}

void
initsignals(void)
{
	signal(SIGHUP, sighandler);
	signal(SIGINT, sighandler);
	signal(SIGQUIT, sighandler);
	signal(SIGABRT, sighandler);
	signal(SIGTERM, sighandler);

	signal(SIGCHLD, SIG_IGN);
	signal(SIGPIPE, SIG_IGN);
}

int
init_netlink_socket(void)
{
	struct sockaddr_nl nls;
	int fd, slen;

	memset(&nls, 0, sizeof(nls));
	nls.nl_family = AF_NETLINK;
	nls.nl_pid = getpid();
	nls.nl_groups = -1;

	fd = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_KOBJECT_UEVENT);
	if (fd < 0)
		edie("socket");

	slen = 64*1024;
	if (setsockopt(fd, SOL_SOCKET, SO_RCVBUFFORCE, &slen,
				sizeof(slen)) < 0) {
		edie("setsockopt");
	}
	slen = 1;
	if (setsockopt(fd, SOL_SOCKET, SO_PASSCRED, &slen,
				sizeof(slen)) < 0) {
		edie("setsockopt");
	}

	if (bind(fd, (void *)&nls, sizeof(nls)))
		edie("bind");

	fcntl(fd, F_SETFD, FD_CLOEXEC);
	return fd;
}

void
usage(void)
{
	die("usage: %s [-hdb] [-ku] [-- run [...]]\n", argv0);
}

int
main(int argc, char *argv[])
{
	struct sockaddr_nl cnls;
	struct pollfd fds[2];
	struct msghdr hdr;
	struct iovec iov;
	char buf[4097], cbuf[CMSG_SPACE(sizeof(struct ucred))];
	struct cmsghdr *chdr;
	struct ucred *cred;
	int len, showudev, showkernel;

	struct handler child;

	showkernel = 1;
	showudev = 1;

	ARGBEGIN {
	case 'b':
		dofork = 1;
		break;
	case 'd':
		dodebug = 1;
		break;
	case 'k':
		showudev = 0;
		break;
	case 'u':
		showkernel = 0;
		break;
	default:
		usage();
	} ARGEND;

	child.fd = -1;
	child.argc = argc;
	child.argv = argv;

	fds[0].events = POLLIN;
	fds[0].fd = init_netlink_socket();

	fds[1].fd = child.fd;
	fds[1].events = POLLERR;

	listfd = fds[0].fd;

	if (dofork) {
		if (daemon(0, 0) < 0)
			edie("daemon");
		umask(022);
	}

	initsignals();

	while (poll(fds, 2, -1) > -1) {
		if (fds[1].revents & POLLERR) {
			dbg("handler exited");
			close(child.fd);
			fds[1].fd = child.fd = -1;
		}

		if (!(fds[0].revents & POLLIN))
			continue;

		iov.iov_base = &buf;
		iov.iov_len = sizeof(buf);
		memset(&hdr, 0, sizeof(hdr));
		hdr.msg_iov = &iov;
		hdr.msg_iovlen = 1;
		hdr.msg_control = cbuf;
		hdr.msg_controllen = sizeof(cbuf);
		hdr.msg_name = &cnls;
		hdr.msg_namelen = sizeof(cnls);

		len = recvmsg(fds[0].fd, &hdr, 0);
		if (len < 0) {
			if (errno == EINTR)
				continue;
			edie("recvmsg");
		}
		if (len < 32 || len >= sizeof(buf))
			continue;

		chdr = CMSG_FIRSTHDR(&hdr);
		if (chdr == NULL || chdr->cmsg_type != SCM_CREDENTIALS)
			continue;

		/*
		 * Don't allow anyone but root to send us messages.
		 *
		 * We will allow users to send us messages, when
		 * udev is enabled. Udev is just a toy you should
		 * only use for testing.
		 */
		cred = (struct ucred *)CMSG_DATA(chdr);
		if (cred->uid != 0 && !showudev)
			continue;

		if (!memcmp(buf, "libudev", 8)) {
			/*
			 * Receiving messages from udev is insecure.
			 */
			if (!showudev)
				continue;
		} else {
			if (!showkernel)
				continue;
			/*
			 * Kernel messages shouldn't come from the
			 * userspace.
			 */
			if (cnls.nl_pid > 0)
				continue;
		}
		write_to_handler(&child, buf, len);
		fds[1].fd = child.fd;
	}

	return 0;
}