#/usr/local/bin/python

from string import join
from tempfile import mkstemp
from subprocess import call
from threading import Lock

from string import join

from fedid import fedid
from remote_service import service_caller
from service_error import service_error
from util import abac_pem_type, abac_split_cert


import ABAC
import pickle

import sys
import os, os.path
import re

class authorizer_base:
    """
    Classes based on this one keep track of authorization attributes for the
    various modules running.  This base class holds some utility functions that
    they all potentially use.
    """

    # general error exception for badly formed names.  
    class bad_name(RuntimeError): pass
    # difficulty creating an attribute
    class attribute_error(RuntimeError): pass

    @staticmethod
    def auth_name(name):
	"""
	Helper to convert a non-unicode local name to a unicode string.  Mixed
	representations can needlessly confuse the authorizer.
	"""
	if isinstance(name, basestring):
	    if not isinstance(name, unicode): return unicode(name)
	    else: return name
	else: return name

    @staticmethod
    def valid_name(name):
	"""
	Ensure that the given name is valid.  A valid name can either be a
	triple of strings and fedids representing one of our generalized Emulab
	names or a single fedid.  Compound names can include wildcards (None)
	and must anchor to a fedid at their highest level (unless they're all
	None)

	This either returns True or throws an exception.  More an assertion
	than a function.
	"""
	if isinstance(name, tuple) and len(name) == 3:
	    for n in name:
		if n: 
		    if not (isinstance(n, basestring) or isinstance(n, fedid)):
			raise authorizer_base.bad_name(
				"names must be either a triple or a fedid")
	    for n in name: 
		if n:
		    if isinstance(n, fedid):
			return True
		    else:
			raise authorizer_base.bad_name(
				"Compound names must be " + \
				"rooted in fedids: %s" % str(name))

	    return True
	elif isinstance(name, fedid):
	    return True
	else:
	    raise authorizer_base.bad_name(
		    "Names must be a triple or a fedid (%s)" % name)


class authorizer(authorizer_base):
    """
    This class keeps track of authorization attributes for the various modules
    running.  When it gets smarter it will be the basis for a real
    attribute-based authentication system.
    """
    def __init__(self, def_attr="testbed"):
	self.attrs = { }
	self.globals=set()

    def set_attribute(self, name, attr):
	"""
	Attach attr to name.  Multiple attrs can be attached.
	"""
	self.valid_name(name)
	if isinstance(name, tuple):
	    aname = tuple([ self.auth_name(n) for n in name])
	else:
	    aname = self.auth_name(name)

	if not self.attrs.has_key(aname):
	    self.attrs[aname] = set()
	self.attrs[aname].add(attr)

    def unset_attribute(self, name, attr):
	"""
	Remove an attribute from name
	"""
	self.valid_name(name)
	if isinstance(name, tuple):
	    aname = tuple([ self.auth_name(n) for n in name])
	else:
	    aname = self.auth_name(name)

	attrs = self.attrs.get(aname, None)
	if attrs: attrs.discard(attr)

    def check_attribute(self, name, attr):
	"""
	Return True if name has attr (or if attr is global).  Tuple names match
	any tuple name that matches all names present and has None entries in
	other fileds.  For tuple names True implies that there is a matching
	tuple name with the attribute.
	"""
	def tup(tup, i, p):
	    mask = 1 << i
	    if p & mask : return authorizer.auth_name(tup[i])
	    else: return None

	self.valid_name(name)
	if attr in self.globals:
	    return True

	if isinstance(name, tuple):
	    for p in range(0,8):
		lookup = ( tup(name, 0, p), tup(name,1, p), tup(name,2,p))
		if self.attrs.has_key(lookup):
		    if attr in self.attrs[lookup]:
			return True
	else:
	    return  attr in self.attrs.get(self.auth_name(name), set())

    def set_global_attribute(self, attr):
	"""
	Set a global attribute.  All names, even those otherwise unknown to the
	authorizer have this attribute.
	"""
	self.globals.add(attr)

    def unset_global_attribute(self, attr):
	"""
	Remove a global attribute
	"""

	self.globals.discard(attr)

    def import_credentials(self, file_list=None, data_list=None):
	return False

    def __str__(self):
	rv = ""
	rv += "attrs %s\n" % self.attrs
	rv += "globals %s" % self.globals
	return rv

    def clone(self):
	rv = authorizer()
	rv.attrs = self.attrs.copy()
	rv.globals = self.globals.copy()
	return rv

    def save(self, fn):
	f = open(fn, "w")
	pickle.dump(self, f)
	f.close()

    def load(self, fn):
	f = open(fn, "r")
	a = pickle.load(f)
	f.close()
	self.attrs = a.attrs
	self.globals = a.globals
	

class abac_authorizer(authorizer_base):
    """
    Use the ABAC authorization system to make attribute decisions.
    """

    clean_attr_re = re.compile('[^A-Za-z0-9_]+')
    cred_file_re = re.compile('.*\.der$')
    bad_name = authorizer_base.bad_name
    attribute_error = authorizer_base.attribute_error
    class no_file_error(RuntimeError): pass
    class bad_cert_error(RuntimeError): pass

    def __init__(self, certs=None, me=None, key=None, load=None, save=None):
	self.creddy = '/usr/local/bin/creddy'
	self.globals = set()
	self.lock = Lock()
	self.me = me
	self.save_dir = load or save
	if self.save_dir:
	    self.save_dir = os.path.abspath(self.save_dir)
	# If the me parameter is a combination certificate, split it into the
	# abac_authorizer save directory (if any) for use with creddy.
	if self.me is not None and abac_pem_type(self.me) == 'both':
	    if self.save_dir:
		self.key, self.me = abac_split_cert(self.me,
			keyfile="%s/key.pem" % self.save_dir, 
			certfile = "%s/cert.pem" % self.save_dir)
	    else:
		raise abac_authorizer.bad_cert_error("Combination " + \
			"certificate and nowhere to split it");
	else:
	    self.key = key
	self.context = ABAC.Context()
	if me:
	    self.fedid = fedid(file=self.me)
	    rv = self.context.load_id_file(self.me)
	    if rv != 0:
		raise abac_authorizer.bad_name(
			'Cannot load identity from %s' % me.cert)
	else:
	    self.fedid = None

	if isinstance(certs, basestring):
	    certs = [ certs ] 

	for dir in certs or []:
	    self.context.load_directory(dir)

	if load:
	    self.load(load)

    @staticmethod
    def clean_attr(attr):
	return abac_authorizer.clean_attr_re.sub('_', attr)

    def import_credentials(self, file_list=None, data_list=None):
	if data_list:
	    return any([self.import_credential(data=d) for d in data_list])
	elif file_list:
	    return any([self.import_credential(file=f) for f in file_list])
	else:
	    return False

    def import_credential(self, file=None, data=None):
	if data:
	    rv = self.context.load_id_chunk(data)
	    print "id %d" % rv
	    if rv == ABAC.ABAC_CERT_SUCCESS: return True
	    rv = self.context.load_attribute_chunk(data)
	    print "attr %d" % rv
	    return rv == ABAC.ABAC_CERT_SUCCESS
	    #if self.context.load_id_chunk(data) != ABAC.ABAC_CERT_SUCCESS:
	#	return self.context.load_attribute_chunk(data) == \
	#		ABAC.ABAC_CERT_SUCCESS
	    ##else:
	#	return True
	elif file:
	    if self.context.load_id_file(file) != ABAC.ABAC_CERT_SUCCESS:
		return self.context.load_attribute_file(file) == \
			ABAC.ABAC_CERT_SUCCESS
	    else:
		return True
	else:
	    return False

    def set_attribute(self, name=None, attr=None, cert=None):
	if name and attr:
	    if isinstance(name, tuple):
		raise abac_authorizer.bad_name(
			"ABAC doesn't understand three-names")
	    # Convert non-string attributes to strings
	    if not isinstance(attr, basestring):
		attr = "%s" % attr
	    if self.me and self.key:
		# Create a credential and insert it into context
		# This will simplify when we have libcreddy
		try:
		    # create temp file
		    f, fn = mkstemp()
		    os.close(f)
		except EnvironmentError, e:
		    raise abac_authorizer.attribute_error(
			    "Cannot create temp file: %s" %e)

		# Create the attribute certificate with creddy
		cmd = [self.creddy, '--attribute', '--issuer=%s' % self.me, 
		    '--key=%s' % self.key, '--role=%s' % self.clean_attr(attr), 
		    '--subject-id=%s' % name, '--out=%s' % fn]
		rv = call(cmd)
		if rv == 0:
		    self.lock.acquire()
		    # load it to context and remove the file
		    rv = self.context.load_attribute_file(fn)
		    self.lock.release()
		    os.unlink(fn)
		else:
		    os.unlink(fn)
		    raise abac_authorizer.attribute_error(
			    "creddy returned %s" % rv)
	    else:
		raise abac_authorizer.attribute_error(
			"Identity and key not specified on creation")
	elif cert:
	    # Insert this credential into the context
	    self.lock.acquire()
	    self.context.load_attribute_chunk(cert)
	    self.lock.release()
	else:
	    raise abac_authorizer.attribute_error(
		    "Neither name/attr nor cert is set")

    def unset_attribute(self, name, attr):
	if isinstance(name, tuple):
	    raise abac_authorizer.bad_name(
		    "ABAC doesn't understand three-names")
	# Convert non-string attributes to strings
	if not isinstance(attr, basestring):
	    attr = "%s" % attr
	cattr = self.clean_attr(attr)
	self.lock.acquire()
	ctxt = ABAC.Context()
	ids = set()
	for c in self.context.credentials():
	    h = c.head()
	    t = c.tail()
	    if h.is_role() and t.is_principal():
		if t.principal() == '%s' % name and \
			h.principal() == '%s' % self.fedid and \
			h.role_name() == cattr:
		    continue

	    id = c.issuer_cert()
	    if id not in ids:
		ctxt.load_id_chunk(id)
		ids.add(id)
	    ctxt.load_attribute_chunk(c.attribute_cert())
	self.context = ctxt
	self.lock.release()


    def check_attribute(self, name, attr):
	# XXX proof soon
	if isinstance(name, tuple):
	    raise abac_authorizer.bad_name(
		    "ABAC doesn't understand three-names")
	else:
	    # Convert non-string attributes to strings
	    if not isinstance(attr, basestring):
		attr = "%s" % attr
	    # Naked attributes are attested by this principal
	    if attr.find('.') == -1: 
		a = "%s.%s" % (self.fedid, self.clean_attr(attr))
	    else: 
		r, a = attr.split('.',1)
		a = "%s.%s" % ( r, self.clean_attr(a))

	    self.lock.acquire()
	    rv, proof = self.context.query(a, "%s" % name)
	    # XXX delete soon
	    if not rv and attr in self.globals: rv = True
	    self.lock.release()

	    return rv

    def set_global_attribute(self, attr):
	"""
	Set a global attribute.  All names, even those otherwise unknown to the
	authorizer have this attribute.
	"""
	self.lock.acquire()
	self.globals.add(self.clean_attr(attr))
	self.lock.release()

    def unset_global_attribute(self, attr):
	"""
	Remove a global attribute
	"""

	self.lock.acquire()
	self.globals.discard(self.clean_attr(attr))
	self.lock.release()

    def clone(self):
	self.lock.acquire()
	rv = abac_authorizer(me=self.me, key=self.key)
	rv.globals = self.globals.copy()
	rv.context = ABAC.Context(self.context)
	self.lock.release()
	return rv

    def save(self, dir=None):
	self.lock.acquire()
	if dir:
	    self.save_dir = os.path.abspath(dir)
	else:
	    dir = self.save_dir
	if dir is None:
	    self.lock.release()
	    raise abac_authorizer.no_file_error("No load directory specified")
	try:
	    if not os.access(dir, os.F_OK):
		os.mkdir(dir)
	    # These are unpicklable, so set them aside
	    context = self.context
	    lock = self.lock
	    self.context = None
	    self.lock = None

	    f = open("%s/state" % dir, "w")
	    pickle.dump(self, f)
	    f.close()

	    if not os.access("%s/certs" %dir, os.F_OK):
		os.mkdir("%s/certs" % dir)
	    seenid = set()
	    seenattr = set()

	    #restore unpicklable state 
	    self.context = context
	    self.lock = lock
	    #remove old certs
	    for fn in [ f for f in os.listdir("%s/certs" % dir) \
		    if abac_authorizer.cred_file_re.match(f)]:
		os.unlink('%s/certs/%s' % (dir, fn))
	    ii = 0
	    ai = 0
	    for c in self.context.credentials():
		id = c.issuer_cert()
		attr = c.attribute_cert()
		# NB: file naming conventions matter here.  The trailing_ID and
		# _attr are required by ABAC.COntext.load_directory()
		if id and id not in seenid:
		    f = open("%s/certs/ID_%03d_ID.der" % (dir, ii), "w")
		    f.write(id)
		    f.close()
		    ii += 1
		    seenid.add(id)
		if attr and attr not in seenattr:
		    f = open("%s/certs/attr_%03d_attr.der" % (dir, ai), "w")
		    f.write(attr)
		    f.close()
		    ai += 1
		    seenattr.add(attr)
	except EnvironmentError, e:
	    # If we've mislaid self.lock, release lock (they're the same object)
	    if self.lock: self.lock.release()
	    elif lock: lock.release()
	    raise e
	except pickle.PickleError, e:
	    # If we've mislaid self.lock, release lock (they're the same object)
	    if self.lock: self.lock.release()
	    elif lock: lock.release()
	    raise e
	self.lock.release()

    def load(self, dir=None):
	self.lock.acquire()
	if dir:
	    self.save_dir = dir
	else:
	    dir = self.save_dir
	if dir is None:
	    self.lock.release()
	    raise abac_authorizer.no_file_error("No load directory specified")
	try:
	    if os.access("%s/state" % dir, os.R_OK):
		f = open("%s/state" % dir, "r")
		st = pickle.load(f)
		f.close()
		# Copy the useful attributes from the pickled state
		for a in ('globals', 'key', 'me', 'cert', 'fedid'):
		    setattr(self, a, getattr(st, a, None))

	    # Initialize the new context with the new identity
	    self.context = ABAC.Context()
	    if self.me:
		self.context.load_id_file(self.me)
	    self.context.load_directory("%s/certs" % dir)
	    self.save_dir = dir
	except EnvironmentError, e:
	    self.lock.release()
	    raise e
	except pickle.PickleError, e:
	    self.lock.release()
	    raise e
	self.lock.release()

    @staticmethod
    def encode_credential(c):
	return '%s <- %s' % (c.head().string(), c.tail().string())

    def get_creds_for_principal(self, fid):
	look_for = set(["%s" % fid])
	found_attrs = set()
	next_look = set()
	found = set([])

	self.lock.acquire()
	while look_for:
	    for c in self.context.credentials():
		tail = c.tail()
		# XXX: This needs to be more aggressive for linked stuff
		if tail.string() in look_for and c not in found:
		    found.add(c)
		    next_look.add(c.head().string())

	    look_for = next_look
	    next_look = set()
	self.lock.release()
	
	return found

    def __str__(self):

	self.lock.acquire()
	rv = "%s" % self.fedid
	add = join([abac_authorizer.encode_credential(c)
	    for c in self.context.credentials()], '\n');
	if add: rv += "\n%s" % add
	self.lock.release()
	return rv