Python ThreadPool.start方法代码示例

# 需要导入模块: from threadpool import ThreadPool [as 别名]
# 或者: from threadpool.ThreadPool import start [as 别名]
class ThreadedRunner(BaseRunner):
    "Run tests using a threadpool"

    def __init__(self, num_threads):
        BaseRunner.__init__(self)
        from threadpool import ThreadPool
        self.pool = ThreadPool(num_threads)
        self.pool.start()

    def run(self, fixture):
        BaseRunner.run(self, fixture)
        self.pool.dispatch(fixture, self.reporter)

    def done(self):
        self.pool.stop()
        BaseRunner.done(self)

# 需要导入模块: from threadpool import ThreadPool [as 别名]
# 或者: from threadpool.ThreadPool import start [as 别名]
class Scanner:
	def __init__( self, kb, cfg, targets, edispatcher ):
		self.kb      = kb
		self.cfg     = cfg
		self.targets = targets
		self.pool	 = ThreadPool( window_size = self.cfg.Threads, prototype = ScannerThread, async = False )
		self.ed		 = edispatcher
		
	def start( self ):
		for target in self.targets:
			for kbitem in self.kb.items:
				self.pool.pushArgs( kbitem, target, self.ed )
				
		self.pool.start()
		
	def stop( self ):
		self.pool.stop()
		
	def running( self ):
		return self.pool.active

# 需要导入模块: from threadpool import ThreadPool [as 别名]
# 或者: from threadpool.ThreadPool import start [as 别名]
	def run( self ):
			if self.kbitem.id in ( 'files', 'dirs' ):
				resp404 = None 
				try:
					non_existent_url      = target.url.copy()
					non_existent_url.path =  "/%s" % self.__genRandom()
					non_existent          = GetRequest(non_existent_url)
					resp404				  = non_existent.fetch()
				except:
					pass
				
				pool = ThreadPool( window_size = 20, prototype = HttpVerifierThread )
				for payload in self.kbitem.payloads:
					pool.pushArgs( self.kbitem, payload, self.target, self.ed, resp404 )
					
				pool.start()
			else:
				for payload in self.kbitem.payloads:
					try:
						random = self.__genRandom()
						if (payload.scope == '*' or payload.scope.lower() == 'get') and isinstance( self.target, GetRequest ):
							for param in self.target.url.params.keys():
								target = self.target.copy()
								target.__class__ = GetRequest
								p 	   = payload.copy()
								if "@RANDOM" in p.data:
									p.data = p.data.replace( "@RANDOM", random )
									
								target.setParam( param, p.data )
								
								response = target.fetch()
								
								for m in self.kbitem.matches:
									m = m.copy()
									if "@PAYLOAD" in m.data:
										m.data = m.data.replace( "@PAYLOAD", p.data )
									if "@RANDOM" in m.data:
										m.data = m.data.replace( "@RANDOM", random )
									if m.match(response):
										self.ed.vulnerability( target, self.kbitem, param )
										return
						elif (payload.scope == '*' or payload.scope.lower() == 'post') and isinstance( self.target, PostRequest ):
							for field in self.target.fields.keys():
								target = self.target.copy()
								target.__class__ = PostRequest
								p 	   = payload.copy()
								if "@RANDOM" in p.data:
									p.data = p.data.replace( "@RANDOM", random )
									
								target.setField( field, p.data )
								
								response = target.fetch()
								
								for m in self.kbitem.matches:
									m = m.copy()
									if "@PAYLOAD" in m.data:
										m.data = m.data.replace( "@PAYLOAD", p.data )
									if "@RANDOM" in m.data:
										m.data = m.data.replace( "@RANDOM", random )
									if m.match(response):
										self.ed.vulnerability( target, self.kbitem, field )
										return
							for param in self.target.url.params.keys():
								target = self.target.copy()
								target.__class__ = PostRequest
								p 	   = payload.copy()
								if "@RANDOM" in p.data:
									p.data = p.data.replace( "@RANDOM", random )
									
								target.setParam( param, p.data )
								
								response = target.fetch()
								
								for m in self.kbitem.matches:
									m = m.copy()
									if "@PAYLOAD" in m.data:
										m.data = m.data.replace( "@PAYLOAD", p.data )
									if "@RANDOM" in m.data:
										m.data = m.data.replace( "@RANDOM", random )
									if m.match(response):
										self.ed.vulnerability( target, self.kbitem, param )
										return
						elif payload.scope.lower() == 'header':
							target 	   = self.target.copy()
							connection = httplib.HTTPConnection( target.url.netloc )
							connection.request( "HEAD", "/" )
							response = connection.getresponse()
							headers  = response.getheaders()
							for header in headers:
								for m in self.kbitem.matches:
									if m.match(header[1]):
										self.ed.vulnerability( target, self.kbitem, None )
										return
					# except HTTPError as e:
					#	self.ed.warning(e)
					except:
						pass

# 需要导入模块: from threadpool import ThreadPool [as 别名]
# 或者: from threadpool.ThreadPool import start [as 别名]
class ThriftThreadServer:
    def __init__(self, module, handler_class, addr, max_process=1, max_conn=50, framed=True):
        module.handler = handler_class()
        global service
        service = module
        self.framed = framed

        self.workers = []
        self.running = True

        def signal_master_handler(signum, frame):
            log.warn("signal %d catched in master %d, wait for kill all worker", signum, os.getpid())
            self.running = False
            for p in self.workers:
                p.terminate()
        
        def signal_worker_handler(signum, frame):
            log.warn("worker %d will exit after all request handled", os.getpid())
            #self.server.close()
            self.running = False
            self.sock.close()

        def server_init():
            signal.signal(signal.SIGTERM, signal_worker_handler)
            log.warn('server started addr=%s:%d pid=%d', addr[0], addr[1], os.getpid())
            if hasattr(service.handler, '_initial'):
                service.handler._initial()
            #self.server.serve_forever()

        def _start_process(index):
            server_name = 'proc-%02d' % index
            p = multiprocessing.Process(target=server_start, name=server_name)
            p.start()
            return p
        
        def signal_child_handler(signum, frame):
            time.sleep(1)
            if self.running:
                log.warn("master recv worker exit, fork one")
                try:
                    pinfo = os.waitpid(-1, 0)
                    pid = pinfo[0]

                    index = -1
                    for i in range(0, len(self.workers)):
                        p = self.workers[i]
                        if p.pid == pid:
                            index = i
                            break
                   
                    if index >= 0:
                        self.workers[index] = _start_process(index)
                except OSError:
                    log.info('waitpid error:')

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(addr)
        self.sock.listen(1024)

        server_init()

        self.tp = ThreadPool(max_conn, max_conn*5) 
        self.tp.start()


    def forever(self):
        class MyTask (Task):
            def run(self):
                return self._func(*self._args, **self._kwargs)

        while self.running:
            try:
                client, addr = self.sock.accept()
                self.tp.add(MyTask(handle, client=client, addr=addr, framed=self.framed))
            except KeyboardInterrupt:
                break
                #os.kill(os.getpid(), 9)
            except Exception, e:
                log.warn(traceback.format_exc())
                log.warn('master exception: %s', str(e))

        self.sock.close()
        self.sock = None
        self.tp.stop_wait()

        log.warn('master exit')