#include <sys/types.h>
#include <sys/param.h>
#include <sys/proc.h>
#include <sys/user.h>
#include <sys/stat.h>
#include <sys/jail.h>
#include <sys/sysctl.h>
#include <sys/file.h>
#include <sys/wait.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>

#include "util.h"

extern char** environ;

size_t get_jail_sysctl(struct xprison** ret)
{
	struct xprison* xp;
	size_t len;
	*ret = NULL;

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

retry:

	if(len <= 0)
		return 0;

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

	if(sysctlbyname("security.jail.list", xp, &len, NULL, 0) == -1)
	{
		if(errno == ENOMEM)
		{
			free(xp);
			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");

	*ret = xp;
	return len / sizeof(*xp);
}

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

	len = get_jail_sysctl(&xp);
	if(len == 0)
		goto done;

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

	/* If it was all a number ... */
	if(!*e)
	{
		if(jid <= 0)
			errx(1, "invalid jail id: %s", str);

		/* Validate the number */
		for(i = 0; i < len; i++)
		{
			if(jid == xp[i].pr_id)
				goto done;
		}

		jid = -1;
	}

	jid = -1;

	for(i = 0; i < len; i++)
	{
		if(strcmp(xp[i].pr_host, str) == 0)
		{
			jid = xp[i].pr_id;
			break;
		}
	}

done:
	if(xp)
		free(xp);

	return jid;
}

/*
 * 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(NULL, NULL, NULL, 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 == EACCES || errno == ELOOP || errno == ENAMETOOLONG ||
			 errno == ENOENT || errno == ENOTDIR)
		{
			warn("%s%scan't execute in jail: %s", jail ? jail : "",
								jail ? ": " : "", cmd);
			return 0;
		}

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

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

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

	return 1;
}

int run_overlay_command(const char* jail, const char* cmd, char* env[],
												char* args[])
{
	if(args)
		execve(cmd, args, env);
	else
		execle(cmd, cmd, NULL, env);

	warn("%s%serror executing: %s: %s", jail ? jail : "",
					jail ? ": " : "", cmd);
	return 0;
}

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

	if(opts & JAIL_RUN_NOFORK)
		return run_overlay_command(jail, cmd, env, args);

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

	/* This is the child here */
	case 0:
		if(args)
			execve(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) == -1)
			err(1, "couldn't wait on child process");

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

		break;
	}

	return 1;
}

/* 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_dup_command(const char* jail, const char* cmd, char* env[],
										char* 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)
				execve(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;
			char buff[256];
			struct timeval timeout = { 0, 10000 };

			FD_ZERO(&readmask);

			/* Open the console file and write the header */
			if(opts & JAIL_RUN_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_RUN_STDOUT)
							write(STDOUT, buff, ret);

						if(opts & JAIL_RUN_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(WEXITSTATUS(status) != 0)
			{
				warnx("%s%serror executing: %s: %s", jail ? jail : "",
								jail ? ": " : "", cmd, strerror(WEXITSTATUS(status)));
				return 0;
			}

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

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

	return 1;
}


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

	memset(env, 0, sizeof(env));

#define MAKE_ENV_VAR(n)		 											\
	t = getenv(n);																\
	if(t != NULL)																	\
	{																							\
		env[j] = alloca(strlen(n) + 2 + strlen(t));	\
		sprintf(env[j], "%s=%s", (char*)(n), t);		\
		j++;																				\
	}

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

	MAKE_ENV_VAR("TERM");
	MAKE_ENV_VAR("COLUMNS");
	MAKE_ENV_VAR("LINES");

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