#include <sys/types.h>
#include <sys/jail.h>
#include <sys/wait.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>

void setup(); 
void teardown();

uid_t unpriv_uid = 5000;
gid_t unpriv_gid = 5000;

void
test(int (*fn)(), int expected, char *msg, char *msg2)
{
	int retval;

	setup();
	retval = fn();
	teardown();

	printf("%s (%s): ", msg, msg2);

	if (retval == expected) {
		printf("OK\n");
	} else {
		printf("FAILED (was: %d, expected: %d)\n", retval, expected);
	}
	fflush(stdout);
}

void 
test_as_root(int (*fn)(), int expected, char *msg)
{
	if (getuid() != 0) {
		fprintf(stderr, "must be run as root\n");
		exit(-1);
	}

	test(fn, expected, msg, "as root");
}

void 
test_as_jailed_root(int (*fn)(), int expected, char *msg)
{
	if (getuid() != 0) {
		fprintf(stderr, "must be run as root\n");
		exit(-1);
	}

	int child = fork();

	if (child == -1) {
		fprintf(stderr, "fork failed\n");
		exit(-2);
	}

	if (child) {
		struct jail j;
		j.version = 1;
		j.path = "/";
		j.hostname = "jail";
		j.n_ips = 0;

		int jid = jail(&j);
		if (jid < 0) {
			fprintf(stderr, "jail failed\n");
			exit(-1); // TODO
		}
		test(fn, expected, msg, "as jailed root");
		exit(0);
	}
	else {
		waitpid(child, NULL, 0);
	}
}

void 
test_as_unpriv(int (*fn)(), int expected, char *msg)
{
	if (getuid() != 0) {
		fprintf(stderr, "must be run as root\n");
		exit(-1);
	}

	int child = fork();

	if (child == -1) {
		fprintf(stderr, "fork failed\n");
		exit(-2);
	}

	if (child) {
		setgid(unpriv_gid);
		setuid(unpriv_uid);

		if (getuid() != unpriv_uid || getgid() != unpriv_gid) {
			fprintf(stderr, "setuid/gid failed\n");
			exit(-1); // TODO
		}
		test(fn, expected, msg, "as unpriv");
		exit(0);
	}
	else {
		waitpid(child, NULL, 0);
	}
}

void 
test_as_jailed_unpriv(int (*fn)(), int expected, char *msg)
{
	if (getuid() != 0) {
		fprintf(stderr, "must be run as root\n");
		exit(-1);
	}

	int child = fork();

	if (child == -1) {
		fprintf(stderr, "fork failed\n");
		exit(-2);
	}

	if (child) {
		struct jail j;
		j.version = 1;
		j.path = "/";
		j.hostname = "jail";
		j.n_ips = 0;

		int jid = jail(&j);
		if (jid < 0) {
			fprintf(stderr, "jail failed\n");
			exit(-1); // TODO
		}

		setgid(unpriv_gid);
		setuid(unpriv_uid);

		if (getuid() != unpriv_uid || getgid() != unpriv_gid) {
			fprintf(stderr, "setuid/gid failed\n");
			exit(-1); // TODO
		}
		test(fn, expected, msg, "as jailed unpriv");
		exit(0);
	}
	else {
		waitpid(child, NULL, 0);
	}
}