#!/usr/local/bin/python

import sys, os
import re
import os.path

import tempfile
import subprocess
import logging 
import time
import signal

import util

import xmlrpclib

import emulab_segment


class start_segment(emulab_segment.start_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None,
	    ops=None, cert=None):
	emulab_segment.start_segment.__init__(self, log=log, keyfile=keyfile,
		debug=debug, boss=boss, ops=ops, cert=cert)

    def __call__(self, parent, eid, pid, user, tclfile, tmpdir, timeout=0, 
	    gid=None):
	"""
	Start a sub-experiment on a federant.

	Get the current state, and terminate the experiment if it exists. The
	group membership of the experiment is difficult to determine or change,
	so start with a clean slate.  Create a new one and ship data
	and configs and start the experiment.  There are small ordering
	differences based on the initial state of the sub-experiment.
	"""

	# exports_segment has established a clean dummy experiment.

	if not self.set_up_experiment_filespace(user, self.ops, 
		pid, eid, tmpdir):
	    return False

	# Put the file into a string to pass to emulab.
	try:
	    tcl = "".join([ l for l in open(tclfile,"r")])
	except EnvironmentError, e:
	    self.log.error("Can't read %s: %s" % (tclfile, e))
	    return False
	
	# Stage the new configuration
	if not self.modify_exp(pid, eid, tcl):
	    self.log.error("modify failed")
	    return False

	if not self.swap_exp(pid, eid, 'in'):
	    self.log.error("swap in failed")
	    return False
	# Everything has gone OK.
	self.get_mapping(pid,eid)
	return True

class exports_segment(emulab_segment.exports_segment):
    '''
    Class to export parameters from this segment.  The shared NAT allocates
    externally accessible ports in 2 steps.  Provisional ports are assigned
    when the experiment is created or modified and then preserved when possible
    across further mods.  This class creates a very simple experiment to make
    those provisional allocations and then installs them in the connInfo list.
    Note that this class takes on some of the initial steps from start_segment
    - clearing out old experiments and creating the dummy and initial
    experiment.
    '''
    def __init__(self, log=None, keyfile=None, debug=False, boss=None,
	    ops=None, cert=None):
	emulab_segment.exports_segment.__init__(self, log=log, keyfile=keyfile,
		debug=debug, boss=boss, ops=ops, cert=cert)

    def make_portal_experiment(self, pid, eid, tmpdir, gid, ports):
	ns_data = \
'''
set ns [new Simulator]
source tb_compat.tcl

'''
	for p in ports:
	    ns_data += 'set %s [$ns node]\n' % p
	    ns_data += \
		'tb-allow-external %s shared nat 0/0/tcp rdr 22/tcp\n' % p
	ns_data += \
'''	
$ns rtproto Session
$ns run 
'''
	# If ports is empty, just create the standard null experiment
	if len(ports) == 0:
	    ns_data = None

	state = self.get_state(pid, eid)

	if state != 'none':
	    self.terminate_exp(pid, eid)

	return self.make_null_experiment(pid, eid, tmpdir, gid, ns_data)

    def get_assignments(self, pid, eid, user, tmpdir, ports):
	'''
	Copy the output of the shared NAT risky experiment setup output to the
	temporary directory and parse it for preliminary (and hopefully
	permanent) port assignments.  That output is a pseudo-xmlrpc response.
	'''
	# No ports? Nothing to do here.
	if len(ports) == 0:
	    return True

	remote_risk = os.path.join('/proj', pid, 'exp', eid, 'tbdata',
		'%s.risk' % eid)
	local_risk = os.path.join(tmpdir, 'risk')

	try:
	    self.scp_file_from(remote_risk, user, self.ops, local_risk)
	    f = open(local_risk)
	    params, method = xmlrpclib.loads(f.read())
	    f.close()
	    # params should be a single element tuple containing the "xmlrpc"
	    # request as a dict
	    if not isinstance(params, tuple):
		self.log.error('misformatted risk file: params not a tuple')
		return False
	    params = params[0]

	    if 'assignments' not in params:
		self.log.error('misformatted risk file: No arguments in params')
		return False

	    for k, v in params['assignments'].items():
		if not k.endswith('/22/tcp'):
		    continue
		# Get the portal name from the SSH port assignment
		p = k[0:k.index('/22/tcp')]
		if p in ports:
		    ports[p]  = v
	except Exception, e:
	    self.log.error('Error getting assignments: %s' % e)
	    return False
	return True


    def __call__(self, parent, eid, pid, user, connInfo, tmpdir, timeout=0,
	    gid=None):
	'''
	Create an experiment containing only portals and read the ssh_ports
	from it.  The shared NAT properly marshals DNS names.
	'''

	ports = { }

	for c in connInfo:
	    for p in c.get('parameter', []):
		# Only set output parameters
		if  p.get('type', '') != 'output':
		    continue
		name = p.get('name', '')
		k = p.get('key', None)
		if name == 'peer':
		    if k is None or k.index('/') == -1:
			self.log.debug('Bad key for peer: %s' % k)
			return False

		    portal_name = k[k.index('/')+1:]

		    if parent.nat_portal:
			value = parent.nat_portal
		    elif k and k.index('/') != -1:
			value = "%s.%s.%s%s" % \
			    (portal_name, eid, pid, parent.domain)
		    else:
			self.log.error('Bad export request: %s' % p)
			continue
		    p['value'] = value
		    ports[portal_name] = None
		    self.log.debug('Assigning %s to %s' % (k, value))
		    self.log.debug('Adding entry for %s to ports' % portal_name)
		elif name == 'ssh_port':
		    # not yet
		    continue
		else:
		    self.log.error("Unknown export parameter: %s" % name)
		    return False

	if not self.make_portal_experiment(pid, eid, tmpdir, gid, ports):
	    return False

	if not self.get_assignments(pid, eid, user, tmpdir, ports):
	    return False

	# Through these again to set ssh_ports
	for c in connInfo:
	    for p in c.get('parameter', []):
		# Only set output parameters
		if  p.get('type', '') != 'output':
		    continue
		name = p.get('name', '')
		k = p.get('key', None)
		if name != 'ssh_port':
		    continue
		if k is None or k.index('/') == -1:
		    self.log.debug('Bad key for ssh port: %s' % k)
		    return False

		# The portal name is the string following the slash in the key,
		# with the trainling (5 character) -port removed.
		portal_name = k[k.index('/')+1:-5]
		if portal_name not in ports or ports[portal_name] is None:
		    self.log.error('No port assigned to %s' % portal_name)
		    return False
		value = ports[portal_name]
		p['value'] = value
		self.log.debug('Assigning %s to %s' % (k, value))
	
	return True

# These are all simple aliases to the main emulab_segment line.
class stop_segment(emulab_segment.stop_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	emulab_segment.stop_segment.__init__(self, log=log, keyfile=keyfile,
		debug=debug, boss=boss, ops=ops, cert=cert)

class info_segment(emulab_segment.info_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	emulab_segment.info_segment.__init__(self, log=log, keyfile=keyfile,
		debug=debug, boss=boss, ops=ops, cert=cert)

class operation_segment(emulab_segment.operation_segment):
    def __init__(self, log=None, keyfile=None, debug=False, boss=None, 
	    ops=None, cert=None):
	emulab_segment.operation_segment.__init__(self, log=log,
		keyfile=keyfile, debug=debug, boss=boss, ops=ops, cert=cert)