#define _XOPEN_SOURCE
#include <errno.h>
#include <getopt.h>
#include <limits.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <pwd.h>
#include <unistd.h>

int chroot(char *path);

struct userbind {
	char *src, *dest;
	struct userbind *next;
};

pid_t child;
char *pidpath = NULL, *root = NULL;
struct userbind *binds = NULL;

static void check(int p, char *ctx) {
	if (p) {
		return;
	}
	fprintf(stderr, "%s: %s\n", ctx, strerror(errno));
	exit(1);
}

static void usage(void) {
	fprintf(stderr, "Usage: qemu-chroot [-b <src>:<dest>] [-p <pidfile>] "
			"/bin/qemu-system-<arch> [args...]\n");
}

static void cleanup() {
	int r;
	char path[PATH_MAX + 1] = "";
	char *paths[] = {
		"/dev/null",
		"/dev/kvm",
	};
	for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); ++i) {
		snprintf(path, sizeof(path), "%s%s", root, paths[i]);
		r = umount2(path, 0);
		if (r != 0) {
			fprintf(stderr, "umount %s: %s\n", path, strerror(errno));
		}
	}

	while (binds) {
		snprintf(path, sizeof(path), "%s%s", root, binds->dest);
		r = umount2(path, 0);
		if (r != 0) {
			fprintf(stderr, "umount %s: %s\n", path, strerror(errno));
		}
		struct userbind *next = binds->next;
		free(binds);
		binds = next;
	}

	r = umount2(root, 0);
	if (r != 0) {
		fprintf(stderr, "umount %s: %s\n", root, strerror(errno));
	}
	remove(root);

	remove(pidpath);
}

static void sig(int s) {
	kill(child, s);
}

int main(int argc, char *argv[]) {
	FILE *pidfile = NULL;

	char c;
	while ((c = getopt(argc, argv, "b:p:h")) != -1) {
		switch (c) {
		case 'b':;
			char *src = strtok(optarg, ":");
			char *dest = strtok(NULL, ":");
			struct userbind *bind = calloc(1, sizeof(struct userbind));
			bind->src = src;
			bind->dest = dest;
			bind->next = binds;
			binds = bind;
			break;
		case 'p':
			if (pidfile != NULL) {
				fprintf(stderr, "-p may only be specified once");
				return 1;
			}
			pidpath = optarg;
			pidfile = fopen(pidpath, "w");
			check(pidfile != NULL, "open pid file for write");
			break;
		case 'h':
			usage();
			return 0;
		default:
			usage();
			return 1;
		}
	}

	argc -= optind;
	argv += optind;

	char tmp[] = "/tmp/qemu-chroot.XXXXXX";
	root = mkdtemp(tmp);
	check(root != NULL, "mkdtemp");

	struct passwd *nobody = getpwnam("nobody");
	check(nobody != NULL, "getpwnam");

	int r;
	r = mount("/usr/lib/qemu-minimal-static", root, "", MS_BIND, NULL);
	check(r == 0, "bind mount /usr/lib/qemu-minimal-static");

	char *paths[] = {
		"/dev/null",
		"/dev/kvm",
	};
	char path[PATH_MAX + 1] = "";
	for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); ++i) {
		snprintf(path, sizeof(path), "%s%s", root, paths[i]);
		r = mount(paths[i], path, "", MS_BIND, NULL);
		if (r != 0) {
			fprintf(stderr, "bind mount %s: %s\n", paths[i], strerror(errno));
			return 1;
		}
	}

	struct userbind *b = binds;
	while (b) {
		snprintf(path, sizeof(path), "%s%s", root, b->dest);
		r = mount(b->src, path, "", MS_BIND, NULL);
		if (r != 0) {
			fprintf(stderr, "bind mount %s: %s\n", b->src, strerror(errno));
			return 1;
		}
		b = b->next;
	}

	struct sigaction sa = {0};
	sigemptyset(&sa.sa_mask);
	sa.sa_handler = sig;
	sa.sa_flags = SA_RESTART;

	r = sigaction(SIGINT, &sa, NULL);
	check(r != -1, "sigaction");

	r = sigaction(SIGTERM, &sa, NULL);
	check(r != -1, "sigaction");

	child = fork();
	check(child != -1, "fork");
	if (child == 0) {
		if (pidfile != NULL) {
			fclose(pidfile);
		}

		r = chroot(root);
		check(r == 0, "chroot");

		r = setgid(nobody->pw_gid);
		check(r == 0, "setgid");
		r = setegid(nobody->pw_gid);
		check(r == 0, "setegid");
		r = setuid(nobody->pw_uid);
		check(r == 0, "setuid");
		r = seteuid(nobody->pw_uid);
		check(r == 0, "seteuid");

		if (setuid(0) != -1) {
			fprintf(stderr, "Unable to drop root, bailing out\n");
			return 1;
		}

		r = execv(argv[0], &argv[0]);
		check(r == 0, "execv");
	}

	if (pidfile != NULL) {
		fprintf(pidfile, "%d\n", child);
		fclose(pidfile);
	}

	int status;
	do {
		pid_t wpid = waitpid(child, &status, WUNTRACED | WCONTINUED);
		check(wpid != -1, "waitpid");
	} while (!WIFEXITED(status) && !WIFSIGNALED(status));

	cleanup();

	return WEXITSTATUS(status);
}