class LoaderConnection(RequestObject):
def init(self):
self.stack = Stack((StackableSocket(sock=self.conn),
StackablePacketAssembler(),
StackablePickler()))
def destroy(self):
self.stack.close()
del self.stack
def receive(self):
try:
obj = self.stack.poll()
if obj != None:
try:
print('[LOADER]', obj['load'], 'to Job', obj['id'])
def respond(x):
if x != None:
mgr.get_job(obj['id']).modules.append(obj['load'])
self.stack.write({'module': x, 'name': obj['load']})
mgr.get_module(obj['load'], respond)
except:
print('[LOADER] Received malformed request:', obj)
return True
except StackableError:
return False
class DispatchPusher(object):
def __init__(self, ip=None, port=None):
self.stack = None
self.com_id = 0
self.cbs = {
'all': [
lambda x: print(json.dumps(x, sort_keys=True, indent=3))
]
}
if ip != None and port != None:
self.connect(ip, port)
def connect(self, ip, port):
self.stack = Stack((StackableSocket(ip=ip, port=port),
StackablePacketAssembler(),
StackablePickler()))
def reply(self, cmd, args):
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def send(self, cmd, args):
self.com_id += 1
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def register_cb(self, cmd, cb):
try:
self.cbs[cmd].append(cb)
except:
self.cbs[cmd] = [cb]
def deregister_cb(self, cmd, cb):
try:
self.cbs[cmd].remove(cb)
except:
pass
def push_module(self, name, bytecode, source, _type, meta):
self.send('push_module', {'name': name, 'bytecode': bytecode, 'source': source, 'type': _type, 'meta': meta})
def dispatch(self, dispatcher, module):
self.send('dispatch', {'name': module, 'targets': [dispatcher]})
def status(self, dispatcher, job):
self.send('get_status', {'target': dispatcher, 'job_id': job})
def get_clients(self):
self.send('get_clients', {})
def get_jobs(self, dispatcher):
self.send('get_jobs', {'target': dispatcher})
def close(self):
self.stack.close()
def monitor(self):
while True:
o = self.stack.read()
if o['cmd'] in self.cbs:
for cb in self.cbs[o['cmd']]:
cb(o['args'])
class BackendConnection(RequestObject):
def init(self):
self.ip = self.conn.getpeername()[0]
print('[B] New connection:\t\t\t', self.ip)
self.stack = Stack((StackableSocket(sock=self.conn),
StackablePacketAssembler(acceptAllMagic=True)))
self.mgc = None
def write(self, obj):
self.stack.write(obj)
def destroy(self):
try:
print('[B] Closing connection:\t\t\t', self.ip)
self.stack.close()
magics[self.mgc].remove(self)
except:
pass
def receive(self):
try:
obj = self.stack.poll()
if obj != None:
if self.mgc == None:
k = self.stack[1]
self.mgc = k.hdr
k.sndhdr = self.mgc
k.magics = [self.mgc]
k.acceptAllMagic = False
if self.mgc not in magics:
magics[self.mgc] = []
print('[B] Magic identified:\t\t\t', self.mgc, "for:\t",
self.ip)
magics[self.mgc].append(self)
x = magics[self.mgc]
for ix in x:
if ix is not self:
try:
print(
"[B] (" + str(self.ip) + "->" + str(ix.ip) +
"):\t", obj)
ix.write(obj)
except StackableError, e:
print("[B] (" + str(ix.ip) + ") Terminating:\t", e)
ix.destroy()
return True
except StackableError, e:
print("[B] (" + str(self.ip) + ") Terminating:\t", e)
return False
class BackendConnection(RequestObject):
def init(self):
self.ip = self.conn.getpeername()[0]
print('[B] New connection:\t\t\t', self.ip)
self.stack = Stack((StackableSocket(sock=self.conn),StackablePacketAssembler(acceptAllMagic=True)))
self.mgc = None
def write(self, obj):
self.stack.write(obj)
def destroy(self):
try:
print('[B] Closing connection:\t\t\t', self.ip)
self.stack.close()
magics[self.mgc].remove(self)
except:
pass
def receive(self):
try:
obj = self.stack.poll()
if obj != None:
if self.mgc == None:
k = self.stack[1]
self.mgc = k.hdr
k.sndhdr = self.mgc
k.magics = [self.mgc]
k.acceptAllMagic = False
if self.mgc not in magics:
magics[self.mgc] = []
print('[B] Magic identified:\t\t\t', self.mgc, "for:\t", self.ip)
magics[self.mgc].append(self)
x = magics[self.mgc]
for ix in x:
if ix is not self:
try:
print("[B] ("+str(self.ip)+"->"+str(ix.ip)+"):\t", obj)
ix.write(obj)
except StackableError,e:
print("[B] ("+str(ix.ip)+") Terminating:\t", e)
ix.destroy()
return True
except StackableError,e:
print("[B] ("+str(self.ip)+") Terminating:\t", e)
return False
auto.clearEvent(p['name'])
return {'type': 'info', 'payload': {'status': 'ok'}}
elif a['type'] == 'disable_event':
auto.disableEvent(p['name'])
return {'type': 'info', 'payload': {'status': 'ok'}}
elif a['type'] == 'enable_event':
auto.enableEvent(p['name'])
return {'type': 'info', 'payload': {'status': 'ok'}}
elif a['type'] == 'on':
auto.on(p['name'])
return {'type': 'info', 'payload': {'status': 'ok'}}
elif a['type'] == 'off':
auto.off(p['name'])
return {'type': 'info', 'payload': {'status': 'ok'}}
return {'type': 'info', 'payload': {'status': 'error'}}
while 1:
try:
stack = Stack(
(StackableSocket(ip=server, port=serverPort),
StackablePacketAssembler(magics=[magic]), StackableJSON()))
stack.write({})
while 1:
stack.write(parse(stack.read()))
except:
stack.close()
traceback.print_exc()
sleep(5)
class Dispatcher(RequestObject):
'The connection objects - represents a connected dispatching node'
def init(self):
self.com_id = 0
self.stack = Stack((StackableSocket(sock=self.conn),
StackablePacketAssembler(),
StackablePickler()))
self.uuid = str(uuid4())
self.jobs = []
root.register(self)
self.send('init', {'uuid': self.uuid})
# self.dispatch('welcome')
self.dispatch('test')
def destroy(self):
root.deregister(self)
self.stack.close()
del self.stack
def dispatch(self, w):
self.send('execute', {'name': w})
def find_job(self, _id):
for i in self.jobs:
if i.id == _id:
return i
def print_status(self, job):
print('[JOB %s-%d] ---- STDOUT ----' % (self.uuid, job.id))
print(job.out)
print('[JOB %s-%d] ---- STDERR ----' % (self.uuid, job.id))
print(job.err)
print('[JOB %s-%d] -- MODULELIST --' % (self.uuid, job.id))
print(', '.join(job.modules))
print('[JOB %s-%d] Completed' % (self.uuid, job.id), job.name)
def reply(self, cmd, args):
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def send(self, cmd, args):
self.com_id += 1
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def handle(self, obj):
try:
cmd = obj['cmd']
args = obj['args']
except:
print('[DISPATCHER] Unknown blob:', obj)
return
if cmd == 'get_module':
try:
mod = root.retrieve(args['name'])
self.reply('module_update', {'module': mod})
self.reply('return', {'status': 0, 'cmd': 'get_module'})
except:
self.reply('module_update', {'not_found': args['name']})
self.reply('return', {'status': -1, 'cmd': 'get_module'})
elif cmd == 'probe_module':
try:
mod = root.retrieve(args['name'])
self.reply('module_probe', {'name': args['name'], 'meta': mod['meta'] if mod != None else None})
self.reply('return', {'status': 0, 'cmd': 'probe_module'})
except:
self.reply('return', {'status': -1, 'cmd': 'probe_module'})
elif cmd == 'push_module':
try:
root.put(args['name'],
args['bytecode'],
args['source'],
args['type'],
args['meta'])
self.reply('return', {'status': 0, 'cmd': 'push_module'})
except:
self.reply('return', {'status': -1, 'cmd': 'push_module'})
elif cmd == 'dispatch':
try:
mod_name = args['name']
targets = args['targets']
for target in targets:
client = root.get(target)
client.dispatch(mod_name)
self.reply('return', {'status': 0, 'cmd': 'dispatch', 'target': target})
except:
self.reply('return', {'status': -1, 'cmd': 'dispatch'})
elif cmd == 'get_status':
try:
client = root.get(args['target'])
job = client.find_job(args['job_id'])
self.reply('status_update', {'status': job.status()})
self.reply('return', {'status': 0, 'cmd': 'get_status'})
except:
self.reply('return', {'status': -1, 'cmd': 'get_status'})
elif cmd == 'get_clients':
try:
clients = root.get_uuids()
self.reply('client_list', {'clients': clients})
self.reply('return', {'status': 0, 'cmd': 'get_clients'})
except:
self.reply('return', {'status': -1, 'cmd': 'get_clients'})
elif cmd == 'get_jobs':
try:
client = root.get(args['target'])
jobs = [{'job_id': i.id, 'name': i.name, 'modules': i.modules, 'alive': i.alive, 'ret': i.ret, 'pid': i.pid} for i in client.jobs]
self.reply('job_list', {'jobs': jobs, 'target': client.uuid})
self.reply('return', {'status': 0, 'cmd': 'get_clients'})
except:
self.reply('return', {'status': -1, 'cmd': 'get_clients'})
# Replies from dispatch
elif cmd == 'dispatched':
try:
mod_name = args['name']
job_id = args['job_id']
pid = args['pid']
self.jobs.append(Job(mod_name, pid, job_id))
except:
pass
elif cmd == 'dispatch_failed':
try:
mod_name = args['name']
print('[JOB %s-?] Dispatch of %s failed' % (self.uuid, mod_name))
except:
pass
elif cmd == 'status_update':
try:
job = self.find_job(args['job_id'])
if job != None:
status = args['status']
job.out = status[0]
job.err = status[1]
job.modules = status[2]
job.alive = status[3]
job.ret = status[4]
if not job.alive:
self.print_status(job)
except:
pass
else:
print('[DISPATCHER] Unknown blob:', obj)
def receive(self):
try:
obj = self.stack.poll()
if obj != None:
self.handle(obj)
return True
except:
return False
