agent.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import os
import re
import sys
import time
import uuid
import socket
import atexit
import logging
import subprocess
try:
    from urlparse import urlparse
except ImportError as e:
    from urllib.parse import urlparse
try:
    import win32process
    import win32con
except ImportError as e:
    pass
try:
    import cPickle as pickle
except ImportError as e:
    import pickle
try:
    from xmlrpclib import ServerProxy
    from xmlrpclib import Error
except ImportError as e:
    from xmlrpc.client import ServerProxy
    from xmlrpc.client import Error

from Peach import xmlrpc
from twisted.web import server

from Peach.Publishers import *
from Peach.Engine.common import *
from Peach.Utilities.common import *


agentProcesses = []


@atexit.register
def cleanup_agent():
    if agentProcesses and len(agentProcesses) > 0:
        logging.debug("Killing agents.")
        for process in agentProcesses:
            process.terminate()
            process.kill()


def Debug(msg):
    #print msg
    pass


def MonitorDebug(monitor, msg):
    print("Monitor[%s]: %s" % (monitor, highlight.info(msg)))


def PeachStr(s):
    """Our implementation of str() which does not convert None to 'None'."""
    return None if s is None else str(s)


class Monitor(object):
    """
    Extend from this to implement a Monitor.
    Monitors are run by an Agent and must operate in an async manner.
    Any blocking tasks must be performed in another thread.

    Agent: onTestStarting()
    Agent: onTestFinished()
    Agent: redoTest()
    Agent: Sending redoTest result [False]
    Agent: detectFault()
    Agent: Detected fault!
    Agent: Sending detectFault result [True]
    Agent: getMonitorData()
    Agent: onFault()
    Agent: stopRun()
    Agent: onShutdown()
    """

    def __init__(self, args):
        """
        Arguments are supplied via the Peach XML file.

        @type	args: Dictionary
        @param	args: Dictionary of parameters
        """
        self._name = None

    def OnTestStarting(self):
        """Called right before start of test case or variation."""
        pass

    def OnTestFinished(self):
        """Called right after a test case or variation."""
        pass

    def GetMonitorData(self):
        """Get any monitored data from a test case."""
        return None

    def RedoTest(self):
        """Should the current test be re-performed."""
        return False

    def DetectedFault(self):
        """Check if a fault was detected."""
        return False

    def OnFault(self):
        """Called when a fault was detected."""
        pass

    def OnShutdown(self):
        """Called when Agent is shutting down, typically at end of a test run or when a Stop-Run occurs."""
        pass

    def StopRun(self):
        """Return True to force test run to fail. This should return True if an unrecoverable error occurs."""
        return False

    def PublisherCall(self, method):
        """Called when a call action is being performed.
        Call actions are used to launch programs, this gives the monitor a chance to determine if it should be
        running the program under a debugger instead. Note: This is a bit of a hack to get this working.
        """
        pass

# Need to define Monitor before we do this include!
from Peach.Agent import *


class _MsgType(object):
    """Type of message"""
    ClientHello = 2  #: Sent by fuzzer, will include password (optional)
    AgentHello = 1  #: Sent if password okay, else drop
    ClientDisconnect = 5
    AgentDisconnect = 6
    AgentReady = 7
    Ack = 17  #: Ack that we completed and finished
    Nack = 18  #: Ack that we did not complete, error occurred
    #: exception in msg.exception
    ## Monitors
    GetMonitorData = 11
    DetectFault = 8
    OnFault = 9
    OnShutdown = 10  #: Test completed, shutting down
    OnTestStarting = 13  #: On Test Case Starting
    OnTestFinished = 14  #: On Test Case Finished
    StopRun = 20
    RedoTest = 30  #: Should we re-perform the current test?
    StartMonitor = 15  #: Startup a monitor
    # Expect a msg.monitorName: str, msg.monitorClass: str and msg.params: dictionary
    StopMonitor = 16  #: Stop a monitor
    ## Publishers
    PublisherStart = 21
    PublisherStop = 22
    PublisherAccept = 23
    PublisherConnect = 24
    PublisherClose = 25
    PublisherCall = 26  #: Notify that we are running a call action
    PublisherProperty = 27  #: Notify that we are running a property action
    PublisherSend = 28
    PublisherReceive = 29


class _Msg(object):
    """This is a message holder that is serialized and sent over the named pipe."""

    def __init__(self, id, type, results=None):
        self.id = id
        self.type = type
        self.results = results
        self.stopRun = False
        self.password = False
        self.pythonPaths = None
        self.imports = None


class Agent(object):
    """A remote or local Agent that listens on a named pipe.
    Each agent can only be connected to by a single Peach Fuzzer.
    """

    def __init__(self, password=None, port=9000):
        """
        Creates and Agent instance and attempts to connect to the AgentMaster.
        If connection works the Client Hello message is sent.

        @type	password: string
        @param	password: Password to use
        """
        from twisted.internet import reactor

        agent = AgentXmlRpc()
        agent._password = password
        agent._monitors = []
        agent._publishers = {}
        agent._id = None
        reactor.listenTCP(port, server.Site(agent))
        if agent._password is not None:
            print("\n[Agent] Listening on [%s] with password [%s]\n" % (port, agent._password))
        else:
            print("\n[Agent] Listening on [%s] with no password\n" % port)
        reactor.run()


class AgentXmlRpc(xmlrpc.XMLRPC):
    def xmlrpc_clientHello(self, msg):
        msg = pickle.loads(msg)
        if msg.password != self._password:
            print("Agent: Incorrect password on clientHello [%s]" % msg.password)
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.ClientHello:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: clientHello()")
        if self._id is not None:
            self._stopAllMonitors()
        self._id = str(uuid.uuid1())
        print("Agent: Session ID: ", self._id)
        # Handle any PythonPath or Imports
        if msg.pythonPaths is not None:
            for p in msg.pythonPaths:
                sys.path.append(p['name'])
        if msg.imports is not None:
            for i in msg.imports:
                self._handleImport(i)
        print("Agent: clientHello() all done")
        return pickle.dumps(_Msg(self._id, _MsgType.AgentHello))

    def GetClassesInModule(self, module):
        """Return array of class names in module."""
        classes = []
        for item in dir(module):
            i = getattr(module, item)
            if isinstance(i, type) and item[0] != '_':
                classes.append(item)
        return classes

    def _handleImport(self, i):
        importStr = i['import']
        if i['from'] is not None and len(i['from']) > 0:
            fromStr = i['from']
            if importStr == "*":
                module = __import__(PeachStr(fromStr), globals(), locals(), [PeachStr(importStr)], -1)
                try:
                    # If we are a module with other modules in us then we have an __all__
                    for item in module.__all__:
                        globals()[item] = getattr(module, item)
                except Exception as e:
                    # Else we just have some classes in us with no __all__
                    for item in self.GetClassesInModule(module):
                        globals()[item] = getattr(module, item)
            else:
                module = __import__(PeachStr(fromStr), globals(), locals(), [PeachStr(importStr)], -1)
                for item in importStr.split(','):
                    item = item.strip()
                    globals()[item] = getattr(module, item)
        else:
            globals()[importStr] = __import__(PeachStr(importStr), globals(), locals(), [], -1)

    def xmlrpc_clientDisconnect(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.ClientDisconnect:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: clientDisconnect()")
        self._stopAllMonitors()
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_stopRun(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.StopRun:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: stopRun()")
        msg = _Msg(None, _MsgType.Ack)
        msg.results = False
        for m in self._monitors:
            if m.StopRun():
                print("Agent: Stop run request!")
                msg.results = True
        return pickle.dumps(msg)

    def xmlrpc_detectFault(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.DetectFault:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: detectFault()")
        msg = _Msg(None, _MsgType.Ack)
        msg.results = False
        for m in self._monitors:
            if m.DetectedFault():
                print("Agent: Detected fault!")
                msg.results = True
        print("Agent: Sending detectFault result [%s]" % repr(msg.results))
        return pickle.dumps(msg)

    def xmlrpc_redoTest(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.RedoTest:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: redoTest()")
        msg = _Msg(None, _MsgType.Ack)
        msg.results = False
        for m in self._monitors:
            if m.RedoTest():
                msg.results = True
        print("Agent: Sending redoTest result [%s]" % repr(msg.results))
        return pickle.dumps(msg)

    def xmlrpc_getMonitorData(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.GetMonitorData:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: getMonitorData()")
        msg = _Msg(None, _MsgType.Ack)
        msg.results = []
        for m in self._monitors:
            try:
                data = m.GetMonitorData()
                if data is not None:
                    msg.results.append(data)
            except Exception as e:
                print("Agent: getMonitorData: Failrue getting data from:", m.monitorName)
                raise
        return pickle.dumps(msg)

    def xmlrpc_onFault(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.OnFault:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: onFault()")
        for m in self._monitors:
            m.OnFault()
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_onTestFinished(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.OnTestFinished:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: onTestFinished()")
        for m in self._monitors:
            m.OnTestFinished()
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_onTestStarting(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.OnTestStarting:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: onTestStarting()")
        for m in self._monitors:
            m.OnTestStarting()
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_onPublisherCall(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.PublisherCall:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: onPublisherCall():", msg.method)
        outRet = None
        for m in self._monitors:
            ret = m.PublisherCall(msg.method)
            if ret is not None:
                outRet = ret
        return pickle.dumps(_Msg(None, _MsgType.Ack, outRet))

    def _stopAllMonitors(self):
        """Stop all monitors. Part of resetting our connection."""
        for m in self._monitors:
            m.OnShutdown()
        self._monitors = []

    def xmlrpc_onShutdown(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.OnShutdown:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: onShutdown()")
        self._stopAllMonitors()
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_stopMonitor(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.StopMonitor:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: stopMonitor(%s)" % msg.monitorName)
        for i in range(len(self._monitors)):
            m = self._monitors[i]
            if m._name == msg.monitorName:
                try:
                    m.OnShutdown()
                except Exception as e:
                    pass
                self._monitors.remove(m)
                break
        return pickle.dumps(_Msg(None, _MsgType.Ack))

    def xmlrpc_startMonitor(self, msg):
        msg = pickle.loads(msg)
        if self._id is None or msg.id != self._id:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        if msg.type != _MsgType.StartMonitor:
            return pickle.dumps(_Msg(None, _MsgType.Nack))
        print("Agent: startMonitor(%s)" % msg.monitorName)
        try:
            code = msg.monitorClass + "(msg.params)"
            print("code:", code)
            monitor = eval(code)
            if monitor is None:
                print("Agent: Unable to create Monitor [%s]" % msg.monitorClass)
                return pickle.dumps(_Msg(self._id, _MsgType.Nack, "Unable to create Monitor [%s]" % msg.monitorClass))
            monitor.monitorName = msg.monitorName
            self._monitors.append(monitor)
            print("Agent: Sending Ack")
            return pickle.dumps(_Msg(None, _MsgType.Ack))
        except Exception as e:
            print("Agent: Unable to create Monitor [%s], exception occured." % msg.monitorClass)
            raise

    ## Publishers ######################################################

    def xmlrpc_publisherInitialize(self, id, name, cls, *args):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherInitialize: Can't validate clients session id")
            return -1
        print("Agent: xmlrpc_publisherInitialize(%s, %s)" % (name, cls))
        try:
            code = cls + "("
            for cnt in range(len(args[0])):
                code += "args[0][%d]," % cnt
            code = code[:-1] + ")"
            print("Agent: Code: %s" % code)
            publisher = eval(code)
            if publisher is None:
                print("Agent: Unable to create Publisher [%s]" % cls)
                return -2
            publisher.publisherName = name
            self._publishers[name] = publisher
            print("Agent: Publisher created okay!")
            return 0
        except Exception as e:
            print("Agent: Unable to create Publisher [%s], exception occurred." % cls)
            return -2

    def xmlrpc_publisherStart(self, id, name):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherStart: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        self._publishers[name].start()
        return 0

    def xmlrpc_publisherStop(self, id, name):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherStop: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        self._publishers[name].stop()
        return 0

    def xmlrpc_publisherAccept(self, id, name):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherAccept: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        self._publishers[name].accept()
        return 0

    def xmlrpc_publisherConnect(self, id, name):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherConnect: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        self._publishers[name].connect()
        return 0

    def xmlrpc_publisherClose(self, id, name):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherClose: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        self._publishers[name].close()
        return 0

    def xmlrpc_publisherCall(self, id, name, method, args):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherCall: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        args = pickle.loads(args)
        ret = self._publishers[name].call(method, args)
        if ret is None:
            return 0
        return ret

    def xmlrpc_publisherProperty(self, id, name, property, value):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherProperty: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        value = pickle.loads(value)
        ret = self._publishers[name].property(property, value)
        if ret is None:
            return 0
        return ret

    def xmlrpc_publisherSend(self, id, name, data):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherSend: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        # TODO: ??
        value = pickle.loads(data)
        ret = self._publishers[name].send(data)
        if ret is None:
            return 0
        return ret

    def xmlrpc_publisherReceive(self, id, name, size):
        if self._id is None or id != self._id:
            print("xmlrpc_publisherReceive: Can't validate clients session id")
            return -1
        if name not in self._publishers:
            return -2
        return self._publishers[name].receive(size)


class AgentClient(object):
    """An Agent client. Clients connect and send/recieve messages with a single remote Agent."""

    def __init__(self, agentUri, password, pythonPaths=None, imports=None, configs=None):
        """
        Creates and Agent instance and attempts to connect to the AgentMaster.
        If connection works the Client Hello message is sent.

        @type	agentUri: string
        @param	agentUri: Url of agent
        @type	password: string
        @param	password: [optional] Password to authenticate to agent.  Warning: CLEAR-TEXT!!
        @type	pythonPaths: list
        @param	pythonPaths: List of paths we should configure on the remote agent
        @type	imports: list
        @param	imports: list of imports that should be performed on the remote agent
        """

        self._pythonPaths = pythonPaths
        self._imports = imports
        self._password = password
        self._monitors = []
        self._id = None
        self._agent = None
        self._agentUri = agentUri

        agentUrl = urlparse(agentUri)
        agentPort = agentUrl.port
        agentHostname = agentUrl.hostname

        # This is nicer, but does not work on darwin:
        #if socket.getfqdn(agentHostname) in ('localhost', socket.gethostname()):
        if agentHostname in ("127.0.0.1", "0.0.0.0", "localhost", socket.gethostname()):
            peachPath = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            macros = []
            if configs:
                macros.append("-macros")
                for kv in configs.iteritems():
                    macros.append("=".join(kv))
            if sys.platform == "win32":
                agentProcess = subprocess.call(['start',
                                                "Peach Agent",
                                                sys.executable,
                                                "%s\peach.py" % peachPath,
                                                "-agent", str(agentPort), password] + macros)
                agentProcesses.append(agentProcess)
            elif sys.platform == "darwin":
                if not self.isPeachRunning(agentPort):
                    peachAgentCommand = [sys.executable, 'peach.py', '-agent', str(agentPort), password] + macros
                    if configs and getBooleanAttribute(configs, "HideAgentWindow"):
                        logging.warning("Agent window will not be visible!")
                        agentProcess = subprocess.Popen(peachAgentCommand,
                                                        cwd=peachPath,
                                                        stdout=open('/dev/null', 'w'))
                        agentProcesses.append(agentProcess)
                    else:
                        logging.info("Opening agent in new terminal.")
                        osxTerminalCommand = \
                            """osascript -e 'tell application "Terminal" to do script "cd %s; %s; exit"'""" % \
                            (peachPath, re.sub(r"""(['"])""", r"\\\1", subprocess.list2cmdline(peachAgentCommand)))
                        agentProcess = subprocess.Popen(osxTerminalCommand,
                                                        stdout=subprocess.PIPE,
                                                        shell=True)
                        agentProcesses.append(agentProcess)
            elif sys.platform == "linux2":
                if not self.isPeachRunning(agentPort):
                    logging.info("Opening agent in new terminal.")
                    peachAgentCommand = [sys.executable, "peach.py", "-agent", str(agentPort), password] + macros
                    if "COLORTERM" in os.environ and 'gnome-terminal' in os.environ["COLORTERM"]:
                        linuxTerminalCommands = ['gnome-terminal', '-x'] + peachAgentCommand
                    else:
                        linuxTerminalCommands = ['xterm', '-hold']
                        if configs and getBooleanAttribute(configs, "AgentTerminalLogging"):
                            linuxTerminalCommands += ["-l"]
                        linuxTerminalCommands += ['-e'] + peachAgentCommand
                    agentProcess = subprocess.Popen(linuxTerminalCommands,
                                                    cwd=peachPath,
                                                    stdout=subprocess.PIPE)
                    agentProcesses.append(agentProcess)
            else:
                raise PeachException("We only support auto starting agents on Windows, OSX and Linux. "
                                     "Please configure all agents with location URIs and launch any Agent manually.")

        # Connect to remote agent
        try:
            m = re.search(r"://([^/]*)", agentUri)
            self._name = m.group(1)
        except:
            raise PeachException("Please make sure your agent location string is a valid http URL.")
        self.Connect()

    def isPeachRunning(self, agentPort):
        running = False
        if sys.platform == "darwin" or sys.platform == "linux2":
            agentList = subprocess.check_output(["ps x | grep peach.py"], shell=True)
            for agent in agentList.splitlines():
                try:
                    port = re.findall("peach\.py \-agent (\d+)", agent)[0]
                except IndexError:
                    continue
                if int(port) == agentPort:
                    running = True
        if running:
            logging.warning("An agent is already running on port %d" % agentPort)
        return running

    def LaunchWin32Process(self, command):
        try:
            StartupInfo = win32process.STARTUPINFO()
            StartupInfo.dwFlags = win32process.STARTF_USESHOWWINDOW
            StartupInfo.wShowWindow = win32con.SW_NORMAL
            win32process.CreateProcess(
                None,
                command,
                None,
                None,
                0,
                win32process.NORMAL_PRIORITY_CLASS,
                None,
                None,
                StartupInfo)
        except Exception as e:
            print(sys.exc_info())
            print("Exception in LaunchWin32Process")
            pass

    def Connect(self):
        """Connect to agent. Will retry the connection 10 times before giving up."""
        for i in range(20):
            try:
                self._agent = ServerProxy(self._agentUri)
                msg = _Msg(None, _MsgType.ClientHello, self._name)
                msg.password = self._password
                msg.pythonPaths = self._pythonPaths
                msg.imports = self._imports
                msg = pickle.loads(self._agent.clientHello(pickle.dumps(msg)))
                if msg.type != _MsgType.AgentHello:
                    raise PeachException("Error connecting to remote agent %s, invalid response." % self._name)
                self._id = msg.id
                return
            except Exception as e:
                if i == 19:
                    raise
            time.sleep(1)
            logging.warning("Agent connection failed, retrying.")

    def Reconnect(self):
        """Reconnect to remote agent"""
        try:
            self.Connect()
            for m in self._monitors:
                self.StartMonitor(m[0], m[1], m[2], True)
        except Exception as e:
            raise PeachException("Unable to reconnect to Agent %s." % self._name)

    def StartMonitor(self, name, classStr, params, restarting=False):
        Debug("> StartMonitor")
        msg = _Msg(self._id, _MsgType.StartMonitor)
        msg.monitorName = name
        msg.monitorClass = classStr
        msg.params = params
        msg = pickle.loads(self._agent.startMonitor(pickle.dumps(msg)))
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during StartMonitor call." % self._name)
        if not restarting:
            self._monitors.append([name, classStr, params])
        Debug("< StartMonitor")

    def StopMonitor(self, name):
        Debug("> StopMonitor")
        msg = _Msg(self._id, _MsgType.StopMonitor)
        msg.monitorName = name
        msg = pickle.loads(self._agent.stopMonitor(pickle.dumps(msg)))
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during StopMonitor call." % self._name)
        for m in self._monitors:
            if m[0] == name:
                self._monitors.remove(m)
        Debug("< StopMonitor")

    def OnTestStarting(self):
        """Called right before start of test."""
        Debug("> OnTestStarting")
        msg = _Msg(self._id, _MsgType.OnTestStarting)
        try:
            msg = pickle.loads(self._agent.onTestStarting(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during OnTestStarting call." % self._name)
        Debug("< OnTestStarting")

    def OnPublisherCall(self, method):
        Debug("> OnPublisherCall")
        msg = _Msg(self._id, _MsgType.PublisherCall)
        msg.method = method
        try:
            msg = pickle.loads(self._agent.onPublisherCall(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during OnPublisherCall call." % self._name)
        Debug("< OnPublisherCall")
        return msg.results

    def OnTestFinished(self):
        """Called right after a test."""
        Debug("> OnTestFinished")
        msg = _Msg(self._id, _MsgType.OnTestFinished)
        try:
            msg = pickle.loads(self._agent.onTestFinished(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during OnTestFinished call." % self._name)
        Debug("< OnTestFinished")

    def GetMonitorData(self):
        """Get any monitored data."""
        Debug("> GetMonitorData")
        msg = _Msg(self._id, _MsgType.GetMonitorData)
        try:
            msg = pickle.loads(self._agent.getMonitorData(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during GetMonitorData call." % self._name)
        Debug("< GetMonitorData")
        return msg.results

    def RedoTest(self):
        """Should we repeat current test."""
        Debug("> RedoTest")
        try:
            msg = _Msg(self._id, _MsgType.RedoTest)
            msg = pickle.loads(self._agent.redoTest(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during RedoTest call." % self._name)
        Debug("< RedoTest")

        return msg.results

    def DetectedFault(self):
        """Check if a fault was detected."""
        Debug("> DetectedFault")
        try:
            msg = _Msg(self._id, _MsgType.DetectFault)
            msg = pickle.loads(self._agent.detectFault(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during GetMonitorData call." % self._name)
        Debug("< DetectedFault")
        return msg.results

    def OnFault(self):
        """Called when a fault was detected."""
        Debug("> OnFault")
        try:
            msg = _Msg(self._id, _MsgType.OnFault)
            msg = pickle.loads(self._agent.onFault(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during GetMonitorData call." % self._name)
        Debug("< OnFault")

    def OnShutdown(self):
        """Called when Agent is shutting down."""
        Debug("> OnShutdown")
        msg = _Msg(self._id, _MsgType.OnShutdown)
        self._agent.onShutdown(pickle.dumps(msg))
        Debug("< OnShutdown")

    def StopRun(self):
        """Return True to force test run to fail. This should return True if an unrecoverable error occurs."""
        Debug("> StopRun")
        try:
            msg = _Msg(self._id, _MsgType.StopRun)
            msg = pickle.loads(self._agent.stopRun(pickle.dumps(msg)))
        except Exception as e:
            self.Reconnect()
            raise RedoTestException("Communication error with Agent %s" % self._name)
        if msg.type != _MsgType.Ack:
            raise PeachException("Lost connection to Agent %s during GetMonitorData call." % self._name)
        Debug("< StopRun")
        return msg.results

    #  Publishers
    def PublisherInitialize(self, name, cls, args):
        ret = self._agent.publisherInitialize(self._id, name, cls, args)
        if ret is not None and ret < 0:
            raise Exception("That sucked")

    def PublisherStart(self, name):
        ret = self._agent.publisherStart(self._id, name)
        if ret < 0:
            raise Exception("That sucked")

    def PublisherStop(self, name):
        ret = self._agent.publisherStop(self._id, name)
        if ret < 0:
            raise Exception("That sucked")

    def PublisherAccept(self, name):
        ret = self._agent.publisherAccept(self._id, name)
        if ret < 0:
            raise Exception("That sucked")

    def PublisherConnect(self, name):
        ret = self._agent.publisherConnect(self._id, name)
        if ret < 0:
            raise Exception("That sucked")

    def PublisherClose(self, name):
        ret = self._agent.publisherClose(self._id, name)
        if ret < 0:
            raise Exception("That sucked")

    def PublisherCall(self, name, method, args):
        ret = self._agent.publisherCall(self._id, name, method, pickle.dumps(args))
        if ret < 0:
            raise Exception("That sucked")
        return ret

    def PublisherProperty(self, name, property, value=None):
        ret = self._agent.publisherProperty(self._id, name, property, pickle.dumps(value))
        if ret < 0:
            raise Exception("That sucked")
        return ret

    def PublisherSend(self, name, data):
        ret = self._agent.publisherSend(self._id, name, pickle.dumps(data))
        if ret < 0:
            raise Exception("That sucked")
        return ret

    def PublisherReceive(self, name, size=None):
        ret = self._agent.publisherReceive(self._id, name, size)
        if ret < 0:
            raise Exception("That sucked")
        return ret


class AgentPlexer(object):
    """Manages communication with one or more agent."""

    def __init__(self):
        self._agents = {}

    def __getitem__(self, key):
        return self._agents[key]

    def __setitem__(self, key, value):
        self._agents[key] = value

    def AddAgent(self, name, agentUri, password=None, pythonPath=None, imports=None, configs=None):
        #m = re.search(r"://([^:/]*)", agentUri)
        #name = m.group(1)
        agent = AgentClient(agentUri, password, pythonPath, imports, configs)
        self._agents[name] = agent
        return agent

    def OnTestStarting(self):
        """Called right before start of test."""
        for name in self._agents.keys():
            self._agents[name].OnTestStarting()

    def OnPublisherCall(self, method):
        ourRet = None
        for name in self._agents.keys():
            ret = self._agents[name].OnPublisherCall(method)
            if ret is not None:
                ourRet = ret
        return ourRet

    def OnTestFinished(self):
        """Called right after a test."""
        for name in self._agents.keys():
            self._agents[name].OnTestFinished()

    def GetMonitorData(self):
        """Get any monitored data."""
        ret = {}
        for name in self._agents.keys():
            arrayOfMonitorData = self._agents[name].GetMonitorData()
            for hashOfData in arrayOfMonitorData:
                for key in hashOfData.keys():
                    ret["%s_%s" % (name, key)] = hashOfData[key]
        return ret

    def RedoTest(self):
        """Check if a fault was detected."""
        ret = False
        for name in self._agents.keys():
            if self._agents[name].RedoTest():
                ret = True
        return ret

    def DetectedFault(self):
        """Check if a fault was detected."""
        ret = False
        for name in self._agents.keys():
            if self._agents[name].DetectedFault():
                ret = True
        return ret

    def OnFault(self):
        """Called when a fault was detected."""
        for name in self._agents.keys():
            self._agents[name].OnFault()

    def OnShutdown(self):
        """Called when Agent is shutting down."""
        for name in self._agents.keys():
            self._agents[name].OnShutdown()
        self._agents = {}

    def StopRun(self):
        """Return True to force test run to fail. This should return True if an unrecoverable error occurs."""
        ret = False
        for name in self._agents.keys():
            if self._agents[name].StopRun():
                ret = True
        return ret

    # Publishers
    def PublisherInitialize(self, name, cls, args):
        for name in self._agents.keys():
            self._agents[name].PublisherInitialize(name, cls, args)

    def PublisherStart(self, name):
        for name in self._agents.keys():
            self._agents[name].PublisherStart(name)

    def PublisherStop(self, name):
        for name in self._agents.keys():
            self._agents[name].PublisherStop(name)

    def PublisherAccept(self, name):
        for name in self._agents.keys():
            self._agents[name].PublisherAccept(name)

    def PublisherConnect(self, name):
        for name in self._agents.keys():
            self._agents[name].PublisherConnect(name)

    def PublisherClose(self, name):
        for name in self._agents.keys():
            self._agents[name].PublisherClose(name)

    def PublisherCall(self, name, method, args):
        for name in self._agents.keys():
            ret = self._agents[name].PublisherCall(name, method, args)
            if ret is not None:
                return ret

    def PublisherProperty(self, name, property, value=None):
        for name in self._agents.keys():
            ret = self._agents[name].PublisherProperty(name, property, value)
            if ret is not None:
                return ret

    def PublisherSend(self, name, data):
        for name in self._agents.keys():
            self._agents[name].PublisherSend(name, data)

    def PublisherReceive(self, name, size=None):
        for name in self._agents.keys():
            ret = self._agents[name].PublisherReceive(name, size)
            if ret is not None:
                return ret