import multiprocessing as mp import logging as log import os import changeme import time from runner import runner from submitter import do_submissions from changeme import submit_flag, 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()