Rooba
/
pupy
mirror of https://github.com/n1nj4sec/pupy.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Rework interruptions

This commit is contained in:
Oleksii Shevchuk 2018-03-11 23:33:09 +02:00
parent a9774a3c67
commit 3b54958c59
2 changed files with 123 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

93
pupy/pupylib/PupyCmd.py
View File

@ -601,7 +601,7 @@ class PupyCmd(cmd.Cmd):
j.stop() j.stop()
self.display_success("job closed") self.display_success("job closed")
else: else:
j.interrupt(wait=False) j.interrupt()
j.stop() j.stop()
self.display_success("job killed") self.display_success("job killed")
self.pupsrv.del_job(modargs.kill) self.pupsrv.del_job(modargs.kill)
@ -763,51 +763,43 @@ class PupyCmd(cmd.Cmd):
modjobs=[x for x in self.pupsrv.jobs.itervalues() if x.pupymodules[0].get_name() == mod.get_name() and x.pupymodules[0].client in l] modjobs=[x for x in self.pupsrv.jobs.itervalues() if x.pupymodules[0].get_name() == mod.get_name() and x.pupymodules[0].client in l]
pj=None pj=None
try:
interactive=False
if mod.daemon and mod.unique_instance and modjobs:
pj=modjobs[0]
else:
pj=PupyJob(self.pupsrv, '{} {}'.format(modargs.module, ' '.join(args)))
if len(l)==1 and not modargs.bg and not mod.daemon:
ps=mod(l[0], pj, stdout=self.stdout, log=modargs.output)
pj.add_module(ps)
interactive=True
else:
for c in l:
ps=mod(c, pj, log=modargs.output)
pj.add_module(ps)
try:
pj.start(args, once=modargs.once)
except Exception as e:
self.display_error(e)
pj.stop()
if not mod.unique_instance: interactive=False
if modargs.bg: if mod.daemon and mod.unique_instance and modjobs:
self.pupsrv.add_job(pj) pj=modjobs[0]
self.display_info("job %s started in background !"%pj) else:
elif mod.daemon: pj=PupyJob(self.pupsrv, '{} {}'.format(modargs.module, ' '.join(args)))
self.pupsrv.add_job(pj) if len(l)==1 and not modargs.bg and not mod.daemon:
self.display_info("job %s started in background !"%pj) ps=mod(l[0], pj, stdout=self.stdout, log=modargs.output)
else: pj.add_module(ps)
error=pj.interactive_wait() interactive=True
if error and not modjobs: else:
pj.stop() for c in l:
ps=mod(c, pj, log=modargs.output)
pj.add_module(ps)
try:
pj.start(args, once=modargs.once)
except Exception as e:
self.display_error(e)
pj.stop()
if not mod.unique_instance:
if modargs.bg:
self.pupsrv.add_job(pj)
self.display_info("job %s started in background !"%pj)
elif mod.daemon:
self.pupsrv.add_job(pj)
self.display_info("job %s started in background !"%pj)
else: else:
if mod.daemon and not modjobs:
self.pupsrv.add_job(pj)
error=pj.interactive_wait() error=pj.interactive_wait()
if error and not modjobs: if error and not modjobs:
pj.stop() pj.stop()
else:
except KeyboardInterrupt: if mod.daemon and not modjobs:
self.display_warning('interrupting job ... (please wait)') self.pupsrv.add_job(pj)
interrupted = pj.interrupt() error=pj.interactive_wait()
if interrupted: if error and not modjobs:
self.display_warning('job was interrupted') pj.stop()
else:
self.display_error('job was sent to background and may consume resources')
if not interactive: if not interactive:
summary = pj.result_summary() summary = pj.result_summary()
@ -1620,6 +1612,25 @@ class PupyCmd(cmd.Cmd):
self.display(PupyCmd.table_format(objects, wl=['IP', 'PORTS'])) self.display(PupyCmd.table_format(objects, wl=['IP', 'PORTS']))
session.open_ports = {} session.open_ports = {}
def do_stacktrace(self, arg):
""" Dump current stacks (debug) """
print >> sys.stderr, "\n*** STACKTRACE - START ***\n"
code = []
for threadId, stack in sys._current_frames().items():
code.append("\n# ThreadID: %s" % threadId)
for filename, lineno, name, line in traceback.extract_stack(stack):
code.append('File: "%s", line %d, in %s' % (filename, lineno, name))
if line:
code.append(" %s" % (line.strip()))
for line in code:
print >> sys.stderr, line
print >> sys.stderr, "\n*** STACKTRACE - END ***\n"
def do_exit(self, arg): def do_exit(self, arg):
""" Quit Pupy Shell """ """ Quit Pupy Shell """
for job in self.pupsrv.jobs.itervalues(): for job in self.pupsrv.jobs.itervalues():

108
pupy/pupylib/PupyJob.py
View File

@ -69,8 +69,8 @@ class ThreadPool(object):
self.thread_pool=[] self.thread_pool=[]
def apply_async(self, func, args): def apply_async(self, func, args):
t=Thread(target=func, args=args) t = Thread(target=func, args=args)
t.daemon=True t.daemon = True
self.thread_pool.append(t) self.thread_pool.append(t)
t.start() t.start()
@ -79,25 +79,39 @@ class ThreadPool(object):
if t.isAlive(): if t.isAlive():
t.stop() t.stop()
def join(self): def join(self, timeout=5, on_interrupt=None):
giveup = False giveup = False
allok = True
interrupt_sent = False
while True: while True:
try: try:
allok=True allok = True
for t in self.thread_pool: for t in self.thread_pool:
if t.isAlive(): if t.isAlive():
t.join(0.5) t.join(timeout)
allok=False allok = False
if allok: if allok:
break break
except KeyboardInterrupt: except KeyboardInterrupt:
if not giveup: if not giveup:
print "{ Press ^C once again to give up on waiting }" if not interrupt_sent:
giveup = True if on_interrupt:
if not on_interrupt():
break
else:
print "{ Press ^C once again to give up on waiting }"
interrupt_sent = True
else:
giveup = True
else: else:
break break
return not giveup and allok
def all_finished(self): def all_finished(self):
for t in self.thread_pool: for t in self.thread_pool:
if t.isAlive(): if t.isAlive():
@ -115,10 +129,11 @@ class PupyJob(object):
self.pupymodules = [] self.pupymodules = []
self.worker_pool = ThreadPool() self.worker_pool = ThreadPool()
self.started = threading.Event() self.started = threading.Event()
self.error_happened = threading.Event() self.error = None
self.jid = None self.jid = None
self.destroyed = False self.destroyed = False
self.id = None self.id = None
self.interrupted = False
def add_module(self, mod): def add_module(self, mod):
self.pupymodules.append(mod) self.pupymodules.append(mod)
@ -130,31 +145,43 @@ class PupyJob(object):
self.interrupt() self.interrupt()
def module_worker(self, module, cmdline, args, once): def module_worker(self, module, cmdline, args, once):
e = None
try: try:
module.import_dependencies() module.import_dependencies()
module.init(cmdline, args) module.init(cmdline, args)
module.run(args) module.run(args)
except PupyModuleExit as e: except PupyModuleExit as e:
self.error = e
return return
except PupyModuleError as e: except PupyModuleError as e:
self.error_happened.set() self.error = e
module.error(str(e)) if not self.interrupted:
module.error(str(e))
except KeyboardInterrupt: except KeyboardInterrupt:
pass pass
except Exception as e: except Exception as e:
self.error_happened.set() self.error = e
module.error(str(e)) if not self.interrupted:
import traceback module.error(str(e))
traceback.print_exc(e, file=module.stdout) import traceback
traceback.print_exc(e, file=module.stdout)
finally: finally:
if once: if not self.interrupted and once:
module.clean_dependencies() module.clean_dependencies()
if self.id is not None: if self.id is not None:
self.pupsrv.handler.display_srvinfo('{}: {} - {}'.format( display = self.pupsrv.handler.display_warning if \
self.id, self.interrupted else self.pupsrv.handler.display_srvinfo
self.name,
'failed' if self.error_happened.is_set() else 'success')) display(
'{}: {} - {}'.format(
self.id, self.name, 'failed' if e else 'success'))
def start(self, args, once=False): def start(self, args, once=False):
#if self.started.is_set(): #if self.started.is_set():
@ -167,59 +194,66 @@ class PupyJob(object):
margs.unknown_args = unknown_args margs.unknown_args = unknown_args
else: else:
margs = m.arg_parser.parse_args(args) margs = m.arg_parser.parse_args(args)
except PupyModuleExit as e: except PupyModuleExit as e:
m.error("Arguments parse error : %s"%e) m.error("Arguments parse error : %s"%e)
continue continue
res = m.is_compatible() res = m.is_compatible()
if type(res) is tuple: if type(res) is tuple:
comp, comp_exp=res comp, comp_exp = res
elif res is None: elif res is None:
comp=True comp = True
comp_exp="" comp_exp = ""
else: else:
comp=res comp=res
comp_exp="reason not precised" comp_exp = "reason not precised"
if not comp: if not comp:
m.error("Compatibility error : %s"%comp_exp) m.error("Compatibility error : %s"%comp_exp)
continue continue
self.worker_pool.apply_async(self.module_worker, (m, args, margs, once)) self.worker_pool.apply_async(self.module_worker, (m, args, margs, once))
self.started.set() self.started.set()
def interrupt(self, wait=True): def interrupt(self):
if not self.started: if not self.started:
raise RuntimeError("can't interrupt. job %s has not been started"%str(self)) raise RuntimeError("can't interrupt. job %s has not been started"%str(self))
if self.interrupted:
return True
self.interrupted = True
#calling the interrupt method is one is defined for the module instead of killing the thread #calling the interrupt method is one is defined for the module instead of killing the thread
if hasattr(self.pupymodules[0],'interrupt'): if hasattr(self.pupymodules[0], 'interrupt'):
self.pupsrv.handler.display_info('Sending interrupt request')
for m in self.pupymodules: for m in self.pupymodules:
m.interrupt() m.interrupt()
return True return True
else: else:
self.pupsrv.handler.display_warning('Module does not support interrupts. Resources may leak!')
self.worker_pool.interrupt_all() self.worker_pool.interrupt_all()
if wait: self.check()
self.wait()
return False return False
def interactive_wait(self): def interactive_wait(self):
while True: self.worker_pool.join(on_interrupt=self.interrupt)
if self.is_finished(): if self.error:
break
time.sleep(0.1)
if self.error_happened.is_set():
return True return True
return False return False
def wait(self): def check(self):
self.worker_pool.join()
for m in self.pupymodules: for m in self.pupymodules:
while True: while True:
if not m.client: if not m.client:
break break
try: try:
m.client.conn._conn.ping(timeout=2, block=True) m.client.conn._conn.ping(timeout=2)
logging.info('connection {} alive'.format(m))
break break
except KeyboardInterrupt: except KeyboardInterrupt: