Handling hostile input at scale

"""Budgets for untrusted input, and a read that honours two of them: total bytes and total time."""
import os
import select
import time
from dataclasses import dataclass
 
 
@dataclass(frozen=True)
class Budget:
    max_bytes: int = 1_000_000     # never read more than this from the source at all
    max_seconds: float = 5.0       # the whole input must arrive within this time
    max_lines: int = 10_000        # a record count a single run should never exceed
    max_line_bytes: int = 8_192    # one record is small; a huge line is an attack or a bug
 
 
class BudgetExceeded(Exception):
    """The input is larger or slower than we agreed to process: reject it, do not truncate it."""
 
 
def read_capped(fd: int, budget: Budget) -> bytes:
    # Chunk by chunk: memory stays bounded even if the source is endless, and every wait draws
    # from one deadline, so a sender that trickles a byte at a time cannot hold the worker.
    give_up_at = time.monotonic() + budget.max_seconds
    chunks, size = [], 0
    while True:
        remaining = give_up_at - time.monotonic()
        if remaining <= 0 or not select.select([fd], [], [], remaining)[0]:
            raise BudgetExceeded(f"input still incomplete after {budget.max_seconds} s")
        chunk = os.read(fd, 65_536)
        if not chunk:                          # end of input
            return b"".join(chunks)
        size += len(chunk)
        if size > budget.max_bytes:
            raise BudgetExceeded(f"input exceeds {budget.max_bytes} bytes")
        chunks.append(chunk)

"""Split capped input into records under two more budgets: line count and single-line length."""
import sys
 
from limits import Budget, BudgetExceeded, read_capped
 
 
def read_records(fd: int, budget: Budget):
    lines = read_capped(fd, budget).split(b"\n")
    if lines[-1] == b"":                       # a final newline ends the last line, it adds none
        lines.pop()
    if len(lines) > budget.max_lines:
        raise BudgetExceeded(f"more than {budget.max_lines} lines")
    for n, raw in enumerate(lines, 1):
        if len(raw) > budget.max_line_bytes:
            raise BudgetExceeded(f"line {n} is {len(raw)} bytes, over {budget.max_line_bytes}")
        if raw:
            yield raw
 
 
if __name__ == "__main__":
    try:
        count = sum(1 for _ in read_records(sys.stdin.fileno(), Budget()))
    except BudgetExceeded as err:
        print(f"rejected: {err}", file=sys.stderr)
        raise SystemExit(1)
    print(f"accepted {count} records")

"""Time the built-in re engine on a nested-quantifier pattern against an almost-matching string.
 
The pattern validates space-separated host names. Each argument is the total length of the input."""
import re
import sys
import time
 
HOSTS = re.compile(r"^(\w+\s?)+$")  # a group that repeats, containing \w+ that also repeats
 
for arg in sys.argv[1:]:
    text = "a" * (int(arg) - 1) + "!"  # one final character the pattern cannot match
    start = time.perf_counter()
    HOSTS.search(text)
    print(f"{len(text):>3} characters: {time.perf_counter() - start:7.3f} s", flush=True)

"""Fix the pattern first, then cap the field: a cap bounds the work only once the pattern is linear."""
import re
import time
 
CAP = 253                                    # a host name is never longer than this
NESTED = re.compile(r"^(\w+\s?)+$")          # the shape from ps-regex: exponential on a near miss
FLAT = re.compile(r"^\w+(?:\s\w+)*\s?$")     # the same language: every repeat starts with a space
ATOMIC = re.compile(r"^(?>\w+\s?)+$")        # or keep the shape and forbid backtracking into it
 
 
def valid_hosts(value: str) -> bool:
    if len(value) > CAP:                     # the budget: refuse an over-long field, never truncate
        return False
    return FLAT.search(value) is not None
 
 
for text in ("a" * 28 + "!", "a" * 40_000 + "!"):
    for name, pattern in (("flat", FLAT), ("atomic", ATOMIC)):
        start = time.perf_counter()
        pattern.search(text)
        print(f"{name:>6} on {len(text):>5} characters: {time.perf_counter() - start:.4f} s")
 
samples = ["web01 web02", "web01 ", "web01  web02", "web01;id"]
print("nested:", [bool(NESTED.search(s)) for s in samples])  # safe here: these are short and benign
print("flat  :", [bool(FLAT.search(s)) for s in samples])
print("capped:", valid_hosts("web01 web02"), valid_hosts("a " * 200))

"""signal.alarm does interrupt re.search on CPython 3.14, but only in the main thread."""
import re
import signal
import threading
import time
 
NESTED = re.compile(r"^(\w+\s?)+$")
 
 
class MatchTimeout(Exception):
    pass
 
 
def on_alarm(signum, frame):
    raise MatchTimeout  # raised in the main thread, wherever it is: here, inside the matcher
 
 
signal.signal(signal.SIGALRM, on_alarm)
start = time.monotonic()
signal.alarm(1)                              # whole seconds only; setitimer() for finer steps
try:
    NESTED.search("a" * 40 + "!")
except MatchTimeout:
    print(f"main thread: re.search interrupted after {time.monotonic() - start:.1f} s")
finally:
    signal.alarm(0)
 
 
def worker() -> None:
    try:
        signal.signal(signal.SIGALRM, on_alarm)
    except ValueError as err:
        print(f"worker thread: {err}")
 
 
thread = threading.Thread(target=worker)
thread.start()
thread.join()

"""Two engines for patterns you cannot rewrite: regex with a timeout, and google-re2 in linear time.
 
Run under the lab venv, which pins both."""
import re
import time
from concurrent.futures import ThreadPoolExecutor
 
import re2
import regex
 
LONG = "a" * 40_000 + "!"
 
 
def bounded_search(pattern: str, text: str) -> str:
    # regex checks a wall-clock timeout as it matches, in any thread, unlike signal.alarm.
    try:
        regex.search(pattern, text, timeout=0.2)
        return "completed"
    except TimeoutError:
        return "TimeoutError after 0.2 s, the worker is free again"
 
 
with ThreadPoolExecutor(max_workers=1) as pool:
    print("regex in a worker thread:", pool.submit(bounded_search, r"(a|a)+$", LONG).result())
 
start = time.perf_counter()
match = re2.search(r"^(\w+\s?)+$", LONG)     # the exponential pattern, on 40,001 characters
print(f"re2 linear: {time.perf_counter() - start:.4f} s on {len(LONG)} chars, match={bool(match)}")
 
# Before swapping re2 in, compare answers, including the inputs where the engines differ.
samples = ["web01", "web01;id", "wéb01", "web01\n"]
print("re  :", [bool(re.search(r"^\w+$", s)) for s in samples])
print("re2 :", [bool(re2.search(r"^\w+$", s)) for s in samples])

"""A payload object. Unpickling it runs code before you read a single field. Shown to be avoided.
 
To keep the lab harmless the payload only writes a marker file in the working directory; an attacker
would put anything runnable here. This is why pickle must never touch data you did not produce."""
from pathlib import Path
 
 
def _run_on_load(path: str) -> None:  # stands in for any command an attacker could choose
    Path(path).write_text("arbitrary code ran during unpickle\n")
 
 
class Report:
    def __reduce__(self):
        # __reduce__ tells pickle which callable to run and with what arguments when it rebuilds
        # the object. pickle obeys it during load, before the result is ever inspected.
        return (_run_on_load, ("pickle-marker",))

"""Load a pickle the way a hundred tutorials show. Unsafe on untrusted input: shown to be avoided."""
import pickle
import sys
 
with open(sys.argv[1], "rb") as f:
    obj = pickle.load(f)  # runs the payload's __reduce__ during the load
print("loaded object:", obj)

"""Validate an untrusted record into a typed object, collecting every error before failing.
 
json.loads gives you dicts and lists, not a checked object. Check shape, type, range and grammar
once, at the boundary, and reject with a clear message instead of a traceback."""
import re
from dataclasses import dataclass
 
SEVERITIES = {"low", "medium", "high", "critical"}
# Letters, digits, dots and hyphens, starting and ending with a letter or digit: no spaces, no
# ';', no leading '-' that a later command would read as an option, no newline, no escapes.
HOSTNAME = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9.-]{0,251}[A-Za-z0-9])?")
 
 
class ValidationError(Exception):
    def __init__(self, errors: list[str]):
        self.errors = errors
        super().__init__("; ".join(errors))
 
 
@dataclass(frozen=True)
class Finding:
    target: str
    severity: str
    score: int
 
 
def parse_finding(raw: object) -> Finding:
    if not isinstance(raw, dict):                # valid JSON can still be a list, string or number
        raise ValidationError(["record must be a JSON object"])
    errors: list[str] = []
    target = raw.get("target")
    if not isinstance(target, str) or len(target) > 253 or not HOSTNAME.fullmatch(target):
        errors.append("target must be a host name (letters, digits, '.', '-'; at most 253)")
    severity = raw.get("severity")
    if severity not in SEVERITIES:
        errors.append(f"severity must be one of {sorted(SEVERITIES)}")
    score = raw.get("score")
    # isinstance(True, int) is True in Python, so reject bool explicitly; a "90" string is not an int
    if isinstance(score, bool) or not isinstance(score, int) or not (0 <= score <= 100):
        errors.append("score must be an integer in 0..100")
    if errors:
        raise ValidationError(errors)
    return Finding(target=target, severity=severity, score=score)

"""The trust boundary for one record: JSON text in, a checked Finding or a ValidationError out."""
import json
import sys
 
from schema import Finding, ValidationError, parse_finding
 
 
def load_finding(text: str | bytes) -> Finding:
    try:
        raw = json.loads(text)
    except (ValueError, RecursionError) as err:  # JSONDecodeError is only one of the ways it fails
        raise ValidationError([f"not usable JSON ({err.__class__.__name__})"]) from None
    return parse_finding(raw)
 
 
if __name__ == "__main__":
    try:
        print("valid:", load_finding(sys.argv[1]))
    except ValidationError as err:
        print("rejected:", err, file=sys.stderr)
        raise SystemExit(1)

"""Write BASE/RELPATH in a tree others can write, three ways. Usage: toctou.py MODE BASE RELPATH
 
naive  follows every symlink on the way: the write lands wherever the links point
leaf   O_NOFOLLOW on the last name only: a symlinked parent directory or '..' still escapes
walk   one component at a time, never following a symlink and refusing '..': stays under BASE"""
import os
import sys
 
WRITE = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
 
 
def open_beneath(base: str, relpath: str, flags: int) -> int:
    parts = relpath.split("/")
    if relpath.startswith("/") or any(part in ("", ".", "..") for part in parts):
        raise ValueError(f"refusing {relpath!r}: absolute, empty, '.' or '..' component")
    fd = os.open(base, os.O_RDONLY | os.O_DIRECTORY)  # BASE comes from your config, not the attacker
    try:
        for part in parts[:-1]:                       # each directory: no symlink, must be a dir
            next_fd = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd)
            os.close(fd)
            fd = next_fd
        return os.open(parts[-1], flags | os.O_NOFOLLOW, 0o644, dir_fd=fd)
    finally:
        os.close(fd)
 
 
mode, base, relpath = sys.argv[1:4]
try:
    if mode == "walk":
        fd = open_beneath(base, relpath, WRITE)
    else:
        base_fd = os.open(base, os.O_RDONLY | os.O_DIRECTORY)
        extra = os.O_NOFOLLOW if mode == "leaf" else 0
        try:
            fd = os.open(relpath, WRITE | extra, 0o644, dir_fd=base_fd)
        finally:
            os.close(base_fd)
    with os.fdopen(fd, "wb") as out:
        out.write(b"written by the ingest job\n")
    print(f"{mode}: wrote {relpath}")
except (OSError, ValueError) as err:
    reason = os.strerror(err.errno) if isinstance(err, OSError) else str(err)
    print(f"{mode}: refused {relpath} ({reason})")
    raise SystemExit(1)

"""A field from an untrusted record ends up in a log line. Two ways to write it, one safe.
 
If you paste the value into a text line, a newline in it forges a second log entry and an ANSI escape
can rewrite your terminal. A structured logger serialises the value, so control characters survive as
data (\\n, \\u001b) and can never break out of the field."""
import json
import logging
import sys
from datetime import datetime, timezone
 
 
class JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        payload = {
            "ts": datetime.fromtimestamp(record.created, timezone.utc).isoformat(),
            "level": record.levelname.lower(),
            "msg": record.getMessage(),
        }
        fields = getattr(record, "fields", None)
        if fields:
            # Nested, not merged: an untrusted "level" or "msg" key cannot replace the real one.
            payload["fields"] = fields
        return json.dumps(payload, ensure_ascii=True)  # escapes newlines and control characters
 
 
# A user name crafted to inject a second, fake "authentication failed" line and hide it with ANSI.
hostile = "alice\n2026-01-01 00:00:00 CRITICAL auth failed for admin\x1b[2K"
 
naive = logging.getLogger("naive")
naive.addHandler(logging.StreamHandler(sys.stderr))
naive.setLevel(logging.INFO)
print("--- naive: value pasted into the message text", file=sys.stderr, flush=True)
naive.info("login user=%s", hostile)  # the newline in `hostile` becomes a real new line
 
structured = logging.getLogger("structured")
h = logging.StreamHandler(sys.stderr)
h.setFormatter(JsonFormatter())
structured.addHandler(h)
structured.setLevel(logging.INFO)
structured.propagate = False
print("--- structured: value carried as a field", file=sys.stderr, flush=True)
structured.info("login", extra={"fields": {"user": hostile, "level": "debug"}})  # "level" came from the record