#include <sys/types.h>
#include <sys/user.h>
#include <sys/proc.h>
#include <sys/param.h>
#include <sys/jail.h>
#include <sys/sysctl.h>
#include <sys/file.h>

#include <err.h>
#include <errno.h>
#include <stdlib.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <kvm.h>
#include <paths.h>

int translate_jail_name(const char* str)
{
	struct xprison* sxp;
	struct xprison* xp;
	size_t len, i;
	char* e;
	int jid;

	jid = strtol(str, &e, 10);

	/* If it was all a number ... */
	if(!*e)
	{
		if(jid > 0)
			return jid;

		errx(1, "invalid jail id: %s", str);
	}

	/* ... otherwise it's a name */

  if(sysctlbyname("security.jail.list", NULL, &len, NULL, 0) == -1)
  	err(1, "couldn't list jails");

retry:

	if(len <= 0)
		return -1;

	sxp = xp = calloc(len, 1);
	if(sxp == NULL)
		err(1, "out of memory");

	if(sysctlbyname("security.jail.list", xp, &len, NULL, 0) == -1)
	{
		if(errno == ENOMEM)
		{
			free(sxp);
			goto retry;
		}

		err(1, "couldn't list jails");
	}

	if(len < sizeof(*xp) || len % sizeof(*xp) || xp->pr_version != XPRISON_VERSION)
		errx(1, "kernel and userland out of sync");

	jid = -1;

	for(i = 0; i < (len / sizeof(*xp)); i++)
	{
		if(strcmp(xp->pr_host, str) == 0)
		{
			jid = xp->pr_id;
			break;
		}
	}

	free(sxp);
	return -1;
}

/*
 * in_jail
 * This code was written by James E. Quick mailto:jq@quick.com
 * The code may be freely re-used under the terms of the BSD copyright,
 * as long as this comment remains intact.
 */

int running_in_jail()
{
	int count;
	kvm_t* kd = 0;
	struct kinfo_proc* kp;
	char errbuf[_POSIX2_LINE_MAX];
	int	result = -1;

	kd = kvm_openfiles(_PATH_DEVNULL, _PATH_DEVNULL, _PATH_DEVNULL,
										 O_RDONLY, errbuf);
	if(kd == NULL)
		errx(1, "couldn't connect to kernel: %s", errbuf);

	kp = kvm_getprocs(kd, KERN_PROC_PID, getpid(), &count);
	if(kp == NULL)
		errx(1, "couldn't list processes: %s", kvm_geterr(kd));

	result = (kp->ki_flag & P_JAILED) ? 1 : 0;
	kvm_close(kd);

	return result;
}


int check_jail_command(const char* jail, const char* cmd)
{
	struct stat sb;

	if(stat(cmd, &sb) == -1)
	{
		if(errno == EACCESS || errno == ELOOP || errno == ENAMETOOLONG ||
			 errno == ENOENT || errno == ENOTDIR)
		{
			warn("%s%scan't execute in jail: %s", , jail ? jail : "",
								jail ? jail : ": ",	cmd);
			return 0;
		}

		err(1, "%s%scouldn't stat file: %s", , jail ? jail : "",
								jail ? jail : ": ",	cmd);
	}

	if(!(sb.st_mode & S_IFREG))
	{
		warnx("%s%snot a regular file: %s", , jail ? jail : "",
								jail ? jail : ": ",	cmd);
		return 0;
	}

	if(!(sb.st_mode & (S_IXUSR | S_IXGRP | S_IXOTH)))
	{
		warnx("%s%snot executable: %s", , jail ? jail : "",
								jail ? jail : ": ",	cmd);
		return 0;
	}

	if(sb.st_uid != 0)
	{
		warnx("%s%snot owned by root: %s", jail ? jail : "",
								jail ? jail : ": ",	cmd);
		return 0;
	}

	return 1;
}

int run_simple_command(const char* jail, const char* cmd, char* args[])
{
	pid_t pid;
	int status;

	switch((pid = fork()))
	{
	case -1:
		err(1, "couldn't fork child process");
		break;

	/* This is the child here */
	case 0:
		if(args)
			exect(cmd, args, env);
		else
			execle(cmd, cmd, NULL, env);

		exit(errno);
		break;

	/* This is the parent process */
	default:

		/* If the processes exited then break out */
		if(waitpid(pid, &status, 0) == pid)
			break;

		/* Return any status codes */
		if(status != 0)
		{
			warnx("%s%serror executing: %s: %s", jail ? jail : "",
								jail ? jail : ": ",	cmd, strerror(status));
			return 0;
		}

		break;
	}

	return 1;
}

int run_dup_command(const char* jail, const char* cmd, char* argv[] args, int opts)
{
	int outpipe[2];
	pid_t pid;

	/*
	 * Special function to duplicate output of a command to two
   * files.
   *
   * NOTE: Yes, I know this may seem like overkill, but system,
   * popen and all those guys would hang with certain rc scripts.
   * Those which opened a daemon in the background ('&') but still
   * kept their output going to the same stdin/stdout handles.
	 */

	/* Create a pipe for the child process */
	if(pipe(outpipe) < 0)
		return -1;

	switch(pid = fork())
	{
	case -1:
  	err(1, "couldn't fork child process");
  	break;

	/* This is the child here */
	case 0:
		{
			/* Fix up our end of the pipe */
			if(dup2(outpipe[WRITE_END], STDOUT) < 0 ||
			   dup2(outpipe[WRITE_END], STDERR) < 0)
				exit(errno);

			/* Okay, now run whatever command it was */
			if(args)
				exect(cmd, args, env);
			else
				execle(cmd, cmd, NULL, env);

			/* In case it returns then have to do this to get
			   children to disconnect from stdout */
			fflush(stdout);
			fclose(stdout);
			close(outpipe[WRITE_END]);

			exit(errno);
		}
		break;


	/* And this is the parent */
	default:
		{
			int console = -1;
			int ret;
			int status = 0;
			fd_set readmask;
			struct timeval timeout = { 0, 10000 };

			FD_ZERO(&readmask);

			/* Open the console file and write the header */
			if(opts & JAIL_OUT_CONSOLE)
				console = open(_PATH_CONSOLE, O_WRONLY | O_APPEND);

			/* No blocking on the child processes pipe */
			fcntl(outpipe[READ_END], F_SETFL, fcntl(outpipe[READ_END], F_GETFL, 0) | O_NONBLOCK);

			/* Loop until the process dies or no more output */
			while(1)
			{
				FD_SET(outpipe[READ_END], &readmask);

				if(select(FD_SETSIZE, &readmask, NULL, NULL, &timeout) == -1)
					err(1, "couldn't select");

				if(FD_ISSET(outpipe[READ_END], &readmask))
				{
					/* Read text */
					while((ret = read(outpipe[READ_END], buff, 256)) > 0)
					{
						if(opts & JAIL_OUT_STDOUT)
							write(STDOUT, buff, ret);

						if(opts & JAIL_OUT_STDERR)
							write(STDERR, buff, ret);

						if(console != -1)
							write(console, buff, ret);
					}
				}

				/* If the processes exited then break out */
				if(waitpid(pid, &status, WNOHANG) == pid)
					break;

				/* Or if there's an error or end of file */
				if(ret == -1 && errno != EAGAIN || ret == 0)
					break;
			}

			/* Return any status codes */
			if(status != 0)
			{
				warnx("%s%serror executing: %s: %s", jail ? jail : "",
								jail ? jail : ": ",	cmd, strerror(status));
				return 0;
			}

			/* Clean up */
			close(outpipe[READ_END]);

			if(console != -1)
				close(console);
		}
		break;
	}

	return 1;
}

#define MAKE_ENV_VAR(v, n) 											\
{																								\
	char* t = getenv(n);													\
	t = t ? t : "";																\
	(v) = alloca(strlen(n) + 2 + strlen(t));			\
	sprintf((v), "%s=%s", (n), (v));							\
}

/* read & write ends of a pipe */
#define  READ_END   0
#define  WRITE_END  1

/* pre-set file descriptors */
#define  STDIN   0
#define  STDOUT  1
#define  STDERR  2

int run_jail_command(const char* jail, const char* cmd, char* args[], int opts)
{
	char* env[5];

	/* Prepare an environment for the cmd */
	env[0] = "PATH=" _PATH_STDPATH;
	env[1] = NULL;

	/*
	 * If we're outputing to the console then we need a
	 * few more environement variables. Processes like
	 * 'ps' depend on them for formatting.
	 */
	if(!opts || opts & JAIL_OUT_STDOUT ||
		 opts & JAIL_OUT_STDERR)
	{
		/* TODO: For fun debugging. Remove later */
		t = getenv(VAR);
		t = t ? t : NULL;
		env[1] = alloca(sizeof("TERM") + 2 + strlen(t));
		sprintf(env[1], "%s=%s", "TERM", t);

		MAKE_ENV_VAR(env[1], "TERM");
		MAKE_ENV_VAR(env[2], "COLUMNS");
		MAKE_ENV_VAR(env[3], "LINES");
		env[4] = NULL;
	}

	if(opts)
		return run_dup_command(jail, cmd, args, opts);
	else
		return run_simple_command(jail, cmd, args);
}