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 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 DispatchLoader(object):
def __init__(self, ip, port, job):
self.stack = Stack((StackableSocket(ip=ip, port=port),
StackablePacketAssembler(),
StackablePickler()))
self.job = job
self.cache = {}
self.path = ''
self.import_lock = threading.Lock()
def get_module(self, fullname):
if fullname in self.cache:
return self.cache[fullname]
else:
with self.import_lock:
self.stack.write({'load': fullname, 'id': self.job})
o = self.stack.read()
mod = o['module']
if mod != None:
if mod['type'] != "python":
return None
self.cache[fullname] = mod
return mod
return None
def exec_module(self, module, _globals):
if module['bytecode']:
exec module['bytecode'] in _globals
else:
exec module['source'] in _globals
def execute(self, fullname, _globals):
return self.exec_module(self.get_module(fullname), _globals)
def find_module(self, fullname, path=None):
if self.get_module(fullname) != None:
self.path = path
return self
return None
def load_module(self, fullname):
if fullname in sys.modules:
return sys.modules[fullname]
mod = self.cache[fullname]
m = sys.modules[fullname] = types.ModuleType(fullname, fullname)
self.exec_module(mod, m.__dict__)
return m
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
def __init__(self, ip, port, job):
self.stack = Stack((StackableSocket(ip=ip, port=port),
StackablePacketAssembler(),
StackablePickler()))
self.job = job
self.cache = {}
self.path = ''
self.import_lock = threading.Lock()
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'])
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 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
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)
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 connect(self, ip, port): 'Connect' self.stack = Stack((StackableSocket(ip=ip, port=port), StackablePacketAssembler(), StackablePickler()))
class DispatchManager(object):
universal_handlers = {
'shell': 'sh',
'ruby': 'ruby'
}
def __init__(self, ip=None, port=None):
self.processes = []
self.plock = Lock()
self.cache = {}
self.mod_cbs = {}
self.stack = None
self.job_cnt = 0
self.id = str(uuid4())
self.com_id = 0
if ip != None and port != None:
self.connect(ip, port)
# Network boiler-plate
def connect(self, ip, port):
'Connect'
self.stack = Stack((StackableSocket(ip=ip, port=port),
StackablePacketAssembler(),
StackablePickler()))
def reply(self, cmd, args):
'Send without incrementing counter'
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def send(self, cmd, args):
'Send'
self.com_id += 1
self.stack.write({'cmd': cmd, 'args': args, 'com_id': self.com_id})
def handle(self, obj):
'Command parser'
cmd = obj['cmd']
args = obj['args']
if cmd == 'init':
self.id = args['uuid']
print("[DISPATCHER] Setting UUID: %s" % self.id)
elif cmd == 'execute':
self.dispatch(args['name'])
elif cmd == 'kill':
job = self.get_job(args['job_id'])
job.kill()
self.send('status_update', {'job_id': job.id, 'status': job.status()})
self.processes.remove(job)
elif cmd == 'module_update':
mod = None
name = ''
if 'not_found' in args:
name = args['not_found']
elif 'module' in args:
mod = args['module']
name = mod['name']
self.cache[name] = mod
if name in self.mod_cbs:
for i in self.mod_cbs[name]:
i(mod)
elif cmd == 'get_status':
job = self.get_job(args['job_id'])
self.send('status_update', {'job_id': job.id, 'status': job.status()})
elif cmd == 'get_jobs':
jobs = [{'name': job.name, 'job_id': job.id, 'modules': job.modules} for job in self.processes]
self.send('job_update', {'jobs': jobs})
# Callbacks
def job_completed(self, job):
'Callback for completed job'
self.send('status_update', {'job_id': job.id, 'status': job.status()})
self.processes.remove(job)
def job_dispatched(self, job):
'Callback for dispatched job'
self.send('dispatched', {'name': job.name, 'pid': job.pid, 'job_id': job.id})
def job_failed_dispatch(self, name):
'Callback for failed dispatching of job'
self.send('dispatch_failed', {'name': name})
# Monitor threads
def net_monitor(self):
while True:
o = self.stack.read()
self.handle(o)
# External interfaces
def get_job(self, _id):
'Retrieve a job from ID'
for i in self.processes:
if i.id == _id:
return i
def get_module(self, m, cb):
'Retrieve a job - Calls cb when done'
if m in self.cache:
cb(self.cache[m])
return
if m not in self.mod_cbs:
self.mod_cbs[m] = []
self.mod_cbs[m].append(cb)
self.send('get_module', {'name': m})
def dispatch(self, name):
'Dispatch job'
def callback(mod):
print('[DISPATCHER] Spawn of type "%s": %s' % (mod['type'], name))
p = None
if mod['type'] == 'python':
p = Job(name, Popen(['pypy', '-c', '''
import sys, loader
__dispatch__ = loader.DispatchLoader('%s', %d, %d)
sys.meta_path = [__dispatch__]
del loader, sys
__dispatch__.execute('%s', globals())''' % ("localhost", 2001, self.job_cnt, name)], stdout=PIPE, stderr=PIPE), self.job_cnt)
elif mod['type'] in self.universal_handlers:
t = NamedTemporaryFile(delete=False)
t.write(mod['source'])
t.close()
p = Job(name, Popen([self.universal_handlers[mod['type']], t.name], stdout=PIPE, stderr=PIPE), self.job_cnt)
p.modules.append(name)
if p != None:
with self.plock:
self.job_cnt += 1
self.processes.append(p)
p.monitor(self.job_completed)
self.job_dispatched(p)
else:
print('[DISPATCHER] Type "%s" not supported' % mod['type'])
self.job_failed_dispatch(name)
self.get_module(name, callback)
def init(self): self.stack = Stack((StackableSocket(sock=self.conn), StackablePacketAssembler(), StackablePickler()))
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
