#!/usr/local/bin/python

import sys
import os
import pwd

from federation import fedid, service_error
from federation.util import fedd_ssl_context, pack_id, unpack_id
from federation.remote_service import service_caller

from optparse import OptionParser, OptionValueError


class IDFormatException(RuntimeError): pass

class access_method:
    """Encapsulates an access method generically."""
    (type_ssh, type_x509, type_pgp) = ('sshPubkey', 'X509', 'pgpPubkey')
    default_type = type_ssh
    def __init__(self, buf=None, type=None, file=None):
	self.buf = buf

	if type != None: self.type = type
	else: self.type = access_method.default_type

	if file != None:
	    self.readfile(file)
    
    def readfile(self, file, type=None):
	f = open(file, "r")
	self.buf = f.read();
	f.close()
	if type == None:
	    if self.type == None:
		self.type = access_method.default_type
	else:
	    self.type = type;
    
class node_desc:
    def __init__(self, image, hardware, count=1):
	if getattr(image, "__iter__", None) == None:
	    if image == None: self.image = [ ]
	    else: self.image = [ image ]
	else:
	    self.image = image

	if getattr(hardware, "__iter__", None) == None: 
	    if hardware == None: self.hardware = [ ]
	    else: self.hardware = [ hardware ]
	else:
	    self.hardware = hardware
	if count != None: self.count = int(count)
	else: self.count = 1

class fedd_client_opts(OptionParser):
    """Encapsulate option processing in this class, rather than in main"""
    def __init__(self):
	OptionParser.__init__(self, usage="%prog [opts] (--help for details)",
		version="0.1")

	self.add_option("-c","--cert", action="store", dest="cert",
		type="string", help="my certificate file")
	self.add_option("-d", "--debug", action="count", dest="debug", 
		default=0, help="Set debug.  Repeat for more information")
	self.add_option("-s", "--serializeOnly", action="store_true", 
		dest="serialize_only", default=False,
		help="Print the SOAP request that would be sent and exit")
	self.add_option("-T","--trusted", action="store", dest="trusted",
		type="string", help="Trusted certificates (required)")
	self.add_option("-u", "--url", action="store", dest="url",
		type="string",default="https://localhost:23235",  
		help="URL to connect to (default %default)")
	self.add_option("-x","--transport", action="store", type="choice",
		choices=("xmlrpc", "soap"), default="soap",
		help="Transport for request (xmlrpc|soap) (Default: %default)")
	self.add_option("--trace", action="store_const", dest="tracefile", 
		const=sys.stderr, help="Print SOAP exchange to stderr")

class fedd_create_opts(fedd_client_opts):
    def __init__(self, access_keys, add_key_callback=None, 
	    add_cert_callback=None):
	fedd_client_opts.__init__(self)
	self.add_option("-e", "--experiment_cert", dest="out_certfile",
		type="string", help="output certificate file")
	self.add_option("-E", "--experiment_name", dest="exp_name",
		type="string", help="output certificate file")
	self.add_option("-F","--useFedid", action="store_true",
		dest="use_fedid", default=False,
		help="Use a fedid derived from my certificate as user identity")
	self.add_option("-f", "--file", dest="file", 
		help="experiment description file")
	self.add_option("-p", "--project", action="store", dest="project", 
		type="string",
		help="Project to export from master")
	if add_key_callback:
	    self.add_option("-k", "--sshKey", action="callback", type="string", 
		    callback=add_key_callback, callback_args=(access_keys,),
		    help="ssh key for access (can be supplied more than once")
	if add_cert_callback:
	    self.add_option("-K", "--x509Key", action="callback",
		    type="string", callback=add_cert_callback,
		    callback_args=(access_keys,),
		    help="X509 certificate for access " + \
			"(can be supplied more than once")
	self.add_option("-m", "--master", dest="master",
		help="Master testbed in the federation")
	self.add_option("-U", "--username", action="store", dest="user",
		type="string", help="Use this username instead of the uid")

class fedd_split_opts(fedd_create_opts):
    def __init__(self, access_keys, add_key_callback=None, 
	    add_cert_callback=None):
	fedd_create_opts.__init__(self, access_keys, add_key_callback,
		add_cert_callback)
	self.add_option('-t','--fedkit', action='store_true', dest='fedkit',
		default=False,
		help="get output suitable for federation kit install")


class fedd_access_opts(fedd_create_opts):
    def __init__(self, access_keys, node_descs, add_key_callback=None, 
	    add_cert_callback=None, add_node_callback=None):
	fedd_create_opts.__init__(self, access_keys, add_key_callback,
		add_cert_callback)
	self.add_option("-a","--anonymous", action="store_true",
		dest="anonymous", default=False,
		help="Do not include a user in the request")
	self.add_option("-l","--label", action="store", dest="label",
		type="string", help="Label for output")
	if add_node_callback:
	    self.add_option("-n", "--node", action="callback", type="string", 
		    callback=add_node_callback, callback_args=(node_descs,),
		    help="Node description: image:hardware[:count]")
	self.add_option("-t", "--testbed", action="store", dest="testbed",
		type="string",
		help="Testbed identifier (URI) to contact (required)")

class fedd_exp_data_opts(fedd_client_opts):
    def __init__(self):
	fedd_client_opts.__init__(self)
	self.add_option("-e", "--experiment_cert", dest="exp_certfile",
		type="string", help="output certificate file")
	self.add_option("-E", "--experiment_name", dest="exp_name",
		type="string", help="output certificate file")

def exit_with_fault(dict, out=sys.stderr):
    """ Print an error message and exit.

    The dictionary contains the FeddFaultBody elements."""
    codestr = ""

    if dict.has_key('errstr'):
	codestr = "Error: %s" % dict['errstr']

    if dict.has_key('code'):
	if len(codestr) > 0 : 
	    codestr += " (%d)" % dict['code']
	else:
	    codestr = "Error Code: %d" % dict['code']

    print>>out, codestr
    print>>out, "Description: %s" % dict['desc']
    sys.exit(dict.get('code', 20))
# Base class that will do a the SOAP/XMLRPC exchange for a request.
class fedd_rpc:
    class RPCException:
	def __init__(self, fb):
	    self.desc = fb.get('desc', None)
	    self.code = fb.get('code', -1)
	    self.errstr = fb.get('errstr', None)

    def __init__(self, pre): 
	"""
	Specialize the class for the pre method
	"""
	self.RequestBody="%sRequestBody" % pre
	self.ResponseBody="%sResponseBody" % pre
	self.method = pre

	self.caller = service_caller(self.method)
	self.RPCException = fedd_rpc.RPCException


    def add_ssh_key(self, option, opt_str, value, parser, access_keys):
	try:
	    access_keys.append(access_method(file=value,
		type=access_method.type_ssh))
	except IOError, (errno, strerror):
	    raise OptionValueError("Cannot generate sshPubkey from %s: "\
		    "%s (%d)" % (value,strerror,errno))

    def add_x509_cert(self, option, opt_str, value, parser, access_keys):
	try:
	    access_keys.append(access_method(file=value,
		type=access_method.type_x509))
	except IOError, (errno, strerror):
	    raise OptionValueError("Cannot read x509 cert from %s: %s (%d)" %
		    (value,strerror,errno))
    def add_node_desc(self, option, opt_str, value, parser, node_descs):
	def none_if_zero(x):
	    if len(x) > 0: return x
	    else: return None

	params = map(none_if_zero, value.split(":"));
	
	if len(params) < 4 and len(params) > 1:
	    node_descs.append(node_desc(*params))
	else:
	    raise OptionValueError("Bad node description: %s" % value)

    def get_user_info(self, access_keys):
	pw = pwd.getpwuid(os.getuid());
	try_cert=None
	user = None

	if pw != None:
	    user = pw[0]
	    try_cert = "%s/.ssl/emulab.pem" % pw[5];
	    if not os.access(try_cert, os.R_OK):
		try_cert = None
	    if len(access_keys) == 0:
		for k in ["%s/.ssh/id_rsa.pub", "%s/.ssh/id_dsa.pub", 
			"%s/.ssh/identity.pub"]:
		    try_key = k % pw[5];
		    if os.access(try_key, os.R_OK):
			access_keys.append(access_method(file=try_key,
			    type=access_method.type_ssh))
			break
	return (user, try_cert)

    def do_rpc(self, req_dict, url, transport, cert, trusted, tracefile=None,
	    serialize_only=False):
	"""
	The work of sending and parsing the RPC as either XMLRPC or SOAP
	"""

	context = None
	while context == None:
	    try:
		context = fedd_ssl_context(cert, trusted)
	    except Exception, e:
		# Yes, doing this on message type is not ideal.  The string
		# comes from OpenSSL, so check there is this stops working.
		if str(e) == "bad decrypt": 
		    print >>sys.stderr, "Bad Passphrase given."
		else: raise

	if transport == "soap":
	    if serialize_only:
		print self.caller.serialize_soap(req_dict) 
		sys.exit(0)
	    else:
		try:
		    resp = self.caller.call_soap_service(url, req_dict, 
			    context=context, tracefile=tracefile)
		except service_error, e:
		    raise self.RPCException( {\
			    'code': e.code, 
			    'desc': e.desc, 
			    'errstr': e.code_string()\
			})
	elif transport == "xmlrpc":
	    if serialize_only:
		ser = dumps((req_dict,))
		print ser
		sys.exit(0)
	    else:
		try:
		    resp = self.caller.call_xmlrpc_service(url, req_dict, 
			    context=context, tracefile=tracefile)
		except service_error, e:
		    raise self.RPCException( {\
			    'code': e.code, 
			    'desc': e.desc, 
			    'errstr': e.code_string()\
			})

	else:
	    raise RuntimeError("Unknown RPC transport: %s" % transport)

	if resp.has_key(self.ResponseBody):
	    return resp[self.ResponseBody]
	else:
	    raise RuntimeError("No body in response??")

# Querying experiment data follows the same control flow regardless of the
# specific data retrieved.  This class encapsulates that control flow.
class exp_data(fedd_rpc):
    def __init__(self, op): 
	"""
	Specialize the class for the type of data requested (op)
	"""

	fedd_rpc.__init__(self, op)
	if op =='Vtopo':
	    self.key="vtopo"
	    self.xml='experiment'
	elif op == 'Vis':
	    self.key="vis"
	    self.xml='vis'
	elif op == 'Info': pass
	else:
	    raise TypeError("Bad op: %s" % op)

    def print_xml(self, d, out=sys.stdout):
	"""
	Print the retrieved data is a simple xml representation of the dict.
	"""
	str = "<%s>\n" % self.xml
	for t in ('node', 'lan'):
	    if d.has_key(t): 
		for x in d[t]:
		    str += "<%s>" % t
		    for k in x.keys():
			str += "<%s>%s</%s>" % (k, x[k],k)
		    str += "</%s>\n" % t
	str+= "</%s>" % self.xml
	print >>out, str

    def __call__(self):
	"""
	The control flow.  Compose the request and print the response.
	"""
	# Process the options using the customized option parser defined above
	parser = fedd_exp_data_opts()

	(opts, args) = parser.parse_args()

	if opts.trusted:
	    if ( not os.access(opts.trusted, os.R_OK) ) :
		sys.exit("Cannot read trusted certificates (%s)" % opts.trusted)

	if opts.debug > 0: opts.tracefile=sys.stderr

	(user, cert) = self.get_user_info([])

	if opts.cert != None: cert = opts.cert

	if cert == None:
	    sys.exit("No certificate given (--cert) or found")

	if os.access(cert, os.R_OK):
	    fid = fedid(file=cert)
	else:
	    sys.exit("Cannot read certificate (%s)" % cert)

	if opts.exp_name and opts.exp_certfile:
	    sys.exit("Only one of --experiment_cert and " +\
		    "--experiment_name permitted");

	if opts.exp_certfile:
	    exp_id = { 'fedid': fedid(file=opts.exp_certfile) }

	if opts.exp_name:
	    exp_id = { 'localname' : opts.exp_name }

	req = { 'experiment': exp_id }

	try:
	    resp_dict = self.do_rpc(req,
		    opts.url, opts.transport, cert, opts.trusted, 
		    serialize_only=opts.serialize_only,
		    tracefile=opts.tracefile)
	except self.RPCException, e:
	    exit_with_fault(\
		    {'desc': e.desc, 'errstr': e.errstr, 'code': e.code})
	except RuntimeError, e:
	    print e
	    sys.exit("Error processing RPC: %s" % e)

	if getattr(self, 'key', None):
	    try:
		if resp_dict.has_key(self.key):
		    self.print_xml(resp_dict[self.key])
	    except RuntimeError, e:
		sys.exit("Bad response. %s" % e.message)
	else:
	    print resp_dict

class terminate(fedd_rpc):
    def __init__(self): 
	"""
	Termination request
	"""

	fedd_rpc.__init__(self, "Terminate")

    def __call__(self):
	"""
	The control flow.  Compose the request and print the response.
	"""
	# Process the options using the customized option parser defined above
	parser = fedd_exp_data_opts()

	(opts, args) = parser.parse_args()

	(user, cert) = self.get_user_info([])
	if opts.trusted:
	    if ( not os.access(opts.trusted, os.R_OK) ) :
		sys.exit("Cannot read trusted certificates (%s)" % opts.trusted)

	if opts.debug > 0: opts.tracefile=sys.stderr

	if opts.cert != None: cert = opts.cert

	if cert == None:
	    sys.exit("No certificate given (--cert) or found")

	if os.access(cert, os.R_OK):
	    fid = fedid(file=cert)
	else:
	    sys.exit("Cannot read certificate (%s)" % cert)

	if opts.exp_name and opts.exp_certfile:
	    sys.exit("Only one of --experiment_cert and " +\
		    "--experiment_name permitted");

	exp_id = None

	if opts.exp_certfile:
	    exp_id = { 'fedid': fedid(file=opts.exp_certfile) }

	if opts.exp_name:
	    exp_id = { 'localname' : opts.exp_name }

	if not exp_id:
	    sys.exit("Must give one of --experiment_cert and " +\
		    "--experiment_name");

	req = { 'experiment': exp_id }

	try:
	    resp_dict = self.do_rpc(req,
		    opts.url, opts.transport, cert, opts.trusted, 
		    serialize_only=opts.serialize_only,
		    tracefile=opts.tracefile)
	except self.RPCException, e:
	    exit_with_fault(\
		    {'desc': e.desc, 'errstr': e.errstr, 'code': e.code})
	except RuntimeError, e:
	    print e
	    sys.exit("Error processing RPC: %s" % e)

	eid = resp_dict.get('experimentID', None)
	if eid:
	    for id in eid:
		for k in id.keys():
		    if k == 'fedid': print "%s: %s" % (k,fedid(bits=id[k]))
		    else: print "%s: %s" % (k, id[k])

class create(fedd_rpc):
    def __init__(self): 
	fedd_rpc.__init__(self, "Create")
    def __call__(self):
	access_keys = []
	# Process the options using the customized option parser defined above
	parser = fedd_create_opts(access_keys, self.add_ssh_key,
		self.add_x509_cert)

	(opts, args) = parser.parse_args()

	if opts.trusted:
	    if ( not os.access(opts.trusted, os.R_OK) ) :
		sys.exit("Cannot read trusted certificates (%s)" % opts.trusted)

	if not opts.project :
	    parser.error('--project is required')

	if opts.debug > 0: opts.tracefile=sys.stderr

	(user, cert) = self.get_user_info(access_keys)

	if opts.user: user = opts.user

	if opts.cert != None: cert = opts.cert

	if cert == None:
	    sys.exit("No certificate given (--cert) or found")

	if os.access(cert, os.R_OK):
	    fid = fedid(file=cert)
	    if opts.use_fedid == True:
		user = fid
	else:
	    sys.exit("Cannot read certificate (%s)" % cert)

	if opts.file:
	    exp_desc = ""
	    try:
		f = open(opts.file, 'r')
		for line in f:
		    exp_desc += line
		f.close()
	    except IOError:
		sys.exit("Cannot read description file (%s)" %opts.file)
	else:
	    sys.exit("Must specify an experiment description (--file)")

	if not opts.master:
	    sys.exit("Must specify a master testbed (--master)")

	out_certfile = opts.out_certfile

	msg = {
		'experimentdescription': { 'ns2description': exp_desc },
		'master': opts.master,
		'exportProject': { 'localname': opts.project },
		'user' : [ {\
			'userID': pack_id(user), \
			'access': [ { a.type: a.buf } for a in access_keys]\
			} ]
		}

	if opts.exp_name:
	    msg['experimentID'] = { 'localname': opts.exp_name }

	if opts.debug > 1: print >>sys.stderr, msg

	try:
	    resp_dict = self.do_rpc(msg, 
		    opts.url, opts.transport, cert, opts.trusted, 
		    serialize_only=opts.serialize_only,
		    tracefile=opts.tracefile)
	except self.RPCException, e:
	    exit_with_fault(\
		    {'desc': e.desc, 'errstr': e.errstr, 'code': e.code})
	except RuntimeError, e:
	    sys.exit("Error processing RPC: %s" % e)

	if opts.debug > 1: print >>sys.stderr, resp_dict

	ea = resp_dict.get('experimentAccess', None)
	if out_certfile and ea and ea.has_key('X509'):
	    try:
		f = open(out_certfile, "w")
		print >>f, ea['X509']
		f.close()
	    except IOError:
		sys.exit('Could not write to %s' %  out_certfile)
	eid = resp_dict.get('experimentID', None)
	if eid:
	    for id in eid:
		for k in id.keys():
		    print "%s: %s" % (k, id[k])

class split(fedd_rpc):
    def __init__(self): 
	fedd_rpc.__init__(self, "Ns2Split")
    def __call__(self):
	access_keys = []
	# Process the options using the customized option parser defined above
	parser = fedd_split_opts(access_keys, self.add_ssh_key,
		self.add_x509_cert)

	(opts, args) = parser.parse_args()

	if opts.trusted:
	    if ( not os.access(opts.trusted, os.R_OK) ) :
		sys.exit("Cannot read trusted certificates (%s)" % opts.trusted)

	if opts.debug > 0: opts.tracefile=sys.stderr

	if opts.cert != None: cert = opts.cert
	else: cert = None

	if cert == None:
	    sys.exit("No certificate given (--cert) or found")

	if os.access(cert, os.R_OK):
	    fid = fedid(file=cert)
	    if opts.use_fedid == True:
		user = fid
	else:
	    sys.exit("Cannot read certificate (%s)" % cert)

	if opts.file:
	    exp_desc = ""
	    try:
		f = open(opts.file, 'r')
		for line in f:
		    exp_desc += line
		f.close()
	    except IOError:
		sys.exit("Cannot read description file (%s)" %opts.file)
	else:
	    sys.exit("Must specify an experiment description (--file)")

	if not opts.master:
	    sys.exit("Must specify a master testbed (--master)")

	out_certfile = opts.out_certfile

	msg = {
		'description': { 'ns2description': exp_desc },
		'master': opts.master,
		'include_fedkit': opts.fedkit,
		}

	if opts.debug > 1: print >>sys.stderr, msg

	try:
	    resp_dict = self.do_rpc(msg, 
		    opts.url, opts.transport, cert, opts.trusted, 
		    serialize_only=opts.serialize_only,
		    tracefile=opts.tracefile)
	except self.RPCException, e:
	    exit_with_fault(\
		    {'desc': e.desc, 'errstr': e.errstr, 'code': e.code})
	except RuntimeError, e:
	    sys.exit("Error processing RPC: %s" % e)

	if opts.debug > 1: print >>sys.stderr, resp_dict

	out = resp_dict.get('output', None)

	for line in out.splitlines():
	    print "%s" % line

class access(fedd_rpc):
    def __init__(self):
	fedd_rpc.__init__(self, "RequestAccess")

    def print_response_as_testbed(self, resp, label, out=sys.stdout):
	"""Print the response as input to the splitter script"""

	e = resp['emulab']
	p = e['project']
	fields = { 
		"Boss": e['boss'],
		"OpsNode": e['ops'],
		"Domain": e['domain'],
		"FileServer": e['fileServer'],
		"EventServer": e['eventServer'],
		"Project": unpack_id(p['name'])
		}
	if (label != None): print >> out, "[%s]" % label

	for l, v in fields.iteritems():
	    print >>out, "%s: %s" % (l, v)

	for u in p['user']:
	    print >>out, "User: %s" % unpack_id(u['userID'])

	for a in e['fedAttr']:
	    print >>out, "%s: %s" % (a['attribute'], a['value'])

    def __call__(self):
	access_keys = []
	node_descs = []
	proj = None

	# Process the options using the customized option parser defined above
	parser = fedd_access_opts(access_keys, node_descs, self.add_ssh_key,
		self.add_x509_cert, self.add_node_desc)

	(opts, args) = parser.parse_args()

	if opts.testbed == None:
	    parser.error("--testbed is required")

	if opts.trusted:
	    if ( not os.access(opts.trusted, os.R_OK) ) :
		sys.exit("Cannot read trusted certificates (%s)" % opts.trusted)

	if opts.debug > 0: opts.tracefile=sys.stderr

	(user, cert) = self.get_user_info(access_keys)

	if opts.user: user = opts.user

	if opts.cert != None: cert = opts.cert

	if cert == None:
	    sys.exit("No certificate given (--cert) or found")

	if os.access(cert, os.R_OK):
	    fid = fedid(file=cert)
	    if opts.use_fedid == True:
		user = fid
	else:
	    sys.exit("Cannot read certificate (%s)" % cert)

	msg = {
		'allocID': pack_id('test alloc'),
		'destinationTestbed': pack_id(opts.testbed),
		'serviceAccess' : [ { a.type: a.buf } for a in access_keys ],
		'createAccess' : [ { a.type: a.buf } for a in access_keys ],
		}

	if len(node_descs) > 0:
	    msg['resources'] = { 
		    'node': [ 
			{ 
			    'image':  n.image ,
			    'hardware':  n.hardware,
			    'count': n.count,
			} for n in node_descs],
		    }

	if opts.project != None:
	    if not opts.anonymous and user != None:
		msg['project'] = {
			'name': pack_id(opts.project),
			'user': [ { 'userID': pack_id(user) } ],
			}
	    else:
		msg['project'] = { 'name': pack_id(opts.project) }
	else:
	    if not opts.anonymous and user != None:
		msg['user'] = [ { 'userID': pack_id(user) } ]
	    else:
		msg['user'] = [];

	if opts.debug > 1: print >>sys.stderr, msg

	try:
	    resp_dict = self.do_rpc(msg, 
		    opts.url, opts.transport, cert, opts.trusted, 
		    serialize_only=opts.serialize_only,
		    tracefile=opts.tracefile)
	except self.RPCException, e:
	    exit_with_fault(\
		    {'desc': e.desc, 'errstr': e.errstr, 'code': e.code})
	except RuntimeError, e:
	    sys.exit("Error processing RPC: %s" % e.message)

	if opts.debug > 1: print >>sys.stderr, resp_dict
	self.print_response_as_testbed(resp_dict, opts.label)

cmds = {\
	'create': create(),\
	'split': split(),\
	'access': access(),\
	'vtopo': exp_data('Vtopo'),\
	'vis': exp_data('Vis'),\
	'info': exp_data('Info'),\
	'terminate': terminate(),\
    }

operation = cmds.get(sys.argv[1], None)
if operation:
    del sys.argv[1]
    operation()
else:
    sys.exit("Bad command: %s.  Valid ones are: %s" % \
	    (sys.argv[1], ", ".join(cmds.keys())))