runner.py

import os
import logging as log
from changeme import team_id_to_ip
import subprocess as sp
import config
import re
import logging as log


def flags_iter_matches(s, flag_regex):
    return re.finditer(b'(?=('+flag_regex+b'))', s)

def flags_iter(s, flag_regex):
    for m in flags_iter_matches(s, flag_regex):
        yield m.group(1)

def runner(runner_id, work_queue, flag_queue):
    log.info("[RUNNER %d] started", runner_id)
    os.chdir(config.exploit_dir)

    while True:
        job = work_queue.get()
        exploit = job['exploit']
        team_id = job['team_id']
        log.info("[RUNNER %d] launching %s against %d", runner_id, exploit, team_id)
        try:
            result = sp.check_output(["./" + exploit, team_id_to_ip(team_id), str(team_id)], stderr=sp.STDOUT)
            for flag in flags_iter(result, config.flag_regex):
                log.info("[RUNNER %d] found flag %s for team %d", runner_id, flag, team_id)
                flag_queue.put({
                    'team_id': team_id,
                    'flag': flag.decode()
                })

        except sp.CalledProcessError as ex:
            log.error("[RUNNER %d] Exploit %s failed at team %d with error code %d", runner_id, exploit, team_id, ex.returncode)