Input guard against prompt injection

"""Input guard against prompt injection, built for the STUD workspace intake.

STUD runs plays from a short intake: each run's record (buyer-entered fields such
as an account id, an amount, a free-text note) flows into an agent prompt. This
guard sits in front of that hop. It enforces the declared field spec (type, max
length, charset), requires exactly the declared fields, and rejects text that
smuggles section-break delimiters or fake instruction/role markers, so an
injected "SYSTEM: ..." or "ignore previous instructions" never reaches the model.

Design choices, in order of why:
- Allowlist first: the per-field charset screens most injections outright (a
  colon or dash run cannot survive an alnum_space field). The marker scan exists
  for the letters-only attacks the charset cannot catch.
- Exact field set: a missing declared field or an extra undeclared field is a
  reject. Undeclared fields are the classic smuggling surface, so the guard
  treats the spec as the whole contract, not a minimum.
- Booleans are not numbers: True passes isinstance(int) in Python, so number
  fields exclude bool explicitly.
- The guard returns only "accept" or "reject". It is a precondition check on a
  fixed rule set, not a general maliciousness judgment.
"""

import re

# Per-charset allowlists. Empty string is allowed everywhere except email,
# whose pattern requires a full local@domain.tld shape.
CHARSETS = {
    "id": re.compile(r"^[A-Za-z0-9_-]*$"),
    "alnum": re.compile(r"^[A-Za-z0-9]*$"),
    "alnum_space": re.compile(r"^[A-Za-z0-9 ]*$"),
    "numeric": re.compile(r"^[0-9.eE+-]*$"),
    "email": re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$"),
}

# Section-break delimiters (---, ===, ***, ###) and fake instruction/role
# markers ("system:", "assistant:", "user:", "llm:", "ignore ... instruction"),
# case-insensitive. This is the published fixed rule set.
MARKERS = re.compile(
    r"-{3,}|={3,}|\*{3,}|#{3,}|\b(?:system|assistant|user|llm)\s*:|\bignore\b.{0,30}\binstruction",
    re.I,
)


def guard(inp):
    spec = inp.get("spec") or []
    record = inp.get("record")

    # The record must be a mapping; a scalar or list cannot satisfy a field spec.
    if not isinstance(record, dict):
        return "reject"

    # Exactly the declared fields: no extras (injection surface), no gaps.
    declared = {f["field"] for f in spec}
    if set(record.keys()) - declared:
        return "reject"

    for f in spec:
        name = f["field"]
        if name not in record:
            return "reject"
        value = record[name]
        ftype = f.get("type", "text")
        max_len = int(f.get("max_len", 0))
        charset = f.get("charset", "alnum")

        if ftype == "number":
            # bool is an int subclass in Python; a number field rejects it.
            if isinstance(value, bool) or not isinstance(value, (int, float)):
                return "reject"
            text = str(value)
        else:
            if not isinstance(value, str):
                return "reject"
            text = value

        if len(text) > max_len:
            return "reject"

        pattern = CHARSETS.get(charset)
        if pattern is not None and not pattern.match(text):
            return "reject"

        # The marker scan runs on the string-carrying types only.
        if ftype in ("text", "id", "email") and MARKERS.search(text):
            return "reject"

    return "accept"