/*
 * Copy me if you can.
 * by 20h
 *
 * Copyright (c) 2015 Natanael Copa <ncopa@alpinelinux.org>
 */

#include <err.h>
#include <errno.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/socket.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>

#include <linux/netlink.h>

#include <libkmod.h>
#include <blkid.h>

#include "arg.h"

#define EVENT_TIMEOUT 2000

int dodebug;
char *argv0;

#if defined(DEBUG)
#include <stdarg.h>
void
dbg(const char *fmt, ...)
{
	va_list fmtargs;

	fprintf(stderr, "%s: ", argv0);
	va_start(fmtargs, fmt);
	vfprintf(stderr, fmt, fmtargs);
	va_end(fmtargs);
	fprintf(stderr, "\n");
}
#else
#define dbg(...)
#endif

struct uevent {
	char *buf;
	size_t bufsize;
	char *message;
	char *subsystem;
	char *action;
	char *modalias;
	char *devname;
	char *major;
	char *minor;
};

struct ueventconf {
	char **program_argv;
	char *search_device;
	char *subsystem_filter;
	int modalias_count;
	int fork_count;
};


void
run_child(char **argv)
{
	pid_t pid;

	if (!(pid = fork())) {
		dbg("running %s", argv[0]);
		if (execv(argv[0], argv) < 0)
			err(1, argv[0]);
		exit(0);
	}
	if (pid < 0)
		err(1,"fork");

	waitpid(pid, NULL, 0);
}


int load_kmod(const char *modalias)
{
	static struct kmod_ctx *ctx = NULL;
	struct kmod_list *list = NULL;
	struct kmod_list *node;
	int r, count=0;

	dbg("loading alias '%s'", modalias);

	if (ctx == NULL) {
		dbg("initializing kmod");
		ctx = kmod_new(NULL, NULL);
		if (ctx == NULL)
			return -1;
		kmod_set_log_fn(ctx, NULL, NULL);
		r = kmod_load_resources(ctx);
		dbg("load kmod resources: %i", r);
	}

	r = kmod_module_new_from_lookup(ctx, modalias, &list);
	if (r < 0) {
		dbg("alias '%s' lookup failure", modalias);
		return r;
	}

	kmod_list_foreach(node, list) {
		struct kmod_module *mod = kmod_module_get_module(node);
		const char *fmt;
		r = kmod_module_probe_insert_module(mod,
						    KMOD_PROBE_APPLY_BLACKLIST,
						    NULL, NULL, NULL, NULL);
		if (r == 0) {
			fmt = "module '%s' inserted";
			count++;
		} else if (r == KMOD_PROBE_APPLY_BLACKLIST) {
			fmt = "module '%s' is blacklisted";
		} else {
			fmt = "module '%s' failed";
		}
		dbg(fmt, kmod_module_get_name(mod));
		kmod_module_unref(mod);
	}
	kmod_module_unref_list(list);
	return count;
}

void start_mdadm(char *devnode)
{
	char *mdadm_argv[] = {
		"/sbin/mdadm",
		"--incremental",
		"--export",
		devnode,
		NULL
	};
	run_child(mdadm_argv);
}

void start_lvm2(char *devnode)
{
	char *lvm2_argv[] = {
		"/sbin/lvm", "vgchange",
		"--activate" , "ay", "--noudevsync", "--sysinit",
		NULL
	};
	run_child(lvm2_argv);
}

int is_searchdev(char *devname, const char *searchdev)
{
	static blkid_cache cache = NULL;
	char *type = NULL, *label = NULL, *uuid = NULL;
	char devnode[256];

	if (searchdev == NULL)
		return 0;

	if (strcmp(devname, searchdev) == 0) {
		dbg("Found %s", searchdev);
		return 1;
	}

	if (cache == NULL)
		blkid_get_cache(&cache, NULL);

	snprintf(devnode, sizeof(devnode), "/dev/%s", devname);
	if (strncmp("LABEL=", searchdev, 6) == 0) {
		label = blkid_get_tag_value(cache, "LABEL", devnode);
		if (label && strcmp(label, searchdev+6) == 0) {
			dbg("Found %s (%s)", searchdev, devnode);
			free(label);
			return 1;
		}
	} else if (strncmp("UUID=", searchdev, 5) == 0) {
		uuid = blkid_get_tag_value(cache, "UUID", devnode);
		if (uuid && strcmp(uuid, searchdev+5) == 0) {
			printf("Found %s (%s)\n", searchdev, devnode);
			free(uuid);
			return 1;
		}
	}

	type = blkid_get_tag_value(cache, "TYPE", devnode);
	if (type || label || uuid) {
		printf("DEBUG:%s: (searching %s)\n"
			"\ttype='%s'\n"
			"\tlabel='%s'\n"
			"\tuuid='%s'\n", devnode, searchdev,
			type ? type : NULL,
			label ? label : NULL,
			uuid ? uuid : NULL);
	}

	if (type) {
		if (strcmp("linux_raid_member", type) == 0) {
			start_mdadm(devnode);
		} else if (strcmp("LVM2_member", type) == 0) {
			start_lvm2(devnode);
		} else {
			load_kmod(type);
		}
		free(type);
	}

	if (label)
		free(label);

	if (uuid)
		free(uuid);

	return 0;
}

int dispatch_uevent(struct uevent *ev, struct ueventconf *conf)
{
	if (conf->subsystem_filter && ev->subsystem
	    && strcmp(ev->subsystem, conf->subsystem_filter) != 0) {
		dbg("subsystem '%s' filtered out (by '%s').",
		    ev->subsystem, conf->subsystem_filter);
		return 0;
	}

	if (ev->action == NULL)
		return 0;

	if (ev->modalias != NULL && strcmp(ev->action, "add") == 0) {
		load_kmod(ev->modalias);
		conf->modalias_count++;

	} else if (ev->devname != NULL) {
		if (conf->program_argv[0] != NULL) {
			run_child(conf->program_argv);
			conf->fork_count++;
		}

		if (ev->subsystem && strcmp(ev->subsystem, "block") == 0
		    && strcmp(ev->action, "add") == 0
		    && is_searchdev(ev->devname, conf->search_device))
			return 1;
	}
	return 0;
}

int process_uevent(char *buf, const size_t len, struct ueventconf *conf)
{
	struct uevent ev;

	int i, slen = 0;
	char *key, *value;

	memset(&ev, 0, sizeof(ev));
	ev.buf = buf;
	ev.bufsize = len;
	clearenv();
	setenv("PATH", "/sbin:/bin", 1);

	for (i = 0; i < len; i += slen + 1) {

		key = buf + i;
		value = strchr(key, '=');
		slen = strlen(buf+i);

		if (i == 0 && slen != 0) {
			/* first line, the message */
			ev.message = key;
			continue;
		}

		if (!slen)
			continue;

		value[0] = '\0';
		value++;

		if (strcmp(key, "MODALIAS") == 0) {
			ev.modalias = value;
		} else if (strcmp(key, "ACTION") == 0) {
			ev.action = value;
		} else if (strcmp(key, "SUBSYSTEM") == 0) {
			ev.subsystem = value;
		} else if (strcmp(key, "DEVNAME") == 0) {
			ev.devname = value;
		} else if (strcmp(key, "MAJOR") == 0) {
			ev.major = value;
		} else if (strcmp(key, "MINOR") == 0) {
			ev.minor = value;
		}

		if (strcmp(key, "PATH")) {
			setenv(key, value, 1);
		//	dbg("%s = \"%s\"", key, value);
		}
	}
	return dispatch_uevent(&ev, conf);
}

void
usage(void)
{
	errx(1, "usage: %s [-dku] [-f subsystem] [-b device] PATH\n", argv0);
}

int main(int argc, char *argv[])
{
	struct pollfd fds;
	int r;
	struct ueventconf conf = {
		.subsystem_filter = NULL,
		.search_device = NULL,
		.modalias_count = 0,
		.fork_count = 0,
	};

	int event_count = 0;
	int showudev, showkernel;
	size_t total_bytes;

	showudev = 1;
	showkernel = 1;

	ARGBEGIN {
	case 'd':
		dodebug = 1;
		break;
	case 'k':
		showudev = 0;
		break;
	case 'u':
		showkernel = 0;
		break;
	case 'f':
		conf.subsystem_filter = EARGF(usage());
		break;
	case 's':
		conf.search_device = EARGF(usage());
		break;
	default:
		usage();
	} ARGEND;

	conf.program_argv = argv;

	chdir("/dev");
	fds.fd = 0; /* stdin */
	fds.events = POLLIN;
	while ((r = poll(&fds, 1, EVENT_TIMEOUT)) > 0) {
		size_t len;
		struct iovec iov;
		char cbuf[CMSG_SPACE(sizeof(struct ucred))];
		char buf[16384];
		struct cmsghdr *chdr;
		struct ucred *cred;
		struct msghdr hdr;
		struct sockaddr_nl cnls;

		if (!(fds.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.fd, &hdr, 0);
		if (len < 0) {
			if (errno == EINTR)
				continue;
			err(1, "recvmsg");
		}
		if (len < 32 || len >= sizeof(buf))
			continue;

		total_bytes += len;
		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;
		}

		if (process_uevent(buf, len, &conf)) {
			dbg("FOUND %s", conf.search_device);
			break;
		}
		event_count++;
		if (fds.revents & POLLHUP) {
			dbg("parent hung up\n");
			break;
		}
	}
	if (r == -1)
		err(1, "poll");

	if (r == 0)
		dbg("exit due to timeout");
	dbg("modaliases: %i, forks: %i, events: %i, total bufsize: %zu",
		conf.modalias_count,
		conf.fork_count,
		event_count, total_bytes);

	return 0;
}