def testBasic(self):
path = os.path.join(modulePath(), 'testtasks', 'Succesful.py')
runner = TaskRunner()
runner.start()
out_q = runner.taskStart([path, "Foo"])
txt = out_q.get(True, 1)
self.assertEquals(txt, "FLAG: foobarbaz")
self.assertEquals(len(out_q.get(True, 1)), 0)
self.assertFalse(runner.taskAlive())
self.assertEquals(0, runner.taskRetcode())
runner.quit()
def testBasic(self): path = os.path.join(modulePath(),'testtasks','Succesful.py') runner = TaskRunner() runner.start() out_q = runner.taskStart([path,"Foo"]) txt = out_q.get(True,1) self.assertEquals(txt,"FLAG: foobarbaz") self.assertEquals(len(out_q.get(True,1)),0) self.assertFalse(runner.taskAlive()) self.assertEquals(0,runner.taskRetcode()) runner.quit()
class ServiceTask():
def __init__(self,script,team,logger=None):
self.logger = logger
if(self.logger != None):
self.logger.debug("Creating Service Task: %r %r" % (team.host,script))
thread_size = threading.stack_size(524288)
self.runner = TaskRunner()
self.runner.start()
self.script = script
self.team = team
self.cookie_txt = None
self.error_txt = None
self.prev_flag = None
self.current_flag = None
self.expected_flag = None
self.line_q = None
threading.stack_size(thread_size)
def launch(self,flag):
cmd = [self.script,self.team.host,flag]
if(self.cookie_txt != None):
cmd.append(self.cookie_txt)
self.expected_flag = self.current_flag
self.current_flag = flag
self.prev_flag = None
self.error_txt = None
if(self.logger != None):
self.logger.debug("Service Task starting for %r (%r)" % (self.team.host,flag))
self.line_q = self.runner.taskStart(cmd)
def isAlive(self):
return self.runner.taskAlive()
def finish(self):
self.runner.taskStop()
if(self.logger != None):
self.logger.debug("Service Task finished for %r" % (self.team.host))
def prevFlag(self):
return self.prev_flag
def retcode(self):
return self.runner.taskRetcode()
def gotExpectedFlag(self):
if(self.prev_flag == None or self.expected_flag == None):
return False
return self.prev_flag == self.expected_flag
def cookie(self):
return self.cookie_txt
def error(self):
return self.error_txt
def quit(self):
self.runner.quit()
def processOutput(self):
assert(self.line_q != None)
for i in xrange(self.line_q.qsize()):
line = self.line_q.get(False)
assert(line != None)
if(line == ""):
continue
if(self.logger != None):
self.logger.debug("(%s) %r" % (self.team.host,line.strip()))
if(line.startswith("FLAG:")):
self.prev_flag = line[5:].strip()
elif(line.startswith("ERROR:")):
self.error_txt = line[6:].strip()
elif(line.startswith("COOKIE:")):
self.cookie_txt = line[7:].strip()
