#!/bin/env python
import argparse
import yaml
import os
import subprocess
import time
import tempfile
import shutil
import re
import stat
import signal
import atexit
import sys
import linecache
import random
import string
import imp
import socket

prev_line = None
def traceit(f, e, a):
	if e == "line":
		lineno = f.f_lineno
		fil = f.f_globals["__file__"]
		if fil.endswith("zdtm.py"):
			global prev_line
			line = linecache.getline(fil, lineno)
			if line == prev_line:
				print "        ..."
			else:
				prev_line = line
				print "+%4d: %s" % (lineno, line.rstrip())

	return traceit


# Root dir for ns and uns flavors. All tests
# sit in the same dir
tests_root = None

def clean_tests_root():
	global tests_root
	if tests_root:
		os.rmdir(tests_root)

def make_tests_root():
	global tests_root
	if not tests_root:
		tests_root = tempfile.mkdtemp("", "criu-root-", "/tmp")
		atexit.register(clean_tests_root)
	return tests_root

# Arch we run on
arch = os.uname()[4]

#
# Flavors
#  h -- host, test is run in the same set of namespaces as criu
#  ns -- namespaces, test is run in itw own set of namespaces
#  uns -- user namespace, the same as above plus user namespace
#

class host_flavor:
	def __init__(self, opts):
		self.name = "host"
		self.ns = False
		self.root = None

	def init(self, test_bin):
		pass

	def fini(self):
		pass

class ns_flavor:
	def __init__(self, opts):
		self.name = "ns"
		self.ns = True
		self.uns = False
		self.root = make_tests_root()

	def init(self, test_bin):
		print "Construct root for %s" % test_bin
		subprocess.check_call(["mount", "--make-private", "--bind", ".", self.root])

		if not os.access(self.root + "/.constructed", os.F_OK):
			for dir in ["/bin", "/etc", "/lib", "/lib64", "/dev", "/tmp"]:
				os.mkdir(self.root + dir)
				os.chmod(self.root + dir, 0777)

			os.mknod(self.root + "/dev/tty", stat.S_IFCHR, os.makedev(5, 0))
			os.chmod(self.root + "/dev/tty", 0666)
			os.mknod(self.root + "/.constructed", stat.S_IFREG | 0600)

		ldd = subprocess.Popen(["ldd", test_bin], stdout = subprocess.PIPE)
		xl = re.compile('^(linux-gate.so|linux-vdso(64)?.so|not a dynamic)')

		# This Mayakovsky-style code gets list of libraries a binary
		# needs minus vdso and gate .so-s
		libs = map(lambda x: x[1] == '=>' and x[2] or x[0],		\
				map(lambda x: x.split(),			\
					filter(lambda x: not xl.match(x),	\
						map(lambda x: x.strip(),	\
							filter(lambda x: x.startswith('\t'), ldd.stdout.readlines())))))
		ldd.wait()

		for lib in libs:
			tlib = self.root + lib
			if not os.access(tlib, os.F_OK):
				# Copying should be atomic as tests can be
				# run in parallel
				dst = tempfile.mktemp(".tso", "", self.root + os.path.dirname(lib))
				shutil.copy2(lib, dst)
				os.rename(dst, tlib)

	def fini(self):
		subprocess.check_call(["mount", "--make-private", self.root])
		subprocess.check_call(["umount", "-l", self.root])

class userns_flavor(ns_flavor):
	def __init__(self, opts):
		ns_flavor.__init__(self, opts)
		self.name = "userns"
		self.uns = True

flavors = { 'h': host_flavor, 'ns': ns_flavor, 'uns': userns_flavor }

#
# Helpers
#

def tail(path):
	p = subprocess.Popen(['tail', '-n1', path],
			stdout = subprocess.PIPE)
	return p.stdout.readline()

def rpidfile(path):
	return open(path).readline().strip()

def wait_pid_die(pid, who, tmo = 4):
	stime = 0.1
	while stime < tmo:
		try:
			os.kill(int(pid), 0)
		except: # Died
			break

		print "Wait for %s to die for %f" % (who, stime)
		time.sleep(stime)
		stime *= 2
	else:
		raise test_fail_exc("%s die" % who)

def test_flag(tdesc, flag):
	return flag in tdesc.get('flags', '').split()

#
# Exception thrown when something inside the test goes wrong,
# e.g. test doesn't start, criu returns with non zero code or
# test checks fail
#

class test_fail_exc:
	def __init__(self, step):
		self.step = step

class test_fail_expected_exc:
	def __init__(self, cr_action):
		self.cr_action = cr_action

#
# A test from zdtm/ directory.
#

class zdtm_test:
	def __init__(self, name, desc, flavor):
		self.__name = name
		self.__desc = desc
		self.__make_action('cleanout')
		self.__pid = 0
		self.__flavor = flavor
		self.auto_reap = True

	def __make_action(self, act, env = None, root = None):
		tpath = self.__name + '.' + act
		s_args = ['make', '--no-print-directory', \
			 	'-C', os.path.dirname(tpath), \
				      os.path.basename(tpath)]

		if env:
			env = dict(os.environ, **env)

		s = subprocess.Popen(s_args, env = env, cwd = root)
		s.wait()

	def __pidfile(self):
		if self.__flavor.ns:
			return self.__name + '.init.pid'
		else:
			return self.__name + '.pid'

	def __wait_task_die(self):
		wait_pid_die(int(self.__pid), self.__name)

	def start(self):
		env = {}
		self.__flavor.init(self.__name)

		print "Start test"

		env['ZDTM_THREAD_BOMB'] = "100"
		if not test_flag(self.__desc, 'suid'):
			env['ZDTM_UID'] = "18943"
			env['ZDTM_GID'] = "58467"
			env['ZDTM_GROUPS'] = "27495 48244"
		else:
			print "Test is SUID"

		if self.__flavor.ns:
			env['ZDTM_NEWNS'] = "1"
			env['ZDTM_PIDFILE'] = os.path.realpath(self.__name + '.init.pid')
			env['ZDTM_ROOT'] = self.__flavor.root

			if self.__flavor.uns:
				env['ZDTM_USERNS'] = "1"

		self.__make_action('pid', env, self.__flavor.root)

		try:
			os.kill(int(self.getpid()), 0)
		except:
			raise test_fail_exc("start")

	def kill(self, sig = signal.SIGKILL):
		if self.__pid:
			os.kill(int(self.__pid), sig)
			self.gone(sig == signal.SIGKILL)

		self.__flavor.fini()

	def stop(self):
		self.getpid() # Read the pid from pidfile back
		self.kill(signal.SIGTERM)

		res = tail(self.__name + '.out')
		if not 'PASS' in res.split():
			raise test_fail_exc("result check")

	def getpid(self):
		if self.__pid == 0:
			self.__pid = rpidfile(self.__pidfile())

		return self.__pid

	def getname(self):
		return self.__name

	def __getcropts(self):
		opts = self.__desc.get('opts', '').split() + ["--pidfile", os.path.realpath(self.__pidfile())]
		if self.__flavor.ns:
			opts += ["--root", self.__flavor.root]
		if test_flag(self.__desc, 'crlib'):
			opts += ["-L", os.path.dirname(os.path.realpath(self.__name)) + '/lib']
		return opts

	def getdopts(self):
		return self.__getcropts()

	def getropts(self):
		return self.__getcropts()

	def gone(self, force = True):
		if not self.auto_reap:
			pid, status = os.waitpid(int(self.__pid), 0)
			if pid != int(self.__pid):
				raise test_fail_exc("kill pid mess")

		self.__wait_task_die()
		self.__pid = 0
		if force or self.__flavor.ns:
			os.unlink(self.__pidfile())

	def print_output(self):
		if os.access(self.__name + '.out', os.R_OK):
			print "Test output: " + "=" * 32
			print open(self.__name + '.out').read()
			print " <<< " + "=" * 32

	def static(self):
		return self.__name.split('/')[2] == 'static'

	def blocking(self):
		return test_flag(self.__desc, 'crfail')



class inhfd_test:
	def __init__(self, name, desc, flavor):
		self.__name = os.path.basename(name)
		print "Load %s" % name
		self.__fdtyp = imp.load_source(self.__name, name)
		self.__my_file = None
		self.__peer_pid = 0
		self.__peer_file = None
		self.__peer_file_name = None
		self.__dump_opts = None

	def start(self):
		self.__message = "".join([random.choice(string.ascii_letters) for _ in range(16)])
		(self.__my_file, peer_file) = self.__fdtyp.create_fds()

		# Check FDs returned for inter-connection
		self.__my_file.write(self.__message)
		self.__my_file.flush()
		if peer_file.read(16) != self.__message:
			raise test_fail_exc("FDs screwup")

		start_pipe = os.pipe()
		self.__peer_pid = os.fork()
		if self.__peer_pid == 0:
			os.setsid()
			os.close(0)
			os.close(1)
			os.close(2)
		 	self.__my_file.close()
			os.close(start_pipe[0])
			os.close(start_pipe[1])
			try:
				data = peer_file.read(16)
			except:
				sys.exit(1)

			sys.exit(data == self.__message and 42 or 2)

		os.close(start_pipe[1])
		os.read(start_pipe[0], 12)
		os.close(start_pipe[0])

		self.__peer_file_name = self.__fdtyp.filename(peer_file)
		self.__dump_opts = self.__fdtyp.dump_opts(peer_file)

	def stop(self):
		self.__my_file.write(self.__message)
		self.__my_file.flush()
		pid, status = os.waitpid(self.__peer_pid, 0)
		if not os.WIFEXITED(status) or os.WEXITSTATUS(status) != 42:
			raise test_fail_exc("test failed with %d" % status)

	def kill(self):
		if self.__peer_pid:
			os.kill(self.__peer_pid, signal.SIGKILL)

	def getname(self):
		return self.__name

	def getpid(self):
		return "%s" % self.__peer_pid

	def gone(self, force = True):
		os.waitpid(self.__peer_pid, 0)
		wait_pid_die(self.__peer_pid, self.__name)
		self.__my_file = None
		self.__peer_file = None

	def getdopts(self):
		return self.__dump_opts

	def getropts(self):
		(self.__my_file, self.__peer_file) = self.__fdtyp.create_fds()
		return ["--restore-sibling", "--inherit-fd", "fd[%d]:%s" % (self.__peer_file.fileno(), self.__peer_file_name)]

	def print_output(self):
		pass

	def static(self):
		return True

	def blocking(self):
		return False


test_classes = { 'zdtm': zdtm_test, 'inhfd': inhfd_test }

#
# CRIU when launched using CLI
#

class criu_cli:
	def __init__(self, opts):
		self.__test = None
		self.__dump_path = None
		self.__iter = 0
		self.__page_server = (opts['page_server'] and True or False)
		self.__restore_sibling = (opts['sibling'] and True or False)
		self.__fault = (opts['fault'])

	def set_test(self, test):
		self.__test = test
		self.__dump_path = "dump/" + test.getname() + "/" + test.getpid()
		os.makedirs(self.__dump_path)

	def cleanup(self):
		if self.__dump_path:
			print "Remvoing %s" % self.__dump_path
			shutil.rmtree(self.__dump_path)

	def __ddir(self):
		return os.path.join(self.__dump_path, "%d" % self.__iter)

	@staticmethod
	def __criu(action, args, fault = None):
		env = None
		if fault:
			print "Forcing %s fault" % fault
			env = dict(os.environ, CRIU_FAULT = fault)
		cr = subprocess.Popen(["../criu", action] + args, env = env)
		return cr.wait()

	def __criu_act(self, action, opts, log = None):
		if not log:
			log = action + ".log"

		s_args = ["-o", log, "-D", self.__ddir(), "-v4"] + opts

		print "Run CRIU: [" + action + " " + " ".join(s_args) + "]"
		ret = self.__criu(action, s_args, self.__fault)
		if ret != 0:
			if self.__fault or self.__test.blocking():
				raise test_fail_expected_exc(action)
			else:
				grep_errors(os.path.join(self.__ddir(), log))
				raise test_fail_exc("CRIU %s" % action)

	def dump(self, action, opts = []):
		self.__iter += 1
		os.mkdir(self.__ddir())

		a_opts = ["-t", self.__test.getpid()]
		if self.__iter > 1:
			a_opts += ["--prev-images-dir", "../%d" % (self.__iter - 1), "--track-mem"]

		if self.__page_server:
			print "Adding page server"
			self.__criu_act("page-server", opts = [ "--port", "12345", \
					"--daemon", "--pidfile", "ps.pid"])
			a_opts += ["--page-server", "--address", "127.0.0.1", "--port", "12345"]

		a_opts += self.__test.getdopts()

		self.__criu_act(action, opts = a_opts + opts)

		if self.__page_server:
			wait_pid_die(int(rpidfile(self.__ddir() + "/ps.pid")), "page server")

	def restore(self):
		r_opts = []
		if self.__restore_sibling:
			r_opts = ["--restore-sibling"]
			self.__test.auto_reap = False
		r_opts += self.__test.getropts()

		self.__criu_act("restore", opts = r_opts + ["--restore-detached"])

	@staticmethod
	def check(feature):
		return criu_cli.__criu("check", ["-v0", "--feature", feature]) == 0


def try_run_hook(test, args):
	hname = test.getname() + '.hook'
	if os.access(hname, os.X_OK):
		print "Running %s(%s)" % (hname, ', '.join(args))
		hook = subprocess.Popen([hname] + args)
		if hook.wait() != 0:
			raise test_fail_exc("hook " + " ".join(args))

#
# Main testing entity -- dump (probably with pre-dumps) and restore
#

def cr(cr_api, test, opts):
	if opts['nocr']:
		return

	cr_api.set_test(test)

	for i in xrange(0, int(opts['iters'] or 1)):
		for p in xrange(0, int(opts['pre'] or 0)):
			cr_api.dump("pre-dump")

		if opts['norst']:
			cr_api.dump("dump", opts = ["--leave-running"])
		else:
			cr_api.dump("dump")
			test.gone()
			try_run_hook(test, ["--pre-restore"])
			cr_api.restore()

# Additional checks that can be done outside of test process

def get_maps(test):
	maps = [[0,0]]
	last = 0
	for mp in open("/proc/%s/maps" % test.getpid()).readlines():
		m = map(lambda x: int('0x' + x, 0), mp.split()[0].split('-'))
		if maps[last][1] == m[0]:
			maps[last][1] = m[1]
		else:
			maps.append(m)
			last += 1
	maps.pop(0)
	return maps

def get_fds(test):
	return map(lambda x: int(x), os.listdir("/proc/%s/fdinfo" % test.getpid()))

def cmp_lists(m1, m2):
	return filter(lambda x: x[0] != x[1], zip(m1, m2))

def get_visible_state(test):
	if test.static():
		fds = get_fds(test)
		maps = get_maps(test)
		return (fds, maps)
	else:
		return ([], [])

def check_visible_state(test, state):
	new = get_visible_state(test)
	if cmp_lists(new[0], state[0]):
		raise test_fail_exc("fds compare")
	if cmp_lists(new[1], state[1]):
		raise test_fail_exc("maps compare")

def do_run_test(tname, tdesc, flavs, opts):
	tcname = tname.split('/')[0]
	tclass = test_classes.get(tcname, None)
	if not tclass:
		print "Unknown test class %s" % tcname
		return

	for f in flavs:
		print
		print_sep("Run %s in %s" % (tname, f))
		flav = flavors[f](opts)
		t = tclass(tname, tdesc, flav)
		cr_api = criu_cli(opts)

		try:
			t.start()
			s = get_visible_state(t)
			try:
				cr(cr_api, t, opts)
			except test_fail_expected_exc as e:
				if e.cr_action == "dump":
					t.stop()
				try_run_hook(t, ["--fault", e.cr_action])
			else:
				check_visible_state(t, s)
				t.stop()
				try_run_hook(t, ["--clean"])
		except test_fail_exc as e:
			print "Test %s FAIL at %s" % (tname, e.step)
			t.print_output()
			t.kill()
			if opts['keep_img'] == 'never':
				cr_api.cleanup()
			# This exit does two things -- exits from subprocess and
			# aborts the main script execution on the 1st error met
			sys.exit(1)
		else:
			if opts['keep_img'] != 'always':
				cr_api.cleanup()
			print_sep("Test %s PASS" % tname)

class launcher:
	def __init__(self, opts):
		self.__opts = opts
		self.__max = int(opts['parallel'] or 0)
		self.__subs = {}
		self.__fail = False

	def run_test(self, name, desc, flavor):
		if self.__max == 0:
			do_run_test(name, desc, flavor, self.__opts)
			return

		if len(self.__subs) >= self.__max:
			self.wait()
			if self.__fail:
				raise test_fail_exc('')

		nd = ('nocr', 'norst', 'pre', 'iters', 'page_server', 'sibling', 'fault', 'keep_img')
		arg = repr((name, desc, flavor, { d: self.__opts[d] for d in nd }))
		log = name.replace('/', '_') + ".log"
		sub = subprocess.Popen(["./zdtm_ct", "zdtm.py"], \
				env = dict(os.environ, CR_CT_TEST_INFO = arg ), \
				stdout = open(log, "w"), stderr = subprocess.STDOUT)
		self.__subs[sub.pid] = { 'sub': sub, 'log': log }

	def __wait_one(self, flags):
		pid, status = os.waitpid(0, flags)
		if pid != 0:
			sub = self.__subs.pop(pid)
			if status != 0:
				self.__fail = True

			print open(sub['log']).read()
			os.unlink(sub['log'])
			return True

		return False

	def wait(self):
		self.__wait_one(0)
		while self.__subs:
			if not self.__wait_one(os.WNOHANG):
				break

	def finish(self):
		while self.__subs:
			self.__wait_one(0)
		if self.__fail:
			sys.exit(1)

def all_tests(opts):
	desc = eval(open(opts['set'] + '.desc').read())
	lst = subprocess.Popen(['find', desc['dir'], '-type', 'f', '-executable' ], \
			stdout = subprocess.PIPE)
	excl = map(lambda x: os.path.join(desc['dir'], x), desc['exclude'])
	tlist = filter(lambda x: \
			not x.endswith('.checkskip') and \
			not x.endswith('.hook') and \
			not x in excl, \
				map(lambda x: x.strip(), lst.stdout.readlines()) \
		)
	lst.wait()
	return tlist


# Descriptor for abstract test not in list
default_test={ }


def get_test_desc(tname):
	d_path = tname + '.desc'
	if os.access(d_path, os.F_OK):
		return eval(open(d_path).read())

	return default_test


def self_checkskip(tname):
	chs = tname  + '.checkskip'
	if os.access(chs, os.X_OK):
		ch = subprocess.Popen([chs])
		return ch.wait() == 0 and False or True

	return False

def print_sep(title, sep = "="):
	sep_len = (80 - len(title) - 2) / 2
	sep = sep * sep_len
	print "%s %s %s" % (sep, title, sep)

def grep_errors(fname):
	print_sep("grep Error")
	for l in open(fname):
		if "Error" in l:
			print l,
	print_sep("ERROR OVER")

def run_tests(opts):
	excl = None
	features = {}

	if opts['all']:
		torun = all_tests(opts)
		run_all = True
	elif opts['tests']:
		r = re.compile(opts['tests'])
		torun = filter(lambda x: r.match(x), all_tests(opts))
		run_all = True
	elif opts['test']:
		torun = opts['test']
		run_all = False
	else:
		print "Specify test with -t <name> or -a"
		return

	if opts['exclude']:
		excl = re.compile(".*(" + "|".join(opts['exclude']) + ")")
		print "Compiled exclusion list"

	l = launcher(opts)
	try:
		for t in torun:
			global arch

			if excl and excl.match(t):
				print "Skipping %s (exclude)" % t
				continue

			tdesc = get_test_desc(t)
			if tdesc.get('arch', arch) != arch:
				print "Skipping %s (arch %s)" % (t, tdesc['arch'])
				continue

			if run_all and test_flag(tdesc, 'noauto'):
				print "Skipping test %s (manual run only)" % t
				continue

			feat = tdesc.get('feature', None)
			if feat:
				if not features.has_key(feat):
					print "Checking feature %s" % feat
					features[feat] = criu_cli.check(feat)

				if not features[feat]:
					print "Skipping %s (no %s feature)" % (t, feat)
					continue

			if self_checkskip(t):
				print "Skipping %s (self)" % t
				continue

			test_flavs = tdesc.get('flavor', 'h ns uns').split()
			opts_flavs = (opts['flavor'] or 'h,ns,uns').split(',')
			run_flavs = set(test_flavs) & set(opts_flavs)

			if run_flavs:
				l.run_test(t, tdesc, run_flavs)
	finally:
		l.finish()


sti_fmt = "%-40s%-10s%s"

def show_test_info(t):
	tdesc = get_test_desc(t)
	flavs = tdesc.get('flavor', '')
	return sti_fmt % (t, flavs, tdesc.get('flags', ''))


def list_tests(opts):
	tlist = all_tests(opts)
	if opts['info']:
		print sti_fmt % ('Name', 'Flavors', 'Flags')
		tlist = map(lambda x: show_test_info(x), tlist)
	print '\n'.join(tlist)

#
# main() starts here
#

if os.environ.has_key('CR_CT_TEST_INFO'):
	# Fork here, since we're new pidns init and are supposed to
	# collect this namespace's zombies
	status = 0
	pid = os.fork()
	if pid == 0:
		tinfo = eval(os.environ['CR_CT_TEST_INFO'])
		do_run_test(tinfo[0], tinfo[1], tinfo[2], tinfo[3])
	else:
		while True:
			wpid, status = os.wait()
			if wpid == pid:
				if not os.WIFEXITED(status) or os.WEXITSTATUS(status) != 0:
					status = 1
				break;

	sys.exit(status)

p = argparse.ArgumentParser("CRIU test suite")
p.add_argument("--debug", help = "Print what's being executed", action = 'store_true')
p.add_argument("--set", help = "Which set of tests to use", default = 'zdtm')

sp = p.add_subparsers(help = "Use --help for list of actions")

rp = sp.add_parser("run", help = "Run test(s)")
rp.set_defaults(action = run_tests)
rp.add_argument("-a", "--all", action = 'store_true')
rp.add_argument("-t", "--test", help = "Test name", action = 'append')
rp.add_argument("-T", "--tests", help = "Regexp")
rp.add_argument("-f", "--flavor", help = "Flavor to run")
rp.add_argument("-x", "--exclude", help = "Exclude tests from --all run", action = 'append')

rp.add_argument("--sibling", help = "Restore tests as siblings", action = 'store_true')
rp.add_argument("--pre", help = "Do some pre-dumps before dump")
rp.add_argument("--nocr", help = "Do not CR anything, just check test works", action = 'store_true')
rp.add_argument("--norst", help = "Don't restore tasks, leave them running after dump", action = 'store_true')
rp.add_argument("--iters", help = "Do CR cycle several times before check")
rp.add_argument("--fault", help = "Test fault injection")

rp.add_argument("--page-server", help = "Use page server dump", action = 'store_true')
rp.add_argument("-p", "--parallel", help = "Run test in parallel")

rp.add_argument("-k", "--keep-img", help = "Whether or not to keep images after test",
		choices = [ 'always', 'never', 'failed' ], default = 'failed')

lp = sp.add_parser("list", help = "List tests")
lp.set_defaults(action = list_tests)
lp.add_argument('-i', '--info', help = "Show more info about tests", action = 'store_true')

opts = vars(p.parse_args())

if opts['debug']:
	sys.settrace(traceit)

opts['action'](opts)
