init

changeme.py

+"""
+These are ctf dependent functions.
+They have to be aligned.
+"""
+
+
+MESSAGES = [
+        "SUCCESS",
+        "INVALID",
+        "ALREADY_SUBMITTED",
+        "ERROR"
+]
+
+RES_FILE_PATH = "results.txt"
+RES_FILE = None
+
+
+# TODO: Make list of team ids
+all_teams = [i for i in range(16)]
+
+
+def team_id_to_ip(team_id):
+    """
+    Somehow map team id to ip adress
+    """
+    return "10.13.37.{}".format(team_id)
+
+
+def submit_flag(team_id, flag):
+    """
+    Submit individual flag.
+    """
+    global RES_FILE
+    if RES_FILE is None:
+        RES_FILE = open(RES_FILE_PATH, "a")
+
+    # TODO: submission logic
+    # e.g. requests.put(...)
+    # set return val to index of appropriate message code
+    return_val = 0
+
+
+    RES_FILE.write("{} {} {}\n".format(return_val, team_id, flag))
+    RES_FILE.flush()  # dirty but works for now
+
+    return MESSAGES[return_val]
+

config.py

+# Has to be byte string since wre looking through byte output (STDOUT)
+flag_regex = b"flag{[A-Za-z0-9_.]+}"
+exploit_dir = "exploits"
+wait_between_runs = 30
+proc_num = 4

exploits/example_exploit.py

+#!/usr/bin/env python3
+import sys
+from random import randint
+
+team_id = sys.argv[1]
+team_ip = sys.argv[2]
+print("flag{%s_%s_%d}" % (team_id, team_ip, randint(0, 133713371337)))

main.py

+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()

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)
+

submitter.py

+import changeme
+import logging as log
+
+def do_submissions(flag_queue):
+    while True:
+        answer = flag_queue.get()
+        flag = answer['flag']
+        team_id = answer['team_id']
+        result = changeme.submit_flag(team_id, flag)
+        log.info("[SUBMITTER] Flag submission %s [team %d] returned %s", flag, team_id, result)
+