main.py

import multiprocessing as mp
import logging as log
import os
import changeme
import time

from runner import runner
from submitter import do_submissions, submit_flag
from changeme import all_teams

EXPLOIT_DIR = changeme.exploit_dir
runners = []
submitter = None


log.basicConfig(level=log.DEBUG)

def is_executable(path):
    return bool(os.stat(path).st_mode & 0o100)

def find_exploits():
    exps = []
    for p in os.listdir(EXPLOIT_DIR):
        if is_executable(os.path.join(EXPLOIT_DIR, p)):
            exps.append(p)
    return exps


def do_scheduling(work_queue, flag_queue):
    while True:
        exploits = find_exploits()
        log.info('Found %d different exploits', len(exploits))
        for team_id in all_teams:
            for exploit in exploits:
                work_queue.put({
                    "team_id": team_id,
                    "exploit": exploit
                })
        if changeme.wait_between_runs:
            time.sleep(changeme.wait_between_runs)

def main():
    procs = changeme.proc_num

    flag_queue = mp.Queue()
    work_queue = mp.Queue()

    log.info("Launching %d exploit runners", procs)
    for i in range(procs):
        p = mp.Process(target=runner, args=(i, work_queue, flag_queue))
        p.start()
        runners.append(p)

    log.info("Launching submission worker")
    submitter = mp.Process(target=do_submissions, args=(flag_queue,))
    submitter.start()

    do_scheduling(work_queue, flag_queue)


if __name__ == '__main__':
    main()