#!/usr/bin/python2
import os, sys, stat, signal, errno, argparse, time

KILL_PROCESS_TIMEOUT = 5
KILL_ALL_PROCESSES_TIMEOUT = 5

LOG_LEVEL_ERROR = 1
LOG_LEVEL_WARN  = 1
LOG_LEVEL_INFO  = 2
LOG_LEVEL_DEBUG = 3

log_level = None

class AlarmException(Exception):
	pass

def error(message):
	if log_level >= LOG_LEVEL_ERROR:
		sys.stderr.write("*** %s\n" % message)

def warn(message):
	if log_level >= LOG_LEVEL_WARN:
		print("*** %s" % message)

def info(message):
	if log_level >= LOG_LEVEL_INFO:
		print("*** %s" % message)

def debug(message):
	if log_level >= LOG_LEVEL_DEBUG:
		print("*** %s" % message)

def ignore_signals_and_raise_keyboard_interrupt(signame):
	signal.signal(signal.SIGTERM, signal.SIG_IGN)
	signal.signal(signal.SIGINT, signal.SIG_IGN)
	raise KeyboardInterrupt(signame)

def raise_alarm_exception():
	raise AlarmException('Alarm')

def listdir(path):
	try:
		result = os.stat(path)
	except OSError:
		return []
	if stat.S_ISDIR(result.st_mode):
		return sorted(os.listdir(path))
	else:
		return []

def is_exe(path):
	try:
		return os.path.isfile(path) and os.access(path, os.X_OK)
	except OSError:
		return False

def waitpid_reap_other_children(pid):
	done = False
	status = None
	while not done:
		this_pid, status = os.waitpid(-1, 0)
		done = this_pid == pid
	return status

def stop_child_process(name, pid, signo = signal.SIGTERM, time_limit = KILL_PROCESS_TIMEOUT):
	info("Shutting down %s (PID %d)..." % (name, pid))
	try:
		os.kill(pid, signo)
	except OSError:
		pass
	signal.alarm(time_limit)
	try:
		try:
			waitpid_reap_other_children(pid)
		except OSError:
			pass
	except AlarmException:
		warn("%s (PID %d) did not shut down in time. Forcing it to exit.")
		try:
			os.kill(pid, signal.SIGKILL)
		except OSError:
			pass
		try:
			waitpid_reap_other_children(pid)
		except OSError:
			pass
	finally:
		signal.alarm(0)

def run_command_killable(*argv):
	filename = argv[0]
	status = None
	pid = os.spawnvp(os.P_NOWAIT, filename, argv)
	try:
		status = waitpid_reap_other_children(pid)
	except BaseException as s:
		warn("An error occurred. Aborting.")
		stop_child_process(filename, pid)
		raise
	if status != 0:
		error("%s failed with exit code %d\n" % (filename, status))
		sys.exit(1)

def kill_all_processes(time_limit):
	info("Killing all processes...")
	try:
		os.kill(-1, signal.SIGTERM)
	except OSError:
		pass
	signal.alarm(time_limit)
	try:
		# Wait until no more child processes exist.
		done = False
		while not done:
			try:
				os.waitpid(-1, 0)
			except OSError as e:
				if e.errno == errno.ECHILD:
					done = True
				else:
					raise
	except AlarmException:
		warn("Not all processes have exited in time. Forcing them to exit.")
		try:
			os.kill(-1, signal.SIGKILL)
		except OSError:
			pass
	finally:
		signal.alarm(0)

def run_startup_files():
	# Run /etc/my_init.d/*
	for name in listdir("/etc/my_init.d"):
		filename = "/etc/my_init.d/" + name
		if is_exe(filename):
			info("Running %s..." % filename)
			run_command_killable(filename)

	# Run /etc/rc.local.
	if is_exe("/etc/rc.local"):
		info("Running /etc/rc.local...")
		run_command_killable("/etc/rc.local")

def start_runit():
	info("Booting runit daemon...")
	pid = os.spawnl(os.P_NOWAIT, "/usr/bin/runsvdir", "/usr/bin/runsvdir",
		"-P", "/etc/service", "log: %s" % ('.' * 395))
	info("Runit started as PID %d" % pid)
	return pid

def wait_for_runit_or_interrupt(pid):
	try:
		status = waitpid_reap_other_children(pid)
		return (True, status)
	except KeyboardInterrupt:
		return (False, None)

def shutdown_runit_services():
	debug("Begin shutting down runit services...")
	os.system("/usr/bin/sv down /etc/service/*")

def wait_for_runit_services():
	debug("Waiting for runit services to exit...")
	done = False
	while not done:
		done = os.system("/usr/bin/sv status /etc/service/* | grep -q '^run:'") != 0
		if not done:
			time.sleep(0.1)

def main(args):
	if not args.skip_startup_files:
		run_startup_files()
	
	runit_exited = False
	exit_code = None

	if not args.skip_runit:
		runit_pid = start_runit()
	try:
		if len(args.main_command) == 0:
			runit_exited, exit_code = wait_for_runit_or_interrupt(runit_pid)
			if runit_exited:
				info("Runit exited with code %d" % exit_code)
		else:
			info("Running %s..." % " ".join(args.main_command))
			pid = os.spawnvp(os.P_NOWAIT, args.main_command[0], args.main_command)
			try:
				exit_code = waitpid_reap_other_children(pid)
				info("%s exited with exit code %d." % (args.main_command[0], exit_code))
			except KeyboardInterrupt:
				stop_child_process(args.main_command[0], pid)
			except BaseException as s:
				warn("An error occurred. Aborting.")
				stop_child_process(args.main_command[0], pid)
				raise
		sys.exit(exit_code)
	finally:
		if not args.skip_runit:
			shutdown_runit_services()
			if not runit_exited:
				stop_child_process("runit daemon", runit_pid)
			wait_for_runit_services()

# Parse options.
parser = argparse.ArgumentParser(description = 'Initialize the system.')
parser.add_argument('main_command', metavar = 'MAIN_COMMAND', type = str, nargs = '*',
	help = 'The main command to run. (default: runit)')
parser.add_argument('--skip-startup-files', dest = 'skip_startup_files',
	action = 'store_const', const = True, default = False,
	help = 'Skip running /etc/my_init.d/* and /etc/rc.local')
parser.add_argument('--skip-runit', dest = 'skip_runit',
	action = 'store_const', const = True, default = False,
	help = 'Do not run runit services')
parser.add_argument('--no-kill-all-on-exit', dest = 'kill_all_on_exit',
	action = 'store_const', const = False, default = True,
	help = 'Don\'t kill all processes on the system upon exiting')
parser.add_argument('--quiet', dest = 'log_level',
	action = 'store_const', const = LOG_LEVEL_WARN, default = LOG_LEVEL_INFO,
	help = 'Only print warnings and errors')
args = parser.parse_args()
log_level = args.log_level

if args.skip_runit and len(args.main_command) == 0:
	error("When --skip-runit is given, you must also pass a main command.")
	sys.exit(1)

# Run main function.
signal.signal(signal.SIGTERM, lambda signum, frame: ignore_signals_and_raise_keyboard_interrupt('SIGTERM'))
signal.signal(signal.SIGINT, lambda signum, frame: ignore_signals_and_raise_keyboard_interrupt('SIGINT'))
signal.signal(signal.SIGALRM, lambda signum, frame: raise_alarm_exception())
try:
	main(args)
except KeyboardInterrupt:
	warn("Init system aborted.")
	exit(2)
finally:
	if args.kill_all_on_exit:
		kill_all_processes(KILL_ALL_PROCESSES_TIMEOUT)