Little proof of concept to send source and pickled objects over the network for remote execution

agentsmith.py

#!/usr/bin/env python
# encoding: utf-8
"""
agentsmith.py

short demonstration on sending pickles over the net for remote execution

usage on agent: 	./agentsmith.py --agent
usage on controller:	./agentsmith.py --connectagent [agent IP address]

scipy/numpy is not really needed for the concept, you just need to
modify the ASWorkUnit code to make it work without ^^

just a little proof of concept, definitely not bug free :-D
"""

import sys
import getopt
from socket import *
import scipy
import scipy.integrate
import numpy
import pickle
import inspect

help_message = '''
usage on agent: 	./agentsmith.py --agent
usage on controller:	./agentsmith.py --connectagent [agent IP address]
'''

class ASWorkUnit(object):
    """docstring for ASWorkUnit"""
    def __init__(self):
        super(ASWorkUnit, self).__init__()
        self.answer = None
        
    def run(self):
        import scipy, numpy, scipy.integrate
        from numpy import pi, sqrt, sin
        L = 1.0
        M = 0.11
        
        f = lambda x: numpy.array([ L/(2*pi) * sqrt(1 - (-(M**2*pi**2)/L**2) * (sin( (2*pi*x)/L) )**2) ])
        for i in xrange(1,1000):
            scipy.integrate.quad(f, 0, (2*pi*0.43784)/L)
        self.answer = (scipy.integrate.quad(f, 0, (2*pi*0.43784)/L))[0]
        
    def doImports(self):
        pass
    
    def isReady(self):
        if self.answer is None:
            return False
        else:
            return True


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


def smithcontroller(agent):
    clientsocket = socket(AF_INET, SOCK_STREAM)
    print "Connecting to agent %s" %(agent)
    clientsocket.connect((agent, 2971))
    source = inspect.getsource(ASWorkUnit)
    p = pickle.dumps(ASWorkUnit, 2)
    clientsocket.send(str(len(source)))
    imports = "scipy,numpy"
    while True:
        response = clientsocket.recv(256)
        if response:
            status = response.split(": ",1)[1]
            if status == "send-source":
                print "sending source..."
                clientsocket.send(source)
            if status == "send-workunit-length":
                print "sending workunit configuration..."
                clientsocket.send(str(len(p)))
            if status == "send-workunit-data":
                print "sending workunit data..."
                clientsocket.send(p)
            if status == "send-ilength":
                print "sending imports configuration..."
                clientsocket.send(str(len(imports)))
            if status == "send-imports":
                print "sending imports..."
                clientsocket.send(imports)
            if status.find("answer") != -1:
                print "Got answer to workunit: %s" %(status.split("-")[1])
            if status.find("failure") != -1:
                print "error occured: %s" %(status)
                break
        else:
            break

    return


def smithagent():
    hostname = "localhost"
    print "Creating socket."
    agentsocket = socket(AF_INET, SOCK_STREAM)
    print "Binding socket to %s:2971." %(hostname)
    agentsocket.bind((hostname, 2971))
    print "Now listening."
    agentsocket.listen(5)

    while True:
        (clientsocket, address) = agentsocket.accept()
        while True:
            length = long(clientsocket.recv(64))
            if length:
                print "source length is %i bytes" %(length)
                clientsocket.send("Status: send-source")
            else:
                clientsocket.send("Status: failure-transfer-source-length")
                break
            source = clientsocket.recv(length)
            if source:
                print "source received"
                clientsocket.send("Status: send-ilength")
            else:
                clientsocket.send("Status: failure-transfer-source")
                break
            importslen = long(clientsocket.recv(64))
            if importslen:
                print "import length is %d bytes" %(importslen)
                clientsocket.send("Status: send-imports")
            else:
                clientsocket.send("Status: failure-transfer-imports-length")
            
            imports = clientsocket.recv(importslen)
            if imports:
                print "imports received"
                clientsocket.send("Status: send-workunit-length")
            else:
                clientsocket.send("Status: failure-transfer-imports")    
                
            mixedpickleslen = long(clientsocket.recv(64))
            if mixedpickleslen:
                print "mixed pickles length is %d bytes" %(mixedpickleslen)
                clientsocket.send("Status: send-workunit-data")
            else:
                clientsocket.send("Status: failure-transfer-workunit-length")
                break            
            data = clientsocket.recv(mixedpickleslen)
            if data:
                imports = imports.split(",")
                print "have to import %s" %(imports)
                for im in imports:
                    __import__(im, globals())
                exec(source, globals())
                workUnit = pickle.loads(data)
                workUnitInstance = workUnit()
                workUnitInstance.doImports()
                workUnitInstance.run()
                if workUnitInstance.isReady():
                    clientsocket.send("Status: answer-" + str(workUnitInstance.answer))
                    break
                else:
                    print "ASWorkUnit is not ready."
                    break
            else:
                clientsocket.send("Status: failure-transfer-workunit")
                break            
        clientsocket.close()

def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
	    opts, args = getopt.getopt(argv[1:], "hac:", ["help", "agent","connectagent="])
        except getopt.error, msg:
            raise Usage(msg)
        
        # option processing
        for option, value in opts:
            if option in ("-h", "--help"):
                raise Usage(help_message)
	    if option in ("-a", "--agent"):
		smithagent()
		return 0
	    if option in ("-c", "--connectagent"):
		smithcontroller(agent=value)
		return 0
	raise Usage(help_message)
    
    except Usage, err:
        print >> sys.stderr, sys.argv[0].split("/")[-1] + ": " + str(err.msg)
        return 2

    return 0


if __name__ == "__main__":
    sys.exit(main())