#!/usr/bin/python
#
# gkrecord - record Gatekeeper rejection activity
#
# gkrecord filename
#
import sys
import os
import signal
import errno
import subprocess
import tempfile
import plistlib


#
# Usage and fail
#
def usage():
	print >>sys.stderr, "Usage: %s outputfile" % sys.argv[0]
	sys.exit(2)

def fail(whatever):
	print >>sys.stderr, "%s: %s" % (sys.argv[0], whatever)
	sys.exit(1)


#
# Argument processing
#
if len(sys.argv) != 2:
	usage()
outputfile = sys.argv[1]


#
# If the output file already exists, bail
#
if os.path.exists(outputfile):
	fail("already exists: %s" % outputfile)


#
# Places and things
#
collect = "/tmp/gke/"


# must be root 
if os.getuid() != 0:
	fail("Must have root privileges")


#
# Make sure Gatekeeper is disabled
#
subprocess.check_call(["/usr/sbin/spctl", "--master-disable"])


#
# make sure we have a fresh syspolicyd and get its pid
#
subprocess.check_call(["/usr/sbin/spctl", "--assess", "--ignore-cache", "/bin/ls"])
try:
	psax = subprocess.check_output("ps ax|grep syspolicyd|grep -v grep", shell=True).split("\n")
	if len(psax) != 2:	# [ found_syspolicyd, '' ]
		fail("Cannot find syspolicyd")
	spd_pid = int(psax[0].split()[0])
except subprocess.CalledProcessError:
	fail("Cannot find syspolicyd")


#
# run collector dtrace script until dtrace dies.
# recorder_mode arguments are (path, type, label, cdhash, flags)
#
DSCRIPT = '''
syspolicy$1:::recorder_mode { printf("RECORD;%d;%d", arg1, arg4); }

self unsigned char *cdhash;

syspolicy$1:::recorder_mode
{
	self->cdhash = copyin(arg3, 20);
	printf(";%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x%02.2x",
		self->cdhash[0], self->cdhash[1], self->cdhash[2], self->cdhash[3], self->cdhash[4],
		self->cdhash[5], self->cdhash[6], self->cdhash[7], self->cdhash[8], self->cdhash[9],
		self->cdhash[10], self->cdhash[11], self->cdhash[12], self->cdhash[13], self->cdhash[14],
		self->cdhash[15], self->cdhash[16], self->cdhash[17], self->cdhash[18], self->cdhash[19]);
	printf(";%s\\n", copyinstr(arg0));
}

syspolicy$1:::recorder_mode_adhoc_path
{
	printf("SIGNATURE;%d;%s;%s\\n", arg1, copyinstr(arg2), copyinstr(arg0));
}

syspolicy$1:::assess-outcome-unsigned
{
	printf("UNSIGNED;%d;%s\\n", arg1, copyinstr(arg0));
}

syspolicy$1:::assess-outcome-broken
{
	printf("BROKEN;%d;%d;%s\\n", arg1, arg2, copyinstr(arg0));
}
'''

def sigint(sig, ctx):
	os.kill(spd_pid, signal.SIGINT)
signal.signal(signal.SIGINT, sigint)

(authfd, authfile) = tempfile.mkstemp()
dtrace = subprocess.Popen(["dtrace", "-qs", "/dev/stdin", str(spd_pid)], stdin=subprocess.PIPE, stdout=authfd, stderr=subprocess.PIPE)
print "Exercise the programs to be whitelisted now. Interrupt this script (^C) when you are done."
(stdout, stderr) = dtrace.communicate(input=DSCRIPT)
signal.signal(signal.SIGINT, signal.SIG_DFL)
if stderr:
	fail("dtrace failed: %s" % stderr)
os.lseek(authfd, os.SEEK_SET, 0)	# rewind


#
# Collect all the data into dicts
#
auth = { }
sigs = { }
unsigned = { }
badsigned = { }
errors = { }

file = os.fdopen(authfd, "r")
for line in file:
	(cmd, s, args) = line.strip().partition(";")
	if s != ";":
		continue	# spurious
#	print cmd, "--->", args
	if cmd == "RECORD":
		(type, status, cdhash, path) = args.split(";", 3)
		auth[path] = dict(
			path=path,
			type=type,
			status=status,
			cdhash=cdhash,
			version=2
		)
	elif cmd == "SIGNATURE":
		(type, sigpath, path) = args.split(";", 2)
		with open(sigpath, "r") as sigfile:
			sigdata = sigfile.read()
		sigs[path] = dict(
			path=path,
			type=type,
			signature=plistlib.Data(sigdata)
		)
	elif cmd == "UNSIGNED":
		(type, path) = args.split(";", 1)
		unsigned[path] = dict(
			path=path,
			type=type
		)
	elif cmd == "BROKEN":
		(type, exception, path) = args.split(";", 2)
		badsigned[path] = dict(
			path=path,
			type=type,
			exception=exception
		)

# unsigned code that had a good detached signature recorded is okay
for rec in sigs:
	if rec in unsigned:
		del unsigned[rec]


#
# Pack them up as a single output (plist) file
#
gkedict = dict(
	authority = auth,
	signatures = sigs
)
plistlib.writePlist(gkedict, outputfile)


#
# Report on any problems found
#
for rec in unsigned.values():
	print >>sys.stderr, "PROBLEM: unsigned type %d object not whitelisted: %s" % (rec["type"], rec["path"])
for rec in badsigned.values():
	print >>sys.stderr, "PROBLEM: broken code signature; object not whitelisted: %s" % rec["path"]


#
# Done
#
print "Recorded %d authorization(s), %d signature(s) in %s" % (len(auth), len(sigs), outputfile)
sys.exit(0)