#!/usr/bin/python
# seccomp_test: test seccomp with bytecode
# Copyright (C) 2004-2006  Andrea Arcangeli <andrea@cpushare.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation;
# only version 2.1 of the License.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import popen2, re, time, struct, signal, os, sys, errno

# local
from cpushare.seccomp_gen import seccomp_gen_class

HEAP_SIZE = 10*1024*1024
STACK_SIZE = 8*1024*1024

class seccomp_loader_class(object):
	def __init__(self, heap_size, stack_size, home, bytecode):
		print 'init'
		self.popen3 = popen2.Popen3(('/usr/lib/cpushare/seccomp-loader',))
		if self.popen3.childerr != None:
			raise Exception('x')

		seccomp_gen_obj = seccomp_gen_class('/usr/lib/cpushare/' + bytecode)
		self.obj = seccomp_gen_obj._build_header(heap_size, stack_size) + seccomp_gen_obj.text_data

	def load(self):
		print 'load'
		self.popen3.tochild.write(self.obj)
		self.popen3.tochild.flush()

	def start(self):
		print 'start'
		sys.stdout.flush()

		seccomp_file = '/proc/' + str(self.popen3.pid) + '/seccomp'

		try:
			if file(seccomp_file, 'r').read(1) != '0':
				raise Exception('seccomp already enabled?')
		except IOError, err:
			if err[0] != errno.ENOENT:
				raise
			print 'prctl API'
			self.popen3.tochild.write(chr(210)) # seccomp enabled by the seccomp-loader
			self.popen3.tochild.flush()
			x = self.popen3.fromchild.read(1) # wait
			assert x == chr(211)
			self.popen3.tochild.write(chr(212))
		else:
			print '/proc/ API'
			self.popen3.tochild.write(chr(209)) # seccomp enabled by python
			self.popen3.tochild.flush()

			x = self.popen3.fromchild.read(1) # wait
			assert x == chr(211)

			file(seccomp_file, 'w').write('1')

			if file(seccomp_file, 'r').read(1) != '1':
				self.kill()
				raise Exception('seccomp enable failure')

			self.popen3.tochild.write(chr(212))

		self.popen3.tochild.flush()

	def stop(self):
		print 'stop'
		os.kill(self.popen3.pid, signal.SIGQUIT)
		# give it 1 seconds to call exit
		# this wait is not really necessary
		# this is only to test that the task can call
		# exit successfully too
		# sending a SIGTERM stright away would be just fine
		# SIGCHLD will wake us up so we're not really
		# going to wait 1 sec
		time.sleep(1)

	def kill(self):
		print 'kill'
		os.kill(self.popen3.pid, signal.SIGTERM)
		self.popen3.fromchild.close()
		self.popen3.tochild.close()

	def receive_data(self):
		try:
			print struct.unpack('!L', self.popen3.fromchild.read(struct.calcsize('!L')))[0], 'counts'
		except:
			print 'receive_data failure'

class ChildIsGone(Exception):
	pass

def sigchld_handler(signal, frame):
	raise ChildIsGone()

def seccomp_load(bytecode):
	seccomp_loader = seccomp_loader_class(HEAP_SIZE, STACK_SIZE, HOME, bytecode)
	try:
		seccomp_loader.load()

		signal.signal(signal.SIGCHLD, sigchld_handler)
		try:
			seccomp_loader.start()
			time.sleep(1) # 1 seconds of computations
			seccomp_loader.stop()
		except ChildIsGone:
			pass
		signal.signal(signal.SIGCHLD, signal.SIG_DFL)

		seccomp_loader.receive_data()
		seccomp_loader.kill()
		seccomp_loader.popen3.wait()
	finally:
		exit_code = seccomp_loader.popen3.poll()
		print 'exit_code %d signal %d' % (exit_code >> 8, exit_code & 0xff)
		return exit_code

def system_check():
	if [int(x) for x in re.search(r'^(\d+)\.(\d+)\.(\d+)', os.uname()[2]).groups()] < [ 2, 6, 8, ]:
		print "You need to upgrade your kernel to version 2.6.8 or more recent to run this."
		sys.exit(1)

	if sys.version_info < ( 2, 3, 3, ):
		print "You need to upgrade python to version 2.3.3 or more recent to run this."
		sys.exit(1)

	if not os.getuid() or not os.geteuid() or not os.getgid() or not os.getegid():
		print "You shouldn't run this program as root."
		sys.exit(1)

def build():
	global HOME
	HOME = os.path.join(os.getcwd(), os.path.dirname(sys.argv[0]))
	return

	if os.system('cd %s && make seccomp_test' % HOME):
		print '`make` executable has failed'
		sys.exit(1)

def seccomp_test():
	#system_check()

	print 'Starting computing some safe bytecode'
	error = seccomp_load('bytecode-safe')
	if error:
		if error == 15 << 8:
			print 'FAILURE: please recompile the Linux Kernel with CONFIG_SECCOMP=y'
		else:
			print 'FAILURE: the safe bytecode failed to compute - please notify andrea@cpushare.com'
		sys.exit(1)
	print 'The safe bytecode computed successfully'

	print 'Starting computing some malicious bytecode'
	if not seccomp_load('bytecode-malicious'):
		print 'FAILURE: the malicious bytecode broke seccomp - please notify andrea@cpushare.com'
		sys.exit(1)
	print 'The malicious bytecode has been killed successfully by seccomp'

	print 'The seccomp test completed successfully, thank you for testing.'

if __name__ == '__main__':
	build()
	seccomp_test()
