def __init__(self, consumer_plugins, w3af_core, thread_name, create_pool=True, max_pool_queued_tasks=0): """ :param consumer_plugins: Instances of base_consumer plugins in a list :param w3af_core: The w3af core that we'll use for status reporting :param thread_name: How to name the current thread :param create_pool: True to create a worker pool for this consumer """ super(BaseConsumer, self).__init__(name='%sController' % thread_name) self.in_queue = QueueSpeed() self._out_queue = Queue.Queue() self._consumer_plugins = consumer_plugins self._w3af_core = w3af_core self._observers = [] self._tasks_in_progress = {} self._poison_pill_sent = False self._threadpool = None if create_pool: self._threadpool = Pool(10, worker_names='%sWorker' % thread_name, max_queued_tasks=max_pool_queued_tasks)
def test_many_items(self): q = QueueSpeed() self.assertEqual(len(q._input_data), 0) for _ in xrange(q.MAX_SIZE * 2): q.put(None) self.assertEqual(len(q._input_data), q.MAX_SIZE - 1) self.assertEqual(len(q._output_data), 0) for _ in xrange(q.MAX_SIZE * 2): q.get() self.assertEqual(len(q._output_data), q.MAX_SIZE - 1)
def __init__(self, consumer_plugins, w3af_core, thread_name, create_pool=True, max_pool_queued_tasks=0, max_in_queue_size=0): """ :param consumer_plugins: Instances of base_consumer plugins in a list :param w3af_core: The w3af core that we'll use for status reporting :param thread_name: How to name the current thread :param create_pool: True to create a worker pool for this consumer """ super(BaseConsumer, self).__init__(name='%sController' % thread_name) self.in_queue = QueueSpeed(maxsize=max_in_queue_size) self._out_queue = Queue.Queue() self._consumer_plugins = consumer_plugins self._w3af_core = w3af_core self._observers = [] self._tasks_in_progress = {} self._poison_pill_sent = False self._threadpool = None if create_pool: self._threadpool = Pool(self.THREAD_POOL_SIZE, worker_names='%sWorker' % thread_name, max_queued_tasks=max_pool_queued_tasks)
def test_wrapper(self): q = QueueSpeed(4) self.assertEqual(q.qsize(), 0) q.put(None) self.assertEqual(q.qsize(), 1) q.get() self.assertEqual(q.qsize(), 0)
def __init__(self, consumer_plugins, w3af_core, thread_name, create_pool=True): """ :param base_consumer_plugins: Instances of base_consumer plugins in a list :param w3af_core: The w3af core that we'll use for status reporting :param thread_name: How to name the current thread :param create_pool: True to create a worker pool for this consumer """ super(BaseConsumer, self).__init__(name='%sController' % thread_name) self.in_queue = QueueSpeed() self._out_queue = Queue.Queue() self._consumer_plugins = consumer_plugins self._w3af_core = w3af_core self._tasks_in_progress = {} self._threadpool = None if create_pool: self._threadpool = Pool(10, worker_names='%sWorker' % thread_name)
class BaseConsumer(Process): """ Consumer thread that takes fuzzable requests from a Queue that's populated by the crawl plugins and identified vulnerabilities by performing various requests. """ def __init__(self, consumer_plugins, w3af_core, thread_name, create_pool=True): """ :param base_consumer_plugins: Instances of base_consumer plugins in a list :param w3af_core: The w3af core that we'll use for status reporting :param thread_name: How to name the current thread :param create_pool: True to create a worker pool for this consumer """ super(BaseConsumer, self).__init__(name='%sController' % thread_name) self.in_queue = QueueSpeed() self._out_queue = Queue.Queue() self._consumer_plugins = consumer_plugins self._w3af_core = w3af_core self._tasks_in_progress = {} self._threadpool = None if create_pool: self._threadpool = Pool(10, worker_names='%sWorker' % thread_name) def run(self): """ Consume the queue items, sending them to the plugins which are then going to find vulnerabilities, new URLs, etc. """ while True: work_unit = self.in_queue.get() if work_unit == POISON_PILL: # Close the pool and wait for everyone to finish self._threadpool.close() self._threadpool.join() del self._threadpool self._teardown() # Finish this consumer and everyone consuming the output self._out_queue.put(POISON_PILL) self.in_queue.task_done() break else: # pylint: disable=E1120 self._consume_wrapper(work_unit) self.in_queue.task_done() def _teardown(self): raise NotImplementedError def _consume(self, work_unit): raise NotImplementedError @task_decorator def _consume_wrapper(self, function_id, work_unit): """ Just makes sure that all _consume methods are decorated as tasks. """ return self._consume(work_unit) def _task_done(self, function_id): """ The task_in_progress_counter is needed because we want to know if the consumer is processing something and let it finish. It is mainly used in the has_pending_work(). For example: * You can have pending work if there are items in the input_queue * You can have pending work if there are still items to be read from the output_queue by one of the consumers that reads our output. * You can have pending work when there are no items in input_queue and no items in output_queue but the threadpool inside the consumer is processing something. This situation is handled by the self._tasks_in_progress attribute and the _add_task and _task_done methods. So, for each _add_task() there has to be a _task_done() even if the task ends in an error or exception. Recommendation: Do NOT set the callback for apply_async to call _task_done, the Python2.7 pool implementation won't call it if the function raised an exception and you'll end up with tasks in progress that finished with an exception. """ try: self._tasks_in_progress.pop(function_id) except KeyError: raise AssertionError('The function %s was not found!' % function_id) def _add_task(self, function_id): """ :param function_id: Just for debugging @see: _task_done()'s documentation. """ self._tasks_in_progress[function_id] = 1 def in_queue_put(self, work): if work is not None: return self.in_queue.put(work) def in_queue_put_iter(self, work_iter): if work_iter is not None: for work in work_iter: self.in_queue_put(work) def has_pending_work(self): """ @see: _task_done() documentation :return: True if the in_queue_size is != 0 OR if one of the pool workers is still doing something that might impact on out_queue. """ if self.in_queue_size() > 0 \ or self.out_queue.qsize() > 0: return True if len(self._tasks_in_progress) > 0: return True # This is a special case which loosely translates to: "If there are any # pending tasks in the threadpool, even if they haven't yet called the # _add_task method, we know we have pending work to do". if hasattr(self, '_threadpool') and self._threadpool is not None: if self._threadpool._inqueue.qsize() > 0 \ or self._threadpool._outqueue.qsize() > 0: return True return False @property def out_queue(self): # # This output queue can contain one of the following: # * POISON_PILL # * (plugin_name, fuzzable_request, AsyncResult) # * An ExceptionData instance return self._out_queue def in_queue_size(self): return self.in_queue.qsize() def join(self): """ Poison the loop and wait for all queued work to finish this might take some time to process. """ if not self.is_alive(): # This return has a long history, follow it here: # https://github.com/andresriancho/w3af/issues/1172 return self.in_queue_put(POISON_PILL) self.in_queue.join() def terminate(self): """ Remove all queued work from in_queue and poison the loop so the consumer exits. Should be very fast and called only if we don't care about the queued work anymore (ie. user clicked stop in the UI). """ while not self.in_queue.empty(): self.in_queue.get() self.in_queue.task_done() self.join() def get_result(self, timeout=0.5): """ :return: The first result from the output Queue. """ return self._out_queue.get(timeout=timeout) def handle_exception(self, phase, plugin_name, fuzzable_request, _exception): """ Get the exception information, and put it into the output queue then, the strategy will get the items from the output queue and handle the exceptions. :param plugin_name: The plugin that generated the exception :param fuzzable_request: The fuzzable request that was sent as input to the plugin when the exception was raised :param _exception: The exception object """ except_type, except_class, tb = sys.exc_info() enabled_plugins = pprint_plugins(self._w3af_core) status = w3af_core_status(self._w3af_core) status.set_running_plugin(phase, plugin_name, log=False) status.set_current_fuzzable_request(phase, fuzzable_request) exception_data = ExceptionData(status, _exception, tb, enabled_plugins) self._out_queue.put(exception_data)
class BaseConsumer(Process): """ Consumer thread that takes fuzzable requests from a Queue that's populated by the crawl plugins and identified vulnerabilities by performing various requests. """ def __init__(self, consumer_plugins, w3af_core, thread_name, create_pool=True, max_pool_queued_tasks=0): """ :param consumer_plugins: Instances of base_consumer plugins in a list :param w3af_core: The w3af core that we'll use for status reporting :param thread_name: How to name the current thread :param create_pool: True to create a worker pool for this consumer """ super(BaseConsumer, self).__init__(name='%sController' % thread_name) self.in_queue = QueueSpeed() self._out_queue = Queue.Queue() self._consumer_plugins = consumer_plugins self._w3af_core = w3af_core self._observers = [] self._tasks_in_progress = {} self._poison_pill_sent = False self._threadpool = None if create_pool: self._threadpool = Pool(10, worker_names='%sWorker' % thread_name, max_queued_tasks=max_pool_queued_tasks) def run(self): """ Consume the queue items, sending them to the plugins which are then going to find vulnerabilities, new URLs, etc. """ while True: try: work_unit = self.in_queue.get() except KeyboardInterrupt: # https://github.com/andresriancho/w3af/issues/9587 # # If we don't do this, the thread will die and will never # process the POISON_PILL, which will end up in an endless # wait for .join() continue if work_unit == POISON_PILL: # Close the pool and wait for everyone to finish self._threadpool.close() self._threadpool.join() del self._threadpool self._teardown() # Finish this consumer and everyone consuming the output self._out_queue.put(POISON_PILL) self.in_queue.task_done() break else: # pylint: disable=E1120 try: self._consume_wrapper(work_unit) finally: self.in_queue.task_done() def _teardown(self): raise NotImplementedError def _consume(self, work_unit): raise NotImplementedError @task_decorator def _consume_wrapper(self, function_id, work_unit): """ Just makes sure that all _consume methods are decorated as tasks. """ return self._consume(work_unit) def _task_done(self, function_id): """ The task_in_progress_counter is needed because we want to know if the consumer is processing something and let it finish. It is mainly used in the has_pending_work(). For example: * You can have pending work if there are items in the input_queue * You can have pending work if there are still items to be read from the output_queue by one of the consumers that reads our output. * You can have pending work when there are no items in input_queue and no items in output_queue but the threadpool inside the consumer is processing something. This situation is handled by the self._tasks_in_progress attribute and the _add_task and _task_done methods. So, for each _add_task() there has to be a _task_done() even if the task ends in an error or exception. Recommendation: Do NOT set the callback for apply_async to call _task_done, the Python2.7 pool implementation won't call it if the function raised an exception and you'll end up with tasks in progress that finished with an exception. """ try: self._tasks_in_progress.pop(function_id) except KeyError: raise AssertionError('The function %s was not found!' % function_id) def _add_task(self, function_id): """ :param function_id: Just for debugging @see: _task_done()'s documentation. """ self._tasks_in_progress[function_id] = 1 def in_queue_put(self, work): if work is not None: return self.in_queue.put(work) def in_queue_put_iter(self, work_iter): if work_iter is not None: for work in work_iter: self.in_queue_put(work) def has_pending_work(self): """ @see: _task_done() documentation :return: True if the in_queue_size is != 0 OR if one of the pool workers is still doing something that might impact on out_queue. """ if self.in_queue_size() > 0 \ or self.out_queue.qsize() > 0: return True if len(self._tasks_in_progress) > 0: return True # This is a special case which loosely translates to: "If there are any # pending tasks in the threadpool, even if they haven't yet called the # _add_task method, we know we have pending work to do". if hasattr(self, '_threadpool') and self._threadpool is not None: if self._threadpool._inqueue.qsize() > 0 \ or self._threadpool._outqueue.qsize() > 0: return True return False @property def out_queue(self): # This output queue can contain one of the following: # * POISON_PILL # * (plugin_name, fuzzable_request, AsyncResult) # * An ExceptionData instance return self._out_queue def in_queue_size(self): return self.in_queue.qsize() def join(self): """ Poison the loop and wait for all queued work to finish this might take some time to process. """ if not self.is_alive(): # This return has a long history, follow it here: # https://github.com/andresriancho/w3af/issues/1172 return if not self._poison_pill_sent: # https://github.com/andresriancho/w3af/issues/9587 self._poison_pill_sent = True self.in_queue_put(POISON_PILL) self.in_queue.join() def terminate(self): """ Remove all queued work from in_queue and poison the loop so the consumer exits. Should be very fast and called only if we don't care about the queued work anymore (ie. user clicked stop in the UI). """ while not self.in_queue.empty(): self.in_queue.get() self.in_queue.task_done() self.join() def get_result(self, timeout=0.5): """ :return: The first result from the output Queue. """ return self._out_queue.get(timeout=timeout) def handle_exception(self, phase, plugin_name, fuzzable_request, _exception): """ Get the exception information, and put it into the output queue then, the strategy will get the items from the output queue and handle the exceptions. :param plugin_name: The plugin that generated the exception :param fuzzable_request: The fuzzable request that was sent as input to the plugin when the exception was raised :param _exception: The exception object """ except_type, except_class, tb = sys.exc_info() enabled_plugins = pprint_plugins(self._w3af_core) status = w3af_core_status(self._w3af_core) status.set_running_plugin(phase, plugin_name, log=False) status.set_current_fuzzable_request(phase, fuzzable_request) exception_data = ExceptionData(status, _exception, tb, enabled_plugins) self._out_queue.put(exception_data) def add_observer(self, observer): self._observers.append(observer)
def test_exceptions(self): q = QueueSpeed(4) self.assertEqual(None, q.get_input_rpm()) self.assertEqual(None, q.get_output_rpm()) for i in xrange(4): q.put(i) # 20 RPM time.sleep(3) for _ in xrange(10): self.assertRaises(Queue.Full, q.put_nowait, None) self.assertEqual(q.qsize(), 4) self.assertGreater(q.get_input_rpm(), 19) self.assertLess(q.get_input_rpm(), 20) for i in xrange(4): q.get() # 60 RPM time.sleep(1) for _ in xrange(10): self.assertRaises(Queue.Empty, q.get_nowait) self.assertGreater(q.get_output_rpm(), 59) self.assertLess(q.get_output_rpm(), 60) self.assertEqual(q.qsize(), 0)
def test_no_data(self): q = QueueSpeed() for _ in xrange(10): self.assertEqual(None, q.get_input_rpm()) self.assertEqual(None, q.get_output_rpm())
def test_simple(self): q = QueueSpeed() self.assertEqual(None, q.get_input_rpm()) self.assertEqual(None, q.get_output_rpm()) for i in xrange(4): q.put(i) # 20 RPM time.sleep(3) self.assertEqual(q.qsize(), 4) self.assertGreater(q.get_input_rpm(), 19) self.assertLess(q.get_input_rpm(), 20) for i in xrange(4): q.get() # 60 RPM time.sleep(1) self.assertGreater(q.get_output_rpm(), 59) self.assertLess(q.get_output_rpm(), 60) self.assertEqual(q.qsize(), 0)