attack-defense-platform/receiver/challenges/Sheesh.py

from .Challenge import Challenge

import subprocess
import time
import re
import os
import binascii

class Sheesh(Challenge):
    flag_location    = 'flags/sheesh.txt'
    history_location = 'history/sheesh.txt'

    _CONTAINER   = os.environ.get("CHALLENGE_CONTAINER_SHEESH", "sheesh_container")
    _SERVICE_CMD = ["docker", "exec", "-i", _CONTAINER, "python3", "/home/ctfuser/chall/src/chall.py"]
    _HEX_RE      = re.compile(r'^[0-9a-fA-F]+$')

    def _read_container_flag(self) -> str:
        # timeout is mandatory: docker exec can block forever on a saturated container
        out = subprocess.run(
            ["docker", "exec", self._CONTAINER, "cat", "/flag.txt"],
            capture_output=True, text=True, timeout=30
        )
        if out.returncode != 0 or not out.stdout.strip():
            raise FileNotFoundError("Flag not found in container (/flag.txt)")
        return out.stdout.strip()

    def _spawn(self):
        return subprocess.Popen(
            self._SERVICE_CMD,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=0,
        )

    def _read_exact_line(self, proc, timeout=5.0):
        start = time.time()
        buf = []
        r = proc.stdout.read
        while True:
            if time.time() - start > timeout:
                raise TimeoutError("Timeout waiting for a line")
            ch = r(1)
            if ch == "" and proc.poll() is not None:
                raise RuntimeError("Process ended unexpectedly")
            buf.append(ch)
            if ch == "\n":
                return "".join(buf)

    def _read_until_prompt(self, proc, token: str, timeout=5.0):
        start = time.time()
        acc = []
        r = proc.stdout.read
        while True:
            if time.time() - start > timeout:
                raise TimeoutError(f"Timeout waiting for '{token}'")
            ch = r(1)
            if ch == "" and proc.poll() is not None:
                raise RuntimeError("Process ended while waiting for prompt")
            acc.append(ch)
            if "".join(acc).endswith(token):
                return "".join(acc)

    def _send_line(self, proc, s: str):
        proc.stdin.write(s + "\n")
        proc.stdin.flush()

    def _expect_hex_field(self, text: str, label: str) -> str:
        m = re.search(rf"{re.escape(label)}\s*:\s*([0-9a-fA-F]+)", text)
        assert m, f"Missing '{label}' in output. Tail:\n{text[-400:]}"
        hx = m.group(1)
        assert self._HEX_RE.match(hx), f"{label} is not hex"
        return hx

    def distribute(self, flag):
        try:
            os.makedirs(os.path.dirname(self.flag_location), exist_ok=True)
            with open(self.flag_location, 'w') as f:
                f.write(flag)

            os.makedirs(os.path.dirname(self.history_location), exist_ok=True)
            with open(self.history_location, 'a') as f:
                f.write(flag + '\n')

            self.logger.info(f'Flag {flag} written to {self.flag_location}')
            return True
        except Exception as e:
            self.logger.error(f'Could not write flag to {self.flag_location}: {e}')
            return False

    def check(self):
        try:
            # parity check: host vs container flag
            with open(self.flag_location, 'r') as f:
                host_flag = f.read().strip()
            container_flag = self._read_container_flag()
            assert host_flag == container_flag, 'Flag mismatch between host and container'
            self.logger.info('[ok] flag parity (sheesh)')

            proc = self._spawn()

            line1 = self._read_exact_line(proc, timeout=5.0)
            line2 = self._read_exact_line(proc, timeout=5.0)
            blank = self._read_exact_line(proc, timeout=5.0)
            assert line1.startswith("iv:") and line2.startswith("ct:"), "Missing banner"
            iv_hex = self._expect_hex_field(line1, "iv")
            ct_hex = self._expect_hex_field(line2, "ct")
            iv0 = bytes.fromhex(iv_hex)
            ct0 = bytes.fromhex(ct_hex)
            assert len(iv0) == 16 and len(ct0) >= 16 and len(ct0) % 16 == 0, "Bad banner sizes"
            self.logger.info("[ok] banner parsed")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)

            qA = "00" * 16
            self._send_line(proc, qA)
            lA = self._read_exact_line(proc, timeout=3.0)
            assert lA.startswith("blk:"), "Expected blk: for 16-byte blob"
            blk1 = self._expect_hex_field(lA, "blk")
            _ = self._read_exact_line(proc, timeout=3.0)

            self._read_until_prompt(proc, "blob: ", timeout=5.0)
            self._send_line(proc, qA)
            lA2 = self._read_exact_line(proc, timeout=3.0)
            blk2 = self._expect_hex_field(lA2, "blk")
            _ = self._read_exact_line(proc, timeout=3.0)
            assert blk1 == blk2, "ECB-masked oracle not deterministic for same input"
            self.logger.info("[ok] blk determinism")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)
            qB = "01" * 16
            self._send_line(proc, qB)
            lB = self._read_exact_line(proc, timeout=3.0)
            blk3 = self._expect_hex_field(lB, "blk")
            _ = self._read_exact_line(proc, timeout=3.0)
            assert blk3 != blk1, "ECB-masked oracle returned same block for different input"
            self.logger.info("[ok] blk diversity")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)
            self._send_line(proc, "zzzz")   # invalid hex
            line_err = self._read_exact_line(proc, timeout=3.0)
            assert "err" in line_err.lower(), "Expected 'err' for non-hex input"
            self.logger.info("[ok] error handling on non-hex")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)
            self._send_line(proc, (iv_hex + ct_hex))
            ans1 = self._read_exact_line(proc, timeout=3.0).strip().lower()
            assert ans1 in ("ok", "no"), "Padding oracle did not return ok/no for iv||ct"
            self.logger.info(f"[ok] padding oracle responds ({ans1}) on banner bundle")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)
            iv_synth = os.urandom(16).hex()
            ct_synth = ("00" * (6 * 16))
            self._send_line(proc, iv_synth + ct_synth)
            ans2 = self._read_exact_line(proc, timeout=3.0).strip().lower()
            assert ans2 in ("ok", "no"), "Padding oracle did not return ok/no on synthetic 6-block ct"
            self.logger.info(f"[ok] padding oracle responds ({ans2}) on synthetic payload")

            self._read_until_prompt(proc, "blob: ", timeout=5.0)

            proc.kill()
            self.logger.info('Check passed for sheesh')
            return True

        except Exception as e:
            self.logger.error(f'Could not check sheesh: {e}')
            return False