#include <sys/types.h>
#include <sys/stat.h>
#include <sys/param.h>
#include <sys/sysctl.h>
#include <sys/proc.h>
#include <sys/user.h>
#include <sys/wait.h>

#include <paths.h>
#include <signal.h>
#include <stdio.h>
#include <syslog.h>
#include <err.h>
#include <errno.h>
#include <stdarg.h>
#include <unistd.h>
#include <kvm.h>
#include <limits.h>
#include <fcntl.h>

#include "util.h"

#ifdef HAVE_CONFIG_H
#include "../config.h"
#endif

/* The timeout to wait between kills */
#define		DEFAULT_TIMEOUT		3
int g_timeout = DEFAULT_TIMEOUT;

/* Supress warnings */
int g_quiet = 0;
int g_verbose = 0;

static void kill_jail_processes(kvm_t* kd, int sig);
static int kill_jail(const char* jail, int usescripts, int restart, int force);
static int check_running_processes(kvm_t* kd);
static void usage();

int main(int argc, char* argv[])
{
	int ch, r, jid;
	int ret = 0;
	int restart = 0;
	int force = 0;
	int verbose = 0;
	int usescripts = 1;
	pid_t child;

	while((ch = getopt(argc, argv, "fhkqrt:v")) != -1)
	{
		switch(ch)
		{
		case 'f':
			force = 1;
			break;

		case 'h':
			/* dummy for compatibility with killjail */
			warnx("the '-h' option has been depreciated");
			break;

		case 'k':
			usescripts = 0;
			break;

		case 'q':
			g_quiet = 1;
			g_verbose = 0;
			break;

		case 'r':
			restart = 1;
			break;

		/* Timeout to use between kills */
		case 't':
			g_timeout = atoi(optarg);
			if(g_timeout <= 0)
				errx(2, "invalid timeout argument: %s", optarg);
			break;

		case 'v':
			g_verbose = 1;
			g_quiet = 0;
			break;

		case '?':
		default:
			usage();
		}
	}

	argc -= optind;
	argv += optind;

	/* Make sure we have a jail id or name */
	if(argc == 0)
		usage();

	if(!usescripts && restart)
		usage();

	if(running_in_jail())
		errx(1, "can't run inside jail");

	/* For each jail */
	for(; argc > 0; argc--, argv++)
	{
		jid = translate_jail_name(argv[0]);
		if(jid == -1)
		{
			warnx("unknown jail host name: %s", argv[0]);
			ret = 1;
			continue;
		}

		/*
		 * We fork and the child goes into the jail and
		 * does the dirty work.
	   */
#ifdef _DEBUG
	  switch((child = fork()))
	  {
	  /* Error condition */
	  case -1:
	  	err(1, "couldn't fork child process");
	  	break;

	  /* The child */
	  case 0:
#endif
	  	if(jail_attach(jid) == -1)
				err(1, "couldn't attach to jail");

	  	r = kill_jail(argv[0], usescripts, restart, force);
	  	exit(r);
#ifdef _DEBUG
	  	break;

	  /* The parent */
	  default:
	  	if(waitpid(child, &r, 0) == -1)
	  		err(1, "error waiting for child process");

	  	if(WEXITSTATUS(r) != 0)
	  		ret = WEXITSTATUS(r);
	  	break;
	  };
#endif

		argc--;
		argv++;
	}

	return ret;
}

#define SHUTDOWN_SCRIPT "/etc/rc.shutdown"
static char* SHUTDOWN_ARGS[] = {	_PATH_BSHELL, SHUTDOWN_SCRIPT };

#define START_SCRIPT "/etc/rc"
static char* START_ARGS[] = {	_PATH_BSHELL, START_SCRIPT };

static int kill_jail(const char* jail, int usescripts, int restart, int force)
{
	kvm_t* kd = NULL;
	char errbuf[_POSIX2_LINE_MAX];
	int pass = 0;
	int timeout = 0;
	int ret = 0;
	int cmdargs = JAIL_RUN_CONSOLE;

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

	if(g_verbose)
		cmdargs |= JAIL_RUN_STDERR;

	/*
	 * Multiple passes are used to do different things.
	 * Each time the jails processes are listed.
	 */
	while(1)
	{
		while(timeout > 0)
		{
			sleep(1);
			timeout--;

			if(!check_running_processes(kd))
				goto done;
		}

		switch(pass)
		{
		/* First pass is an orderly shutdown */
		case 0:

			/* Check if we have an executable shutdown script */
			if(usescripts && check_jail_command(jail, SHUTDOWN_SCRIPT))
				run_jail_command(jail, SHUTDOWN_ARGS[0], SHUTDOWN_ARGS, cmdargs);

			break;

		/* Okay now quit all processes in jail */
		case 1:
			kill_jail_processes(kd, SIGTERM);
			timeout = g_timeout;
			break;

		/* ... and again ... */
		case 2:
			kill_jail_processes(kd, SIGTERM);
			timeout = g_timeout;
			break;

		/* Okay now we force kill the processes if necessary */
		case 3:

			if(force)
			{
				/* If we get here, jailer looks like it's really irresponsive */
				if(!g_quiet)
					warnx("%s: jail won't stop. forcing jail termination...", jail);

				kill_jail_processes(kd, SIGKILL);
				timeout = g_timeout;
			}

			break;

		case 4:

			/* And if that didn't do it, well then give up */
			if(!g_quiet)
				warnx("%s: couldn't stop jail, processes wouldn't die", jail);

			ret = 1;
			goto done;
		}

		pass++;

		if(!check_running_processes(kd))
			goto done;
	}

done:
	if(restart)
	{
		/* Check if we have an executable shutdown script */
		if(check_jail_command(jail, START_SCRIPT))
			run_jail_command(jail, START_ARGS[0], START_ARGS, cmdargs);
	}

	if(kd != NULL)
		kvm_close(kd);

	return ret;
}

static void kill_jail_processes(kvm_t* kd, int sig)
{
	struct kinfo_proc* kp;
	int nentries, i;
	pid_t cur;

	cur = getpid();

	/* Get a process listing */
	if((kp = kvm_getprocs(kd, KERN_PROC_ALL, 0, &nentries)) == 0)
		errx(1, "couldn't list processes: %s", kvm_geterr(kd));

	/* Okay now loop and look at each process' jail */
	for(i = 0; i < nentries; i++)
	{
		if(kp[i].ki_pid == cur)
			continue;

		if(kill(kp[i].ki_pid, sig) == -1)
		{
			if(errno != ESRCH)
				errx(1, "couldn't signal process: %d", (int)kp[i].ki_pid);
		}
	}
}

static int check_running_processes(kvm_t* kd)
{
	struct kinfo_proc* kp;
	int nentries, i;
	pid_t cur;

	cur = getpid();

	/* Get a process listing */
	if((kp = kvm_getprocs(kd, KERN_PROC_ALL, 0, &nentries)) == 0)
		errx(1, "couldn't list processes: %s", kvm_geterr(kd));

	if(nentries != 1)
		return 1;

	/* Okay now loop and look at each process' jail */
	for(i = 0; i < nentries; i++)
	{
		if(kp[i].ki_pid != cur)
			return 1;
	}

	return 0;
}

static void usage()
{
	fprintf(stderr, "usage: jkill [-fkqv] [-t timeout] jail ...\n");
	fprintf(stderr, "usage: jkill -r [-fqv] [-t timeout] jail ...\n");
	exit(2);
}