#!/usr/local/bin/python

import os, sys
import subprocess
import tempfile
import logging
import copy

from M2Crypto import SSL, X509, EVP
from M2Crypto.m2xmlrpclib import SSL_Transport
import M2Crypto.httpslib
from pyasn1.codec.der import decoder

from fedd_services import *
from fedd_internal_services import *
from service_error import *
from xmlrpclib import ServerProxy, dumps, loads, Fault, Error, Binary


# The version of M2Crypto on users is pretty old and doesn't have all the
# features that are useful.  The legacy code is somewhat more brittle than the
# main line, but will work.
if "as_der" not in dir(EVP.PKey):
    from asn1_raw import get_key_bits_from_file, get_key_bits_from_cert
    legacy = True
else:
    legacy = False

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 = password
	    password =lambda *args: pwd

	if password != None:
	    self.load_cert(my_cert, callback=password)
	else:
	    self.load_cert(my_cert)
	if trusted_certs != None: self.load_verify_locations(trusted_certs)
	self.set_verify(SSL.verify_peer, 10)

class fedid:
    """
    Wrapper around the federated ID from an X509 certificate. 
    """
    HASHSIZE=20
    def __init__(self, bits=None, hexstr=None, cert=None, file=None):
	if bits != None:
	    self.set_bits(bits)
	elif hexstr != None:
	    self.set_hexstr(hexstr)
	elif cert != None:
	    self.set_cert(cert)
	elif file != None:
	    self.set_file(file)
	else:
	    self.buf = None

    def __hash__(self):
	return hash(self.buf)

    def __eq__(self, other):
	if isinstance(other, type(self)):
	    return self.buf == other.buf
	elif isinstance(other, type(str())):
	    return self.buf == other;
	else:
	    return False

    def __ne__(self, other): return not self.__eq__(other)

    def __str__(self):
	if self.buf != None:
	    return str().join([ "%02x" % ord(x) for x in self.buf])
	else: return ""

    def __repr__(self):
	return "fedid(hexstr='%s')" % self.__str__()

    def pack_soap(self):
	return self.buf

    def pack_xmlrpc(self):
	return self.buf

    def digest_bits(self, bits):
	"""Internal function.  Compute the fedid from bits and store in buf"""
	d = EVP.MessageDigest('sha1')
	d.update(bits)
	self.buf = d.final()


    def set_hexstr(self, hexstr):
	h = hexstr.replace(':','')
	self.buf= str().join([chr(int(h[i:i+2],16)) \
		for i in range(0,2*fedid.HASHSIZE,2)])

    def get_hexstr(self):
	"""Return the hexstring representation of the fedid"""
	return __str__(self)

    def set_bits(self, bits):
	"""Set the fedid to bits(a 160 bit buffer)"""
	self.buf = bits

    def get_bits(self):
	"""Get the 160 bit buffer from the fedid"""
	return self.buf

    def set_file(self, file):
	"""Get the fedid from a certificate file

	Calculate the SHA1 hash over the bit string of the public key as
	defined in RFC3280.
	"""
	self.buf = None
	if legacy: self.digest_bits(get_key_bits_from_file(file))
	else: self.set_cert(X509.load_cert(file))

    def set_cert(self, cert):
	"""Get the fedid from a certificate.

	Calculate the SHA1 hash over the bit string of the public key as
	defined in RFC3280.
	"""

	self.buf = None
	if (cert != None):
	    if legacy:
		self.digest_bits(get_key_bits_from_cert(cert))
	    else:
		b = []
		k = cert.get_pubkey()

		# Getting the key was easy, but getting the bit string of the
		# key requires a side trip through ASN.1
		dec = decoder.decode(k.as_der())

		# kv is a tuple of the bits in the key.  The loop below
		# recombines these into bytes and then into a buffer for the
		# SSL digest function.
		kv =  dec[0].getComponentByPosition(1)
		for i in range(0, len(kv), 8):
		    v = 0
		    for j in range(0, 8):
			v = (v << 1) + kv[i+j]
		    b.append(v)
		# The comprehension turns b from a list of bytes into a buffer
		# (string) of bytes
		self.digest_bits(str().join([chr(x) for x in b]))

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, type(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"""
    if id.has_key("fedid"): return fedid(id["fedid"])
    else:
	for k in ("localname", "uri", "kerberosUsername"):
	    if id.has_key(k): return id[k]
    return None

def pack_soap(container, name, contents):
    """
    Convert the dictionary in contents into a tree of ZSI classes.

    The holder classes are constructed from factories in container and assigned
    to either the element or attribute name.  This is used to recursively
    create the SOAP message.
    """
    if getattr(contents, "__iter__", None) != None:
	attr =getattr(container, "new_%s" % name, None)
	if attr: obj = attr()
	else:
	    raise TypeError("%s does not have a new_%s attribute" % \
		    (container, name))
	for e, v in contents.iteritems():
	    assign = getattr(obj, "set_element_%s" % e, None) or \
		    getattr(obj, "set_attribute_%s" % e, None)
	    if isinstance(v, type(dict())):
		assign(pack_soap(obj, e, v))
	    elif getattr(v, "__iter__", None) != None:
		assign([ pack_soap(obj, e, val ) for val in v])
	    elif getattr(v, "pack_soap", None) != None:
		assign(v.pack_soap())
	    else:
		assign(v)
	return obj
    else: return contents

def unpack_soap(element):
    """
    Convert a tree of ZSI SOAP classes intro a hash.  The inverse of pack_soap

    Elements or elements that are empty are ignored.
    """
    methods = [ m for m in dir(element) \
	    if m.startswith("get_element") or m.startswith("get_attribute")]
    if len(methods) > 0:
	rv = { }
	for m in methods:
	    if m.startswith("get_element_"): n = m.replace("get_element_","",1)
	    else: n = m.replace("get_attribute_", "", 1)
	    sub = getattr(element, m)()
	    if sub != None:
		if isinstance(sub, basestring):
		    rv[n] = sub
		elif getattr(sub, "__iter__", None) != None:
		    if len(sub) > 0: rv[n] = [unpack_soap(e) for e in sub]
		else:
		    rv[n] = unpack_soap(sub)
	return rv
    else: 
	return element
def apply_to_tags(e, map):
    """
    Map is an iterable of ordered pairs (tuples) that map a key to a function.
    This function walks the given message and replaces any object with a key in
    the map with the result of applying that function to the object.
    """
    dict_type = type(dict())
    list_type = type(list())
    str_type = type(str())

    if isinstance(e, dict_type):
	for k in e.keys():
	    for tag, fcn in map:
		if k == tag:
		    if isinstance(e[k], list_type):
			e[k] = [ fcn(b) for b in e[k]]
		    else:
			e[k] = fcn(e[k])
		elif isinstance(e[k], dict_type):
		    apply_to_tags(e[k], map)
		elif isinstance(e[k], list_type):
		    for ee in e[k]:
			apply_to_tags(ee, map)
    # Other types end the recursion - they should be leaves
    return e

# These are all just specializations of apply_to_tags
def fedids_to_obj(e, tags=('fedid',)):
    """
    Turn the fedids in a message that are encoded as bitstrings into fedid
    objects.
    """
    map = [ (t, lambda x: fedid(bits=x)) for t in tags]
    return apply_to_tags(e, map)

def encapsulate_binaries(e, tags):
    """Walk through a message and encapsulate any dictionary entries in 
    tags into a binary object."""

    def to_binary(o):
	pack = getattr(o, 'pack_xmlrpc', None)
	if callable(pack): return Binary(pack())
	else: return Binary(o)

    map = [ (t, to_binary) for t in tags]
    return apply_to_tags(e, map)

def decapsulate_binaries(e, tags):
    """Walk through a message and encapsulate any dictionary entries in 
    tags into a binary object."""

    map = [ (t, lambda x: x.data) for t in tags]
    return apply_to_tags(e, map)
#end of tag specializations

def strip_unicode(obj):
    """Walk through a message and convert all strings to non-unicode strings"""
    if isinstance(obj, dict):
	for k in obj.keys():
	    obj[k] = strip_unicode(obj[k])
	return obj
    elif isinstance(obj, basestring):
	return str(obj)
    elif getattr(obj, "__iter__", None):
	return [ strip_unicode(x) for x in obj]
    else:
	return obj

def make_unicode(obj):
    """Walk through a message and convert all strings to unicode"""
    if isinstance(obj, dict):
	for k in obj.keys():
	    obj[k] = make_unicode(obj[k])
	return obj
    elif isinstance(obj, basestring):
	return unicode(obj)
    elif getattr(obj, "__iter__", None):
	return [ make_unicode(x) for x in obj]
    else:
	return obj


def generate_fedid(subj, bits=2048, log=None, dir=None, trace=sys.stderr):
    """
    Create a new certificate and derive a fedid from it.

    The fedid and the certificate are returned as a tuple.
    """

    keypath = None
    certpath = None
    try:
	try:
	    kd, keypath = tempfile.mkstemp(dir=dir, prefix="key",
		    suffix=".pem")
	    cd, certpath = tempfile.mkstemp(dir=dir, prefix="cert",
		    suffix=".pem")

	    cmd = ["/usr/bin/openssl", "req", "-text", "-newkey", 
		    "rsa:%d" % bits, "-keyout", keypath,  "-nodes", 
		    "-subj", "/CN=%s" % subj, "-x509", "-days", "30", 
		    "-out", certpath]

	    if log:
		log.debug("[generate_fedid] %s" % " ".join(cmd))

	    if trace: call_out = trace
	    else: 
		call_out = open("/dev/null", "w")

	    rv = subprocess.call(cmd, stdout=call_out, stderr=call_out)
	    log.debug("rv = %d" % rv)
	    if rv == 0:
		cert = ""
		for p in (certpath, keypath):
		    f = open(p)
		    for line in f:
			cert += line
		
		fid = fedid(file=certpath)
		return (fid, cert)
	    else:
		return (None, None)
	except IOError, e:
	    raise e
    finally:
	if keypath: os.remove(keypath)
	if certpath: os.remove(certpath)


class soap_handler:
    """
    Encapsulate the handler code to unpack and pack SOAP requests and responses
    and call the given method.

    The code to decapsulate and encapsulate parameters encoded in SOAP is the
    same modulo a few parameters.  This is a functor that calls a fedd service
    trhough a soap interface.  The parameters are the typecode of the request
    parameters, the method to call (usually a bound instance of a method on a
    fedd service providing class), the constructor of a response packet and the
    name of the body element of that packet.  The handler takes a ParsedSoap
    object (the request) and returns an instance of the class created by
    constructor containing the response.  Failures of the constructor or badly
    created constructors will result in None being returned.
    """
    def __init__(self, typecode, method, constructor, body_name):
	self.typecode = typecode
	self.method = method
	self.constructor = constructor
	self.body_name = body_name

    def __call__(self, ps, fid):
	req = ps.Parse(self.typecode)

	msg = self.method(fedids_to_obj(unpack_soap(req)), fid)

	resp = self.constructor()
	set_element = getattr(resp, "set_element_%s" % self.body_name, None)
	if set_element and callable(set_element):
	    try:
		set_element(pack_soap(resp, self.body_name, msg))
		return resp
	    except (NameError, TypeError):
		return None
	else:
	    return None

class xmlrpc_handler:
    """
    Generate the handler code to unpack and pack XMLRPC requests and responses
    and call the given method.

    The code to marshall and unmarshall XMLRPC parameters to and from a fedd
    service is largely the same.  This helper creates such a handler.  The
    parameters are the method name, and the name of the body struct that
    contains the response.  A handler is created that takes the params response
    from an xmlrpclib.loads on the incoming rpc and a fedid and responds with
    a hash representing the struct ro be returned to the other side.  On error
    None is returned.  Fedid fields are decapsulated from binary and converted
    to fedid objects on input and encapsulated as Binaries on output.
    """
    def __init__(self, method, body_name):
	self.method = method
	self.body_name = body_name
	# A map used by apply_to_tags to convert fedids from xmlrpclib.Binary
	# objects to fedid objects in one sweep.
	self.decap_fedids = (('fedid', lambda x: fedid(bits=x.data)),)

    def __call__(self, params, fid):
	msg = None

	p = apply_to_tags(params[0], self.decap_fedids)
	try:
	    msg = self.method(p, fid)
	except service_error, e:
	    raise Fault(e.code_string(), e.desc)
	if msg != None:
	    return make_unicode(encapsulate_binaries(\
		    { self.body_name: msg }, ('fedid',)))
	else:
	    return None

class service_caller:
    def __init__(self, service_name, port_name, locator, request_message, 
	    request_body_name, tracefile=None):
	self.service_name = service_name
	self.port_name = port_name
	self.locator = locator
	self.request_message = request_message
	self.request_body_name = request_body_name
	self.tracefile = tracefile
	self.__call__ = self.call_service

    def call_xmlrpc_service(self, url, req, cert_file=None, cert_pwd=None, 
	    trusted_certs=None, context=None, tracefile=None):
	"""Send an XMLRPC request.  """
	decap_fedids = (('fedid', lambda x: fedid(bits=x.data)),)


	# If a context is given, use it.  Otherwise construct one from
	# components.  The construction shouldn't call out for passwords.
	if context:
	    ctx = context
	else:
	    try:
		ctx = fedd_ssl_context(cert_file, trusted_certs, 
			password=cert_pwd)
	    except SSL.SSLError:
		raise service_error(service_error.server_config,
			"Certificates misconfigured")

	# Of all the dumbass things.  The XMLRPC library in use here won't
	# properly encode unicode strings, so we make a copy of req with
	# the unicode objects converted.  We also convert the url to a
	# basic string if it isn't one already.
	r = strip_unicode(copy.deepcopy(req))
	url = str(url)
	
	transport = SSL_Transport(ctx)
	port = ServerProxy(url, transport=transport)
	# Make the call, and convert faults back to service_errors
	try:
	    remote_method = getattr(port, self.service_name, None)
	    resp = remote_method(encapsulate_binaries(\
		    { self.request_body_name: r}, ('fedid',)))
	except Fault, f:
	    raise service_error(None, f.faultString, f.faultCode)
	except Error, e:
	    raise service_error(service_error.protocol, 
		    "Remote XMLRPC Fault: %s" % e)

	return apply_to_tags(resp, decap_fedids) 

    def call_soap_service(self, url, req, cert_file=None, cert_pwd=None,
	    trusted_certs=None, context=None, tracefile=None):
	"""
	Send req on to the real destination in dt and return the response

	Req is just the requestType object.  This function re-wraps it.  It
	also rethrows any faults.
	"""

	tf = tracefile or self.tracefile or None

	# If a context is given, use it.  Otherwise construct one from
	# components.  The construction shouldn't call out for passwords.
	if context:
	    ctx = context
	else:
	    try:
		ctx = fedd_ssl_context(cert_file, trusted_certs, 
			password=cert_pwd)
	    except SSL.SSLError:
		raise service_error(service_error.server_config,
			"Certificates misconfigured")
	loc = self.locator()
	get_port = getattr(loc, self.port_name, None)
	if not get_port:
	    raise service_error(service_error.internal, 
		    "Cannot get port %s from locator" % self.port_name)
	port = get_port(url,
		transport=M2Crypto.httpslib.HTTPSConnection, 
		transdict={ 'ssl_context' : ctx },
		tracefile=tf)
	remote_method = getattr(port, self.service_name, None)
	if not remote_method:
	    raise service_error(service_error.internal,
		    "Cannot get service from SOAP port")

	# Reconstruct the full request message
	msg = self.request_message()
	set_element = getattr(msg, "set_element_%s" % self.request_body_name,
		None)
	if not set_element:
	    raise service_error(service_error.internal,
		    "Cannot get element setting method for %s" % \
			    self.request_body_name)
	set_element(pack_soap(msg, self.request_body_name, req))
	try:
	    resp = remote_method(msg)
	except ZSI.ParseException, e:
	    raise service_error(service_error.protocol,
		    "Bad format message (XMLRPC??): %s" % e)
	except ZSI.FaultException, e:
	    ee = unpack_soap(e.fault.detail[0]).get('FeddFaultBody', { })
	    if ee:
		raise service_error(ee['code'], ee['desc'])
	    else:
		raise service_error(service_error.internal,
			"Unexpected fault body")
	r = make_unicode(fedids_to_obj(unpack_soap(resp)))
	return r

    def call_service(self, url, req, cert_file=None, cert_pwd=None, 
	trusted_certs=None, context=None, tracefile=None):
	p_fault = None	# Any SOAP failure (sent unless XMLRPC works)
	resp = None
	try:
	    # Try the SOAP request
	    resp = self.call_soap_service(url, req, 
		    cert_file, cert_pwd, trusted_certs, context, tracefile)
	    return resp
	except service_error, e:
	    if e.code == service_error.protocol: p_fault = None
	    else: raise
	except ZSI.FaultException, f:
	    p_fault = f.fault.detail[0]
		

	# If we could not get a valid SOAP response to the request above,
	# try the same address using XMLRPC and let any faults flow back
	# out.
	if p_fault == None:
	    resp = self.call_xmlrpc_service(url, req, cert_file,
		    cert_pwd, trusted_certs, context, tracefile)
	    return resp
	else:
	    # Build the fault
	    ee = unpack_soap(p_fault).get('FeddFaultBody', { })
	    if ee:
		raise service_error(ee['code'], ee['desc'])
	    else:
		raise service_error(service_error.internal,
			"Unexpected fault body")


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)