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

REMAINING_FILE = "remaining_flags.json"
EXPLOIT_DIR = changeme.exploit_dir
runners = []
submitter = None


log.basicConfig(level=log.DEBUG)

def submit_remaining():
    """
    Submit flags from an earlier run that weren't submitted.
    """
    if os.path.isfile(REMAINING_FILE):
        with open(REMAINING_FILE) as f:
            log.info("There seem to be remaining flags from a previous run")
            log.info("Submitting flags from a previous run")
            flags = json.load(f)
            for team_id, flags in flags.items():
                for flag in flags:
                    submit_flag(team_id, flag)
            os.unlink(REMAINING_FILE)


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():
    submit_remaining()
    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()