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

from .Challenge import Challenge

import select
import subprocess
import time
import re
import os


def _has_data(proc) -> bool:
    """True if the child's pipe still holds buffered output."""
    import fcntl
    try:
        fd = proc.stdout.fileno()
        fl = fcntl.fcntl(fd, fcntl.F_GETFL)
        fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        data = proc.stdout.read()
        if data:
            return True
        return False
    except Exception:
        return False


class Phew(Challenge):
    flag_location    = 'flags/phew.txt'
    history_location = 'history/phew.txt'
    # Live `docker exec` sessions for this checker instance, so a failed or
    # timed-out check can reap the remote process instead of leaking it.
    _children = []

    _CONTAINER   = os.environ.get("CHALLENGE_CONTAINER_PHEW", "phew_container")
    # PYTHONUNBUFFERED is mandatory: chall.py prints its menu to stdout, and the
    # checker reads that pipe interactively. Python block-buffers stdout when it
    # is not a tty, so without it the child never flushes the "1. encrypt ... > "
    # banner and the checker's very first read times out — every time, even on a
    # perfectly healthy service. `python3 -u` would do the same thing.
    _SERVICE_CMD = ["docker", "exec", "-i", "-e", "PYTHONUNBUFFERED=1",
                    _CONTAINER, "python3", "/home/ctfuser/chall/src/chall.py"]
    _HEX_RE      = re.compile(r'^[0-9a-fA-F]+$')
    # chall.py generates a fresh Pailier keypair (os.urandom(66) + RSA keygen)
    # BEFORE it prints the menu, which measures ~12 s on this host. The first
    # read must outlast that or the check fails on a healthy service. Later
    # exchanges reuse the same key, so they can stay short.
    _BOOT_TIMEOUT = 45.0
    # Budget for a single cipher operation. Encrypting the raw 528-bit key
    # (menu option 4) is measurably slower than encrypting a small plaintext,
    # and a saturated host makes even the small ones slower — 5 s was too tight
    # and produced a false DOWN.
    _CRYPTO_TIMEOUT = 30.0

    def _read_container_flag(self) -> str:
        # NB: a timeout is mandatory here. `docker exec` against a container
        # whose process table is saturated (accumulated chall.py zombies) can
        # block forever and take the whole SLA check loop down with it.
        try:
            out = subprocess.run(["docker", "exec", self._CONTAINER, "cat", "/flag.txt"],
                                 capture_output=True, text=True, timeout=30)
        except subprocess.TimeoutExpired:
            raise TimeoutError("docker exec cat /flag.txt timed out (container overloaded?)")
        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):
        proc = subprocess.Popen(
            self._SERVICE_CMD,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=0,
        )
        self._children.append(proc)
        return proc

    def _reap(self, proc):
        """Kill a spawned service session, INSIDE the container too.

        proc.kill() only kills the local `docker exec` CLIENT. The chall.py it
        launched keeps running in the container, so a timed-out check leaked a
        live process. After a few failures the container had 26 concurrent
        chall.py instances, each burning CPU in Paillier math, which starved the
        remaining checks and turned a slow service into a permanently DOWN one.
        Reaping must therefore also pkill the remote process.
        """
        if proc.poll() is None:
            try:
                proc.kill()
                proc.wait(timeout=5)
            except Exception:
                pass
        try:
            subprocess.run(["docker", "exec", self._CONTAINER, "sh", "-c",
                            "pkill -f chall.py 2>/dev/null || true"],
                           capture_output=True, timeout=20)
        except Exception:
            pass
        if proc in self._children:
            self._children.remove(proc)

    def _reap_all(self):
        for p in list(self._children):
            self._reap(p)

    def _read_until(self, proc, token, timeout=5.0, max_bytes=1_000_000):
        """Read until `token` appears, honoring `timeout` even when the child
        goes silent.

        The previous implementation used a blocking `stdout.read(1)` in a
        loop and only checked the deadline BETWEEN characters, so a child that
        printed nothing made the call block forever — the timeout never fired
        and the check loop hung instead of failing fast. select() makes the
        deadline authoritative.
        """
        start = time.time()
        buf = []
        deadline = start + timeout
        while True:
            if time.time() > deadline:
                tail = ''.join(buf)[-500:]
                raise TimeoutError(f"Timeout waiting for '{token}'. Got so far:\n{tail}")
            remaining = deadline - time.time()
            ready, _, _ = select.select([proc.stdout], [], [], min(remaining, 1.0))
            if not ready:
                if proc.poll() is not None and not _has_data(proc):
                    raise RuntimeError(
                        f"Process ended while waiting for '{token}'. Output:\n{''.join(buf)}")
                continue
            ch = proc.stdout.read(1)
            if ch == "":
                raise RuntimeError(
                    f"Process ended while waiting for '{token}'. Output:\n{''.join(buf)}")
            buf.append(ch)
            if len(buf) > max_bytes:
                raise RuntimeError("Exceeded max read size")
            if token in "".join(buf):
                return "".join(buf)

    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
            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 (phew)')

            def run_encrypt_once(pt_hex: str) -> str:
                proc = self._spawn()
                try:
                    self._read_until(proc, "> ", timeout=self._BOOT_TIMEOUT)
                    self._send_line(proc, "1")
                    self._read_until(proc, "pt (hex)", timeout=3.0)
                    self._read_until(proc, "> ", timeout=3.0)
                    self._send_line(proc, pt_hex)
                    out = self._read_until(proc, "> ", timeout=self._CRYPTO_TIMEOUT)
                    ct_hex = self._expect_hex_field(out, "ct")
                    self._send_line(proc, "9")
                    try:
                        proc.wait(timeout=2.0)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        raise AssertionError("Program did not exit after exit command (encrypt)")
                    return ct_hex
                finally:
                    self._reap(proc)

            def run_decrypt_once(ct_hex: str) -> str:
                proc = self._spawn()
                try:
                    self._read_until(proc, "> ", timeout=self._BOOT_TIMEOUT)
                    self._send_line(proc, "3")
                    self._read_until(proc, "ct (hex)", timeout=3.0)
                    self._read_until(proc, "> ", timeout=3.0)
                    self._send_line(proc, ct_hex)
                    out = self._read_until(proc, "> ", timeout=self._CRYPTO_TIMEOUT)
                    pt_hex = self._expect_hex_field(out, "pt")
                    self._send_line(proc, "9")
                    try:
                        proc.wait(timeout=2.0)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        raise AssertionError("Program did not exit after exit command (decrypt)")
                    return pt_hex
                finally:
                    self._reap(proc)

            def run_keyct_once() -> str:
                proc = self._spawn()
                try:
                    self._read_until(proc, "> ", timeout=self._BOOT_TIMEOUT)
                    self._send_line(proc, "4")
                    # Option 4 runs `cipher.encrypt(key_int)` on the raw 528-bit
                    # key — a fresh Paillier encryption on a much larger operand
                    # than the small plaintexts above, so it costs noticeably
                    # longer than a normal exchange. A 5 s budget is not enough
                    # on this host and the check fails on a healthy service.
                    out = self._read_until(proc, "> ", timeout=self._CRYPTO_TIMEOUT)
                    ct_hex = self._expect_hex_field(out, "ct")
                    self._send_line(proc, "9")
                    try:
                        proc.wait(timeout=2.0)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        raise AssertionError("Program did not exit after exit command (keyct)")
                    return ct_hex
                finally:
                    self._reap(proc)

            def run_bingo_reject_wrong_key():
                wrong_key_hex = "00" * 66
                proc = self._spawn()
                try:
                    self._read_until(proc, "> ", timeout=self._BOOT_TIMEOUT)
                    self._send_line(proc, "2")
                    self._read_until(proc, "key (hex)", timeout=3.0)
                    self._read_until(proc, "> ", timeout=3.0)
                    self._send_line(proc, wrong_key_hex)
                    out = self._read_until(proc, "> ", timeout=self._CRYPTO_TIMEOUT)
                    assert "nope" in out.lower(), f"bingo did not reject wrong key; got:\n{out[-300:]}"
                    self._send_line(proc, "9")
                    try:
                        proc.wait(timeout=2.0)
                    except subprocess.TimeoutExpired:
                        proc.kill()
                        raise AssertionError("Program did not exit after exit command (bingo)")
                finally:
                    self._reap(proc)

            ct1 = run_encrypt_once("414243444546") 
            assert ct1 and self._HEX_RE.match(ct1), "encrypt(1) did not return hex"
            self.logger.info("[ok] encrypt produced hex")

            pt_back = run_decrypt_once(ct1)
            assert pt_back.strip() != "", "decrypt returned empty output"
            assert self._HEX_RE.match(pt_back), "decrypt(3) did not return hex"
            self.logger.info("[ok] decrypt produced hex (custom mapping accepted)")

            pt_same = "01" * 8
            ct_a = run_encrypt_once(pt_same)
            ct_b = run_encrypt_once(pt_same)
            assert ct_a.lower() != ct_b.lower(), "Encryption appears deterministic for same plaintext"
            self.logger.info("[ok] encrypt randomness")

            k1 = run_keyct_once()
            k2 = run_keyct_once()
            assert k1.lower() != k2.lower(), "key? ciphertexts reused randomness"
            self.logger.info("[ok] key? randomness")

            # run_bingo_reject_wrong_key()
            # self.logger.info("[ok] bingo rejects wrong key")

            return True

        except Exception as e:
            self.logger.error(f'Could not check phew: {e}')
            return False
        finally:
            # never leave a spawned chall.py behind, whatever happened above
            self._reap_all()