#!/usr/local/bin/python

import re
import os
import string
import logging
import pickle

import httplib

from optparse import OptionParser

from socket import sslerror
from tempfile import mkstemp

from M2Crypto import SSL
from M2Crypto.SSL import SSLError
from fedid import fedid
from service_error import service_error
from urlparse import urlparse
from M2Crypto import m2


# If this is an old enough version of M2Crypto.SSL that has an
# ssl_verify_callback that doesn't allow 0-length signed certs, create a
# version of that callback that does.  This is edited from the original in
# M2Crypto.SSL.cb.  This version also elides the printing to stderr.
if not getattr(SSL.cb, 'ssl_verify_callback_allow_unknown_ca', None):
    from M2Crypto.SSL.Context import map

    def fedd_ssl_verify_callback(ssl_ctx_ptr, x509_ptr, errnum, errdepth, ok):
	unknown_issuer = [
	    m2.X509_V_ERR_UNABLE_TO_GET_ISSUER_CERT_LOCALLY,
	    m2.X509_V_ERR_UNABLE_TO_VERIFY_LEAF_SIGNATURE,
	    m2.X509_V_ERR_CERT_UNTRUSTED,
	    m2.X509_V_ERR_DEPTH_ZERO_SELF_SIGNED_CERT
	    ]
	# m2.X509_V_ERR_SELF_SIGNED_CERT_IN_CHAIN should also be allowed
	if getattr(m2, 'X509_V_ERR_SELF_SIGNED_CERT_IN_CHAIN', None):
	    unknown_issuer.append(getattr(m2, 
		'X509_V_ERR_SELF_SIGNED_CERT_IN_CHAIN', None))
	ssl_ctx = map()[ssl_ctx_ptr]

	if errnum in unknown_issuer: 
	    if ssl_ctx.get_allow_unknown_ca():
		ok = 1
	# CRL checking goes here...
	if ok:
	    if ssl_ctx.get_verify_depth() >= errdepth:
		ok = 1
	    else:
		ok = 0
	return ok
else:
    def fedd_ssl_verify_callback(ok, store):
	'''
	m2.X509_V_ERR_SELF_SIGNED_CERT_IN_CHAIN should also be allowed
	'''
	errnum = store.get_error()
	if errnum == m2.X509_V_ERR_SELF_SIGNED_CERT_IN_CHAIN:
	    ok = 1
	    return ok
	else:
	    return SSL.cb.ssl_verify_callback_allow_unknown_ca(ok, store)

class fedd_ssl_context(SSL.Context):
    """
    Simple wrapper around an M2Crypto.SSL.Context to initialize it for fedd.
    """
    def __init__(self, my_cert, trusted_certs=None, password=None):
	"""
	construct a fedd_ssl_context

	@param my_cert: PEM file with my certificate in it
	@param trusted_certs: PEM file with trusted certs in it (optional)
	"""
	SSL.Context.__init__(self)

	# load_cert takes a callback to get a password, not a password, so if
	# the caller provided a password, this creates a nonce callback using a
	# lambda form.
	if password != None and not callable(password):
	    # This is cute.  password = lambda *args: password produces a
	    # function object that returns itself rather than one that returns
	    # the object itself.  This is because password is an object
	    # reference and after the assignment it's a lambda.  So we assign
	    # to a temp.
	    pwd = str(password)
	    password =lambda *args: pwd

	# The calls to str below (and above) are because the underlying SSL
	# stuff is intolerant of unicode.
	if password != None:
	    self.load_cert(str(my_cert), callback=password)
	else:
	    self.load_cert(str(my_cert))

	# If no trusted certificates are specified, allow unknown CAs.
	if trusted_certs: 
	    self.load_verify_locations(trusted_certs)
	    self.set_verify(SSL.verify_peer, 10)
	else:
	    # Install the proper callback to allow self-signed certs
	    self.set_allow_unknown_ca(True)
	    self.set_verify(SSL.verify_peer, 10, 
		    callback=fedd_ssl_verify_callback)

class file_expanding_opts(OptionParser):
    def expand_file(self, option, opt_str, v, p):
	"""
	Store the given value to the given destination after expanding home
	directories.
	"""
	setattr(p.values, option.dest, os.path.expanduser(v))

    def __init__(self, usage=None, version=None):
	OptionParser.__init__(self)


def read_simple_accessdb(fn, auth, mask=[]):
    """
    Read a simple access database.  Each line is a fedid (of the form
    fedid:hexstring) and a comma separated list of atributes to be assigned to
    it.  This parses out the fedids and adds the attributes to the authorizer.
    comments (preceded with a #) and blank lines are ignored.  Exceptions (e.g.
    file exceptions and ValueErrors from badly parsed lines) are propagated.
    """

    rv = [ ]
    lineno = 0
    fedid_line = re.compile("fedid:([" + string.hexdigits + "]+)\s+" +\
	    "(\w+\s*(,\s*\w+)*)")

    # If a single string came in, make it a list
    if isinstance(mask, basestring): mask = [ mask ]

    f = open(fn, 'r')
    for line in f:
	lineno += 1
	line = line.strip()
	if line.startswith('#') or len(line) == 0: 
	    continue
	m = fedid_line.match(line)
	if m :
	    fid = fedid(hexstr=m.group(1))
	    for a in [ a.strip() for a in m.group(2).split(",") \
		    if not mask or a.strip() in mask ]:
		auth.set_attribute(fid, a.strip())
	else:
	    raise ValueError("Badly formatted line in accessdb: %s line %d" %\
		    (fn, lineno))
    f.close()
    return rv
	

def pack_id(id):
    """
    Return a dictionary with the field name set by the id type.  Handy for
    creating dictionaries to be converted to messages.
    """
    if isinstance(id, fedid): return { 'fedid': id }
    elif id.startswith("http:") or id.startswith("https:"): return { 'uri': id }
    else: return { 'localname': id}

def unpack_id(id):
    """return id as a type determined by the key"""
    for k in ("localname", "fedid", "uri", "kerberosUsername"):
	if id.has_key(k): return id[k]
    return None

def set_log_level(config, sect, log):
    """ Set the logging level to the value passed in sect of config."""
    # The getattr sleight of hand finds the logging level constant
    # corrersponding to the string.  We're a little paranoid to avoid user
    # mayhem.
    if config.has_option(sect, "log_level"):
	level_str = config.get(sect, "log_level")
	try:
	    level = int(getattr(logging, level_str.upper(), -1))

	    if  logging.DEBUG <= level <= logging.CRITICAL:
		log.setLevel(level)
	    else:
		log.error("Bad experiment_log value: %s" % level_str)

	except ValueError:
	    log.error("Bad experiment_log value: %s" % level_str)

def copy_file(src, dest, size=1024):
    """
    Exceedingly simple file copy.  Throws an EnvironmentError if there's a problem.
    """
    s = open(src,'r')
    d = open(dest, 'w')

    buf = s.read(size)
    while buf != "":
	d.write(buf)
	buf = s.read(size)
    s.close()
    d.close()

def get_url(url, cf, destdir, fn=None, max_retries=5):
    """
    Get data from a federated data store.  This presents the client cert/fedid
    to the http server.  We retry up to max_retries times.
    """
    po = urlparse(url)
    if not fn:
	fn = po.path.rpartition('/')[2]
    retries = 0
    ok = False
    failed_exception = None
    while not ok and retries < 5:
	try:
	    conn = httplib.HTTPSConnection(po.hostname, port=po.port, 
		    cert_file=cf, key_file=cf)
	    conn.putrequest('GET', po.path)
	    conn.endheaders()
	    response = conn.getresponse()

	    lf = open("%s/%s" % (destdir, fn), "w")
	    buf = response.read(4096)
	    while buf:
		lf.write(buf)
		buf = response.read(4096)
	    lf.close()
	    ok = True
	except EnvironmentError, e:
	    failed_excpetion = e
	    retries += 1
	except httplib.HTTPException, e:
	    failed_exception = e
	    retries += 1
	except sslerror, e:
	    failed_exception = e
	    retries += 1
	except SSLError, e:
	    failed_exception = e
	    retries += 1

    if retries > 5 and failed_exception:
	raise failed_excpetion

# Functions to manipulate composite testbed names
def testbed_base(tb):
    """
    Simple code to get the base testebd name.
    """
    i = tb.find('/')
    if i == -1: return tb
    else: return tb[0:i]

def testbed_suffix(tb):
    """
    Simple code to get a testbed suffix, if nay.  No suffix returns None.
    """
    i = tb.find('/')
    if i != -1: return tb[i+1:]
    else: return None

def split_testbed(tb):
    """
    Return a testbed and a suffix as a tuple.  No suffix returns None for that
    field
    """

    i = tb.find('/')
    if i != -1: return (tb[0:i], tb[i+1:])
    else: return (tb, None)

def join_testbed(base, suffix=None):
    """
    Build a testbed with suffix.  If base is itself a tuple, combine them,
    otherwise combine the two.
    """
    if isinstance(base, tuple):
	if len(base) == 2:
	    return '/'.join(base)
	else:
	    raise RuntimeError("Too many tuple elements for join_testbed")
    else:
	if suffix:
	    return '/'.join((base, suffix))
	else:
	    return base

def abac_pem_type(cert):
    key_re = re.compile('\s*-----BEGIN RSA PRIVATE KEY-----$')
    cert_re = re.compile('\s*-----BEGIN CERTIFICATE-----$') 
    type = None

    f = open(cert, 'r')
    for line in f:
	if key_re.match(line):
	    if type is None: type = 'key'
	    elif type == 'cert': type = 'both'
	elif cert_re.match(line):
	    if type is None: type = 'cert'
	    elif type == 'key': type = 'both'
	if type == 'both': break
    f.close()
    return type

def abac_split_cert(cert, keyfile=None, certfile=None):
    """
    Split the certificate file in cert into a certificate file and a key file
    in cf and kf respectively.  The ABAC tools generally cannot handle combined
    certificates/keys.  If kf anc cf are given, they are used, otherwise tmp
    files are created.  Created tmp files must be deleted.  Problems opening or
    writing files will cause exceptions.
    """
    class diversion:
	'''
	Wraps up the reqular expression to start and end a diversion, as well as
	the open file that gets the lines.  If fd is passed in, use that system
	file (probably from a mkstemp.  Otherwise open the given filename.
	'''
	def __init__(self, start, end, fn=None, fd=None):
	    self.start = re.compile(start)
	    self.end = re.compile(end)

	    if not fd:
		# Open the file securely with minimal permissions. NB file
		# cannot exist before this call.
		fd = os.open(fn,
		    (os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_EXCL), 0600)

	    self.f = os.fdopen(fd, 'w')

	def close(self):
	    self.f.close()

    if not keyfile:
	kf, rkeyfile = mkstemp(suffix=".pem")
    else:
	kf, rkeyfile = None, keyfile

    if not certfile:
	cf, rcertfile = mkstemp(suffix=".pem")
    else:
	cf, rcertfile = None, certfile

    # Initialize the diversions
    divs = [diversion(s, e, fn=pfn, fd=pfd ) for s, e, pfn, pfd in (
	('\s*-----BEGIN RSA PRIVATE KEY-----$', 
	    '\s*-----END RSA PRIVATE KEY-----$',
	    keyfile, kf),
	('\s*-----BEGIN CERTIFICATE-----$', 
	    '\s*-----END CERTIFICATE-----$',
	    certfile, cf))]

    # walk through the file, beginning a diversion when a start regexp
    # matches until the end regexp matches.  While in the two regexps,
    # print each line to the open diversion file (including the two
    # matches).
    active = None
    f = open(cert, 'r')
    for l in f:
	if active:
	    if active.end.match(l):
		print >>active.f, l,
		active = None
	else:
	    for d in divs:
		if d.start.match(l):
		    active = d
		    break
	if active: print >>active.f, l,

    # This is probably unnecessary.  Close all the diversion files.
    for d in divs: d.close()
    return rkeyfile, rcertfile

def abac_context_to_creds(context):
    """
    Pull all the credentials out of the context and return 2  lists of the
    underlying credentials in an exportable format, IDs and attributes.
    There are no duplicates in the lists.
    """
    ids, attrs = set(), set()
    # This should be a one-iteration loop
    for c in context.credentials():
	ids.add(str(c.issuer_cert()))
	attrs.add(str(c.attribute_cert()))

    return list(ids), list(attrs)

def find_pickle_problem(o, st=None):
    """
    Debugging routine to figure out what doesn't pickle from a dict full of
    dicts and lists.  It tries to walk down the lists and dicts and pickle each
    atom.  If something fails to pickle, it prints an approximation of a stack
    trace through the data structure.
    """
    if st is None: st = [ ]

    if isinstance(o, dict):
	for k, i in o.items():
	    st.append(k)
	    find_pickle_problem(i, st)
	    st.pop()
    elif isinstance(o, list):
	st.append('list')
	for i in o:
	    find_pickle_problem(i, st)
	st.pop()
    else:
	try:
	    pickle.dumps(o)
	except pickle.PicklingError, e:
	    print >>sys.stderr, "<PicklingError>"
	    print >>sys.stderr, st
	    print >>sys.stderr, o
	    print >>sys.stderr, e
	    print >>sys.stderr, "</PicklingError>"