Skip to content
Snippets Groups Projects
main.py 2.04 KiB
Newer Older
Ruben Anthony Gonzalez's avatar
Ruben Anthony Gonzalez committed
import multiprocessing as mp
import logging as log
import os
import config
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 = config.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 config.wait_between_runs:
            time.sleep(config.wait_between_runs)

def main():
    submit_remaining()
    procs = config.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()