queued_submit.py

import queue
import threading
from submitter import CustomSubmitter
import time


class QueuedSubmitter(threading.Thread):
    def __init__(self, settings, _globals):
        super().__init__()
        self.queue = queue.Queue()
        self.submit_url = settings.submit_url
        self.submit_token = settings.submit_token
        self.submitter = CustomSubmitter(settings, _globals)

    def enqueue_flag(self, flag):
        self.queue.put(flag)

    def run(self):
        while True:
            n_flags = 0
            flags = set()
            while not self.queue.empty():
                flag = self.queue.get()
                flags.add(flag)
                n_flags += 1
            flags = list(flags)
            self.submitter.submit(flags, self.submit_url, self.submit_token)
            # task_done useful for future purposes in which attackers join the submitter
            for _ in range(n_flags):
                self.queue.task_done()
            if n_flags > 0:
                print("[+] {} flags submitted by the queued submitter".format(n_flags))
            time.sleep(2)


if __name__ == "__main__":
    # test
    q = queue.Queue()

    def producer(_id):
        while True:
            q.put("msg_{}".format(_id))
            time.sleep(1)

    def consumer():
        while True:
            n_msg = 0
            msg = []
            while not q.empty():
                m = q.get()
                msg.append(m)
                n_msg += 1
            print(" ".join(msg))
            for _ in range(n_msg):
                q.task_done()
            print("Number of messages: {}".format(n_msg))
            time.sleep(2)

    n_threads, threads = 5, []
    for i in range(n_threads):
        t = threading.Thread(target=producer, args=[i+1])
        t.start()
        threads.append(t)
    t = threading.Thread(target=consumer, args=[])
    t.start()
    threads.append(t)
    for t in threads:
        t.join()