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