#!/usr/bin/python -E
#
# Author(s):	Caleb Case <ccase@tresys.com>
#
# Adapted from the bash/awk scripts mkflask.sh and mkaccess_vector.sh
#

import getopt
import os
import sys
import re

class ParseError(Exception):
	def __init__(self, type, file, line):
		self.type = type
		self.file = file
		self.line = line
	def __str__(self):
		typeS = self.type
		if type(self.type) is not str: typeS = Flask.CONSTANT_S[self.type]
		return "Parse Error: Unexpected %s on line %d of %s." % (typeS, self.line, self.file)

class DuplicateError(Exception):
	def __init__(self, type, file, line, symbol):
		self.type = type
		self.file = file
		self.line = line
		self.symbol = symbol
	def __str__(self):
		typeS = self.type
		if type(self.type) is not str: typeS = Flask.CONSTANT_S[self.type]
		return "Duplicate Error: Duplicate %s '%s' on line %d of %s." % (typeS, self.symbol, self.line, self.file)

class UndefinedError(Exception):
	def __init__(self, type, file, line, symbol):
		self.type = type
		self.file = file
		self.line = line
		self.symbol = symbol
	def __str__(self):
		typeS = self.type
		if type(self.type) is not str: typeS = Flask.CONSTANT_S[self.type]
		return "Undefined Error: %s '%s' is not defined but used on line %d of %s." % (typeS, self.symbol, self.line, self.file)

class UnusedError(Exception):
	def __init__(self, info):
		self.info = info
	def __str__(self):
		return "Unused Error: %s" % self.info

class Flask:
	'''
	FLASK container class with utilities for parsing definition
	files and creating c header files.
	'''

	#Constants used in definitions parsing.
	WHITE    = re.compile(r'^\s*$')
	COMMENT  = re.compile(r'^\s*#')
	USERFLAG = re.compile(r'# userspace')
	CLASS    = re.compile(r'^class (?P<name>\w+)')
	COMMON   = re.compile(r'^common (?P<name>\w+)')
	INHERITS = re.compile(r'^inherits (?P<name>\w+)')
	OPENB    = re.compile(r'^{')
	VECTOR   = re.compile(r'^\s*(?P<name>\w+)')
	CLOSEB   = re.compile(r'^}')
	SID      = re.compile(r'^sid (?P<name>\w+)')
	EOF      = "end of file"

	#Constants used in header generation.
	USERSPACE = 0
	KERNEL    = 1

	CONSTANT_S = { \
		#parsing constants
		WHITE    : "whitespace", \
		COMMENT  : "comment", \
		USERFLAG : "userspace flag", \
		CLASS    : "class definition", \
		COMMON   : "common definition", \
		INHERITS : "inherits definition", \
		OPENB    : "'{'", \
		VECTOR   : "access vector definition", \
		CLOSEB   : "'}'", \
		SID      : "security identifier", \
		EOF      : "end of file", \
		#generation constants
		USERSPACE : "userspace mode", \
		KERNEL    : "kernel mode", \
	}

	def __init__(self, warn = True):
		self.WARN = warn
		self.autogen   = "/* This file is automatically generated.  Do not edit. */\n"
		self.commons   = []
		self.common    = {}
		self.classes   = []
		self.vectors   = []
		self.vector    = {}
		self.userspace = {}
		self.sids      = []
		self.inherits  = {}
	
	def warning(self, msg):
		'''
		Prints a warning message out to stderr if warnings are enabled.
		'''
		if self.WARN: sys.stderr.write("Warning: %s\n" % msg)

	def parseClasses(self, path):
		'''
		Parses security class definitions from the given path.
		'''
		classes = []
		input = open(path, 'r')

		number = 0
		for line in input:
			number += 1
			m = self.COMMENT.search(line)
			if m: continue

			m = self.WHITE.search(line)
			if m: continue

			m = self.CLASS.search(line)
			if m:
				g = m.groupdict()
				c = g['name']
				if c in classes: raise DuplicateError, (self.CLASS, path, number, c)
				classes.append(c)
				if self.USERFLAG.search(line):
					self.userspace[c] = True
				continue

			raise ParseError, ("data.  Was expecting either a comment, whitespace, or class definition. ", path, number)

		self.classes = classes
		return classes

	def parseSids(self, path):
		'''
		Parses initial SID definitions from the given path.
		'''

		sids = []
		input = open(path, 'r')
		for line in input:
			m = self.COMMENT.search(line)
			if m: continue

			m = self.WHITE.search(line)
			if m: continue

			m = self.SID.search(line)
			if m:
				g = m.groupdict()
				s = g['name']
				if s in sids: raise DuplicateError, (self.SID, path, number, s)
				sids.append(s)
				continue
			
			raise ParseError, ("data. Was expecting either a comment, whitespace, or security identifier. ", path, number)

		self.sids = sids
		return sids

	def parseVectors(self, path):
		'''
		Parses access vector definitions from the given path.
		'''
		vectors = []
		vector  = {}
		commons = []
		common = {}
		inherits = {}
		input = open(path, 'r')

		# states
		NONE    = 0
		COMMON  = 1
		CLASS   = 2
		INHERIT = 3
		OPEN    = 4

		state = NONE
		state2 = NONE
		number = 0
		for line in input:
			number += 1
			m = self.COMMENT.search(line)
			if m: continue

			m = self.WHITE.search(line)
			if m: 
				if state == INHERIT:
					state = NONE
				continue

			m = self.COMMON.search(line)
			if m:
				if state != NONE: raise ParseError, (self.COMMON, path, number)
				g = m.groupdict()
				c = g['name']
				if c in commons: raise DuplicateError, (self.COMMON, path, number, c)
				commons.append(c)
				common[c] = []
				state = COMMON
				continue

			m = self.CLASS.search(line)
			if m:
				if state != NONE: raise ParseError, (self.CLASS, number)
				g = m.groupdict()
				c = g['name']
				if c in vectors: raise DuplicateError, (self.CLASS, path, number, c)
				if c not in self.classes: raise UndefinedError, (self.CLASS, path, number, c)
				vectors.append(c)
				vector[c] = []
				state = CLASS
				continue
			
			m = self.INHERITS.search(line)
			if m:
				if state != CLASS: raise ParseError, (self.INHERITS, number)
				g = m.groupdict()
				i = g['name']
				if c in inherits: raise DuplicateError, (self.INHERITS, path, number, c)
				if i not in common: raise UndefinedError, (self.COMMON, path, number, i)
				inherits[c] = i
				state = INHERIT
				continue

			m = self.OPENB.search(line)
			if m:
				if (state != CLASS \
				and state != INHERIT \
				and state != COMMON) \
				or state2 != NONE: 
					raise ParseError, (self.OPENB, path, number)
				state2 = OPEN
				continue

			m = self.VECTOR.search(line)
			if m:
				if state2 != OPEN: raise ParseError, (self.VECTOR, path, number)
				g = m.groupdict()
				v = g['name']
				if state == CLASS or state == INHERIT:
					if v in vector[c]: raise DuplicateError, (self.VECTOR, path, number, v)
					vector[c].append(v)
				elif state == COMMON:
					if v in common[c]: raise DuplicateError, (self.VECTOR, path, number, v)
					common[c].append(v)
				continue

			m = self.CLOSEB.search(line)
			if m:
				if state2 != OPEN: raise ParseError, (self.CLOSEB, path, number)
				state = NONE
				state2 = NONE
				c = None
				continue
			
			raise ParseError, ("data", path, number)

		if state != NONE and state2 != NONE: raise ParseError, (self.EOF, path, number)

		cvdiff = set(self.classes) - set(vectors)
		if cvdiff: raise UnusedError, "Not all security classes were used in access vectors: %s" % cvdiff # the inverse of this will be caught as an undefined class error

		self.commons = commons
		self.common = common
		self.vectors = vectors
		self.vector = vector
		self.inherits = inherits
		return vector

	def createHeaders(self, path, mode = USERSPACE):
		'''
		Creates the C header files in the specified MODE and outputs
		them to give PATH.
		'''
		headers = { \
			'av_inherit.h'            : self.createAvInheritH(mode), \
			'av_perm_to_string.h'     : self.createAvPermToStringH(mode), \
			'av_permissions.h'        : self.createAvPermissionsH(mode), \
			'class_to_string.h'       : self.createClassToStringH(mode), \
			'common_perm_to_string.h' : self.createCommonPermToStringH(mode), \
			'flask.h'                 : self.createFlaskH(mode), \
			'initial_sid_to_string.h' : self.createInitialSidToStringH(mode) \
		}

		for key, value in headers.items():
			of = open(os.path.join(path, key), 'w')
			of.writelines(value)
			of.close()

	def createUL(self, count):
		fields = [1, 2, 4, 8]
		return "0x%08xUL" % (fields[count % 4] << 4 * (count / 4))

	def createAvInheritH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		for c in self.vectors:
			if self.inherits.has_key(c):
				i = self.inherits[c]
				count = len(self.common[i])
				user = self.userspace.has_key(c)
				if mode == self.KERNEL and user:
					results.append("   S_(0, 0, 0)\n")
				else:
					results.append("   S_(SECCLASS_%s, %s, %s)\n" % (c.upper(), i, self.createUL(count)))
		return results

	def createAvPermToStringH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		for c in self.vectors:
			for p in self.vector[c]:
				user = self.userspace.has_key(c)
				if (mode == self.KERNEL and not user) or (mode == self.USERSPACE):
					results.append("   S_(SECCLASS_%s, %s__%s, \"%s\")\n" % (c.upper(), c.upper(), p.upper(), p))

		return results

	def createAvPermissionsH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)

		width = 57
		count = 0
		for common in self.commons:
			count = 0
			shift = 0
			for p in self.common[common]:
				columnA = "#define COMMON_%s__%s " % (common.upper(), p.upper())
				columnA += "".join([" " for i in range(width - len(columnA))])
				results.append("%s%s\n" % (columnA, self.createUL(count)))
				count += 1

		width = 50 # broken for old tools whitespace
		for c in self.vectors:
			count = 0

			ps = []
			if self.inherits.has_key(c):
				ps += self.common[self.inherits[c]]
			ps += self.vector[c]
			for p in ps: 
				columnA = "#define %s__%s " % (c.upper(), p.upper())
				columnA += "".join([" " for i in range(width - len(columnA))])
				user = self.userspace.has_key(c)
				if not (mode == self.KERNEL and user):
					results.append("%s%s\n" % (columnA, self.createUL(count)))
				count += 1

		return results

	def createClassToStringH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		results.append("/*\n * Security object class definitions\n */\n")

		if mode == self.KERNEL:
			results.append("    S_(NULL)\n")
		else:
			results.append("    S_(\"null\")\n")

		for c in self.classes:
			user = self.userspace.has_key(c)
			if mode == self.KERNEL and user:
				results.append("    S_(NULL)\n")
			else:
				results.append("    S_(\"%s\")\n" % c)
		return results

	def createCommonPermToStringH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		for common in self.commons:
			results.append("TB_(common_%s_perm_to_string)\n" % common)
			for p in self.common[common]:
				results.append("    S_(\"%s\")\n" % p)
			results.append("TE_(common_%s_perm_to_string)\n\n" % common)
		return results
	
	def createFlaskH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		results.append("#ifndef _SELINUX_FLASK_H_\n")
		results.append("#define _SELINUX_FLASK_H_\n")
		results.append("\n")
		results.append("/*\n")
		results.append(" * Security object class definitions\n")
		results.append(" */\n")

		count = 0
		width = 57
		for c in self.classes:
			count += 1
			columnA = "#define SECCLASS_%s " % c.upper()
			columnA += "".join([" " for i in range(width - len(columnA))])
			user = self.userspace.has_key(c)
			if not (mode == self.KERNEL and user):
				results.append("%s%d\n" % (columnA, count))

		results.append("\n")
		results.append("/*\n")
		results.append(" * Security identifier indices for initial entities\n")
		results.append(" */\n")
		
		count = 0
		width = 56 # broken for old tools whitespace
		for s in self.sids:
			count += 1
			columnA = "#define SECINITSID_%s " % s.upper()
			columnA += "".join([" " for i in range(width - len(columnA))])
			results.append("%s%d\n" % (columnA, count))

		results.append("\n")
		columnA = "#define SECINITSID_NUM "
		columnA += "".join([" " for i in range(width - len(columnA))])
		results.append("%s%d\n" % (columnA, count))

		results.append("\n")
		results.append("#endif\n")
		return results



	def createInitialSidToStringH(self, mode = USERSPACE):
		'''
		'''
		results = []
		results.append(self.autogen)
		results.append("static char *initial_sid_to_string[] =\n")
		results.append("{\n")
		results.append("    \"null\",\n")
		for s in self.sids:
			results.append("    \"%s\",\n" % s)
		results.append("};\n")
		results.append("\n")

		return results

def usage():
	'''
	Returns the usage string.
	'''
	usage  = 'Usage: %s -a ACCESS_VECTORS -i INITIAL_SIDS -s SECURITY_CLASSES -o OUTPUT_DIRECTORY -k|-u [-w]\n' % os.path.basename(sys.argv[0])
	usage += '\n'
	usage += ' -a --access_vectors\taccess vector definitions\n'
	usage += ' -i --initial_sids\tinitial sid definitions\n'
	usage += ' -s --security_classes\tsecurity class definitions\n'
	usage += ' -o --output\toutput directory for generated files\n'
	usage += ' -k --kernel\toutput mode set to kernel (kernel headers contain empty blocks for all classes specified with # userspace in the security_classes file)\n'
	usage += ' -u --user\toutput mode set to userspace\n'
	usage += ' -w --nowarnings\tsupresses output of warning messages\n'
	return usage

########## MAIN ##########
if __name__ == '__main__':
	
	# Parse command line args
	try:
		opts, args = getopt.getopt(sys.argv[1:], 'a:i:s:o:kuwh', ['access_vectors=', 'initial_sids=', 'security_classes=', 'output=', 'kernel', 'user', 'nowarnings', 'help'])
	except getopt.GetoptError:
		print(usage())
		sys.exit(2)
	
	avec = None
	isid = None
	secc = None
	outd = None
	mode = None
	warn = True
	for o, a in opts:
		if o in ('-h', '--help'):
			print(usage())
			sys.exit(0)
		elif o in ('-a', '--access_vectors'):
			avec = a
		elif o in ('-i', '--initial_sids'):
			isid = a
		elif o in ('-s', '--security_classes'):
			secc = a
		elif o in ('-o', '--output'):
			outd = a
		elif o in ('-k', '--kernel'):
			if mode != None:
				print(usage())
				sys.exit(2)
			mode = Flask.KERNEL
		elif o in ('-u', '--user'):
			if mode != None:
				print(usage())
				sys.exit(2)
			mode = Flask.USERSPACE
		elif o in ('-w', '--nowarnings'):
			warn = False
		else:
			print(usage())
			sys.exit(2)

	if avec == None or \
	   isid == None or \
	   secc == None or \
	   outd == None:
		   print(usage())
		   sys.exit(2)

	try:
		f = Flask(warn)
		f.parseSids(isid)
		f.parseClasses(secc)
		f.parseVectors(avec)
		f.createHeaders(outd, mode)
	except Exception, e:
		print(e)
		sys.exit(2)