Source code for pysap.utils

# encoding: utf-8
# pysap - Python library for crafting SAP's network protocols packets
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# GNU General Public License for more details.
# Author:
#   Martin Gallo (@martingalloar)
#   Code contributed by SecureAuth to the OWASP CBAS project

# Standard imports
from Queue import Queue
from threading import Thread, Event

[docs]class Worker(Thread): """Thread Worker It runs a function into a new thread. """ def __init__(self, decoder, function): Thread.__init__(self) self.decoder = decoder self.function = function self.stopped = Event()
[docs] def run(self): while not self.stopped.is_set(): self.function()
[docs] def stop(self): self.stopped.set()
# Simple Thread Pool implementation based on # WorkerQueue class
[docs]class WorkerQueue(Thread): """Thread executing tasks from a given tasks queue""" def __init__(self, tasks): Thread.__init__(self) self.tasks = tasks self.daemon = True self.start()
[docs] def run(self): while True: func, args, kargs = self.tasks.get() try: func(*args, **kargs) except Exception as e: print(e) self.tasks.task_done()
# ThreadPool class
[docs]class ThreadPool(object): """Pool of threads consuming tasks from a queue""" def __init__(self, num_threads): self.tasks = Queue(num_threads) for _ in range(num_threads): WorkerQueue(self.tasks)
[docs] def add_task(self, func, *args, **kargs): """Add a task to the queue""" self.tasks.put((func, args, kargs))
[docs] def wait_completion(self): """Wait for completion of all the tasks in the queue""" self.tasks.join()