ai-agent-security-toolkit 0.4.1__py3-none-any.whl

agentsec/__init__.py

@@ -0,0 +1,3 @@
+"""ai-agent-security-toolkit: fuzz, contain, validate and audit LLM agents."""
+
+__version__ = "0.4.1"

agentsec/cli.py

@@ -0,0 +1,219 @@
+"""``agentsec`` command line interface."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+from agentsec import __version__
+from agentsec.fuzzer import CORPUS, MUTATORS, FuzzConfig, load_target, run
+from agentsec.fuzzer.harness import SCENARIOS, build_cases
+from agentsec.middleware.audit import verify_file
+from agentsec.sandbox import Policy, PolicyError, ToolGuard
+from agentsec.threatmodel import (
+    SYSTEM_TEMPLATE,
+    SystemSpecError,
+    crosswalk_markdown,
+    load_catalog,
+    load_system,
+    render_markdown,
+)
+from agentsec.types import ToolCall
+
+
+def _csv(value: str | None) -> list[str] | None:
+    return [v.strip() for v in value.split(",") if v.strip()] if value else None
+
+
+def cmd_fuzz(args: argparse.Namespace) -> int:
+    config = FuzzConfig(
+        seed=args.seed,
+        scenarios=_csv(args.scenarios) or list(SCENARIOS),
+        mutators=_csv(args.mutators),
+        categories=_csv(args.categories),
+        max_cases=args.max_cases,
+    )
+    if args.list:
+        cases = build_cases(config)
+        print(f"{len(CORPUS)} payloads x {len(MUTATORS)} mutators x {len(SCENARIOS)} scenarios")
+        print(f"selected: {len(cases)} cases")
+        print("mutators:", ", ".join(MUTATORS))
+        return 0
+    try:
+        target = load_target(args.target)
+        report = run(target, config)
+    except (ValueError, ImportError, AttributeError) as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+    if args.json:
+        Path(args.json).write_text(report.to_json(), encoding="utf-8", newline="\n")
+    if args.md:
+        Path(args.md).write_text(report.to_markdown(), encoding="utf-8", newline="\n")
+    if args.sarif:
+        Path(args.sarif).write_text(
+            report.to_sarif(artifact_uri=args.sarif_artifact, target=args.target),
+            encoding="utf-8",
+            newline="\n",
+        )
+    if args.junit:
+        Path(args.junit).write_text(report.to_junit(), encoding="utf-8", newline="\n")
+    print(
+        f"{report.successes}/{report.total - report.errors} attacks succeeded "
+        f"(ASR {report.asr:.1%}); {report.blocked} blocked, {report.errors} errors"
+    )
+    if not args.quiet:
+        print()
+        print(report.to_markdown(max_bypasses=8))
+    if args.max_asr is not None and report.exceeds(args.max_asr):
+        print(f"FAIL: ASR {report.asr:.1%} exceeds --max-asr {args.max_asr:.1%}", file=sys.stderr)
+        return 1
+    return 0
+
+
+def cmd_policy(args: argparse.Namespace) -> int:
+    try:
+        policy = Policy.from_yaml(args.file)
+    except (PolicyError, OSError, ValueError) as exc:
+        print(f"invalid policy: {exc}", file=sys.stderr)
+        return 2
+    if args.action == "validate":
+        notes = policy.lint()
+        print(f"OK: {len(policy.tools)} tool(s) defined")
+        for n in notes:
+            print(f"  warning: {n}")
+        return 1 if (notes and args.strict) else 0
+    try:
+        call_args: Any = json.loads(args.args or "{}")
+    except json.JSONDecodeError as exc:
+        print(f"--args is not valid JSON: {exc}", file=sys.stderr)
+        return 2
+    guard = ToolGuard(policy)
+    if args.tainted:
+        guard.default_session.mark_untrusted("cli")
+    decision = guard.evaluate(ToolCall(name=args.tool, arguments=call_args))
+    print(f"{decision.verdict.value.upper()}: {decision.message()}")
+    return 0 if decision.allowed else 1
+
+
+def cmd_audit(args: argparse.Namespace) -> int:
+    try:
+        result = verify_file(args.file)
+    except OSError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+    if result.ok:
+        print(f"OK: {result.records_checked} record(s), chain intact")
+        return 0
+    print(
+        f"TAMPERED: {result.reason} (first bad seq: {result.first_bad_seq}, "
+        f"{result.records_checked} valid before it)",
+        file=sys.stderr,
+    )
+    return 1
+
+
+def cmd_threatmodel(args: argparse.Namespace) -> int:
+    if args.action == "init":
+        out = Path(args.output) if args.output else None
+        if out:
+            out.write_text(SYSTEM_TEMPLATE, encoding="utf-8", newline="\n")
+            print(f"wrote {out}")
+        else:
+            print(SYSTEM_TEMPLATE, end="")
+        return 0
+    if args.action == "crosswalk":
+        print(crosswalk_markdown(), end="")
+        return 0
+    if args.action == "catalog":
+        for t in load_catalog():
+            print(f"{t.id}  {t.stride:<24} {t.title}")
+        return 0
+    try:
+        doc = render_markdown(load_system(args.system), load_catalog())
+    except (SystemSpecError, OSError, KeyError) as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+    if args.output:
+        Path(args.output).write_text(doc, encoding="utf-8", newline="\n")
+        print(f"wrote {args.output}")
+    else:
+        print(doc)
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="agentsec", description="Security toolkit for LLM agents.")
+    p.add_argument("--version", action="version", version=f"agentsec {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    f = sub.add_parser("fuzz", help="Run the prompt-injection fuzzer against a target.")
+    f.add_argument(
+        "--target",
+        default="builtin:guarded",
+        help="builtin:naive | builtin:guarded | builtin:guarded-noscanner | http(s)://URL | "
+        "anthropic:MODEL | package.module:callable",
+    )
+    f.add_argument("--scenarios", help=f"comma list from {list(SCENARIOS)} (default: all)")
+    f.add_argument("--mutators", help=f"comma list from {list(MUTATORS)} (default: all)")
+    f.add_argument("--categories", help="comma list of payload categories (default: all)")
+    f.add_argument("--seed", type=int, default=0)
+    f.add_argument("--max-cases", type=int, help="randomly sample at most N cases")
+    f.add_argument("--max-asr", type=float, help="exit 1 if attack-success rate exceeds this (0-1)")
+    f.add_argument("--json", help="write full JSON report to this path")
+    f.add_argument("--md", help="write Markdown report to this path")
+    f.add_argument("--sarif", help="write a SARIF 2.1.0 report (GitHub code scanning)")
+    f.add_argument(
+        "--sarif-artifact",
+        default="agent",
+        help="repo-relative path of the file that defines the agent; SARIF alerts attach to it",
+    )
+    f.add_argument("--junit", help="write a JUnit XML report to this path")
+    f.add_argument("--quiet", action="store_true", help="only print the one-line summary")
+    f.add_argument("--list", action="store_true", help="show the test matrix and exit")
+    f.set_defaults(func=cmd_fuzz)
+
+    pol = sub.add_parser("policy", help="Validate a tool policy or test a call against it.")
+    pol_sub = pol.add_subparsers(dest="action", required=True)
+    pv = pol_sub.add_parser("validate", help="Load and lint a policy file.")
+    pv.add_argument("file")
+    pv.add_argument("--strict", action="store_true", help="exit 1 on lint warnings")
+    pc = pol_sub.add_parser("check", help="Evaluate one tool call against a policy.")
+    pc.add_argument("file")
+    pc.add_argument("--tool", required=True)
+    pc.add_argument("--args", help="JSON object of arguments")
+    pc.add_argument(
+        "--tainted", action="store_true", help="evaluate as if untrusted content was read"
+    )
+    pol.set_defaults(func=cmd_policy)
+
+    a = sub.add_parser("audit", help="Audit-log utilities.")
+    a_sub = a.add_subparsers(dest="action", required=True)
+    av = a_sub.add_parser("verify", help="Verify a hash-chained JSONL audit log.")
+    av.add_argument("file")
+    a.set_defaults(func=cmd_audit)
+
+    tm = sub.add_parser("threatmodel", help="STRIDE-for-agents threat modeling.")
+    tm_sub = tm.add_subparsers(dest="action", required=True)
+    ti = tm_sub.add_parser("init", help="Print or write a starter system description.")
+    ti.add_argument("-o", "--output")
+    tr = tm_sub.add_parser("render", help="Generate a threat model from a system description.")
+    tr.add_argument("system", help="system description YAML")
+    tr.add_argument("-o", "--output")
+    tm_sub.add_parser("catalog", help="List the threat catalog.")
+    tm_sub.add_parser(
+        "crosswalk", help="Map the catalog to OWASP Agentic Top 10 and MITRE ATLAS (Markdown)."
+    )
+    tm.set_defaults(func=cmd_threatmodel)
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    return int(args.func(args))
+
+
+if __name__ == "__main__":
+    sys.exit(main())

agentsec/fuzzer/__init__.py

@@ -0,0 +1,44 @@
+"""Prompt-injection test harness / fuzzer."""
+
+from agentsec.fuzzer.harness import (
+    Case,
+    CaseResult,
+    FuzzConfig,
+    build_cases,
+    detect,
+    run,
+)
+from agentsec.fuzzer.mock_agents import GuardedAgent, NaiveAgent, demo_policy
+from agentsec.fuzzer.mutators import MUTATORS
+from agentsec.fuzzer.payloads import CORPUS, Payload
+from agentsec.fuzzer.report import FuzzReport
+from agentsec.fuzzer.targets import (
+    AnthropicTarget,
+    HttpTarget,
+    OllamaTarget,
+    load_target,
+    sync_target,
+    text_target,
+)
+
+__all__ = [
+    "CORPUS",
+    "MUTATORS",
+    "AnthropicTarget",
+    "Case",
+    "CaseResult",
+    "FuzzConfig",
+    "FuzzReport",
+    "GuardedAgent",
+    "HttpTarget",
+    "NaiveAgent",
+    "OllamaTarget",
+    "Payload",
+    "build_cases",
+    "demo_policy",
+    "detect",
+    "load_target",
+    "run",
+    "sync_target",
+    "text_target",
+]

agentsec/fuzzer/formats.py

@@ -0,0 +1,222 @@
+"""CI-friendly renderings of a fuzz report: SARIF 2.1.0 and JUnit XML.
+
+SARIF lets results appear in GitHub code scanning (and other SARIF viewers); JUnit
+XML is understood by nearly every CI system's test-report UI.
+
+Both embed text that came back from the system under test. An attacker may have
+steered that text, so it is bounded and stripped of control characters before use.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from typing import TYPE_CHECKING, Any
+from xml.etree import ElementTree as ET  # noqa: S405 - only used to *generate* XML
+
+from agentsec import __version__
+
+if TYPE_CHECKING:
+    from agentsec.fuzzer.harness import CaseResult
+    from agentsec.fuzzer.report import FuzzReport
+
+SARIF_SCHEMA = "https://json.schemastore.org/sarif-2.1.0.json"
+INFORMATION_URI = "https://github.com/DustyStudy/ai-agent-security-toolkit"
+MAX_TEXT = 300
+
+# Unicode categories removed from quoted text: control characters (incl. NUL and ESC), format
+# characters (zero-width and bidirectional overrides that can visually spoof a report),
+# line/paragraph separators and lone surrogates. XML 1.0 cannot carry most of these anyway.
+_STRIP_CATEGORIES = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"})
+
+# goal -> (rule id, name, severity level, GitHub security-severity, OWASP LLM Top 10 (2025), help)
+_RULES: dict[str, tuple[str, str, str, str, str, str]] = {
+    "tool_misuse": (
+        "AGSEC-TOOL-MISUSE",
+        "InjectedInstructionCalledForbiddenTool",
+        "error",
+        "8.5",
+        "LLM06",
+        "A prompt injection made the agent request a tool it must never call. Enforce a "
+        "deny-by-default tool policy (agentsec.sandbox.ToolGuard) so a hijacked model "
+        "cannot act on the request.",
+    ),
+    "exfil": (
+        "AGSEC-EXFILTRATION",
+        "InjectedInstructionCausedDataExfiltration",
+        "error",
+        "8.5",
+        "LLM02",
+        "A prompt injection made the agent send data toward an attacker-controlled host "
+        "(tool arguments or a rendered markdown/HTML image or link). Validate output and "
+        "restrict outbound destinations (agentsec.middleware.OutputGuard, ToolGuard URL rules).",
+    ),
+    "prompt_leak": (
+        "AGSEC-PROMPT-LEAK",
+        "InjectedInstructionLeakedSystemPrompt",
+        "error",
+        "7.0",
+        "LLM07",
+        "A prompt injection made the agent disclose a secret planted in its system prompt. "
+        "Keep secrets out of prompts and add a protected-string output check.",
+    ),
+    "canary": (
+        "AGSEC-INSTRUCTION-FOLLOWING",
+        "AgentFollowedInjectedInstruction",
+        "warning",
+        "5.0",
+        "LLM01",
+        "The agent followed an instruction planted in untrusted content. On its own this is "
+        "low impact, but it is the first step of every real injection. Contain the blast "
+        "radius with tool policy and output validation.",
+    ),
+}
+_FALLBACK = (
+    "AGSEC-OTHER",
+    "InjectionSucceeded",
+    "warning",
+    "5.0",
+    "LLM01",
+    "A prompt-injection test case succeeded against the target.",
+)
+
+
+def clean(text: str, limit: int = MAX_TEXT) -> str:
+    """Bound and de-fang text that originated in the system under test."""
+    text = "".join(
+        " " if unicodedata.category(ch) in _STRIP_CATEGORIES or ch in "\ufffe\uffff" else ch
+        for ch in text
+    )
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def _rule_for(goal: str) -> tuple[str, str, str, str, str, str]:
+    return _RULES.get(goal, _FALLBACK)
+
+
+# ------------------------------------------------------------------------------ SARIF
+
+
+def to_sarif(
+    report: FuzzReport, *, artifact_uri: str = "agent", target: str | None = None
+) -> dict[str, Any]:
+    """Render successful attacks as a SARIF 2.1.0 log.
+
+    Fuzz findings have no source location, but code-scanning UIs need one to attach an
+    alert to, so every result points at ``artifact_uri`` (line 1). Pass the path of the
+    file that defines your agent, relative to the repository root.
+    """
+    used = sorted({_rule_for(r.goal) for r in report.results if r.success})
+    rules = [
+        {
+            "id": rid,
+            "name": name,
+            "shortDescription": {"text": name},
+            "fullDescription": {"text": help_text},
+            "help": {"text": help_text},
+            "defaultConfiguration": {"level": level},
+            "properties": {
+                "tags": ["security", "prompt-injection", "OWASP-LLM-2025-" + owasp],
+                "security-severity": sev,
+                "precision": "high",
+            },
+        }
+        for rid, name, level, sev, owasp, help_text in used
+    ]
+    results = [_sarif_result(r, artifact_uri) for r in report.results if r.success]
+    run: dict[str, Any] = {
+        "tool": {
+            "driver": {
+                "name": "agentsec",
+                "version": __version__,
+                "informationUri": INFORMATION_URI,
+                "rules": rules,
+            }
+        },
+        "results": results,
+        "properties": {
+            "seed": report.seed,
+            "attackSuccessRate": report.asr,
+            "cases": report.total,
+            "targetErrors": report.errors,
+        },
+    }
+    if target:
+        run["properties"]["target"] = clean(target, 200)
+    return {"$schema": SARIF_SCHEMA, "version": "2.1.0", "runs": [run]}
+
+
+def _sarif_result(r: CaseResult, artifact_uri: str) -> dict[str, Any]:
+    rid, _, level, _, _, _ = _rule_for(r.goal)
+    detail = f" Evidence: {clean(r.evidence)}" if r.evidence else ""
+    return {
+        "ruleId": rid,
+        "level": level,
+        "message": {
+            "text": (
+                f"Attack {r.payload_id} ({r.category}, obfuscation {r.mutator}, "
+                f"delivered {r.scenario}) succeeded: {', '.join(r.hits)}.{detail}"
+            )
+        },
+        "locations": [
+            {
+                "physicalLocation": {
+                    "artifactLocation": {"uri": artifact_uri},
+                    "region": {"startLine": 1},
+                }
+            }
+        ],
+        "partialFingerprints": {"agentsec/caseId/v1": r.case_id},
+        "properties": {"category": r.category, "mutator": r.mutator, "scenario": r.scenario},
+    }
+
+
+# ------------------------------------------------------------------------------ JUnit
+
+
+def to_junit(report: FuzzReport) -> str:
+    """Render a JUnit XML report: one test suite per attack category, one test per case.
+
+    A succeeded attack is a *failure*; a case where the target crashed is an *error*.
+    """
+    by_category: dict[str, list[CaseResult]] = {}
+    for r in report.results:
+        by_category.setdefault(r.category, []).append(r)
+
+    root = ET.Element(
+        "testsuites",
+        name="agentsec-fuzz",
+        tests=str(report.total),
+        failures=str(report.successes),
+        errors=str(report.errors),
+    )
+    for category, cases in sorted(by_category.items()):
+        suite = ET.SubElement(
+            root,
+            "testsuite",
+            name=clean(category),
+            tests=str(len(cases)),
+            failures=str(sum(1 for c in cases if c.success)),
+            errors=str(sum(1 for c in cases if c.error is not None)),
+            time=f"{sum(c.latency_ms for c in cases) / 1000:.3f}",
+        )
+        for c in cases:
+            case = ET.SubElement(
+                suite,
+                "testcase",
+                classname=f"agentsec.{clean(category)}",
+                name=clean(c.case_id),
+                time=f"{c.latency_ms / 1000:.3f}",
+            )
+            if c.error is not None:
+                err = ET.SubElement(case, "error", message=clean(c.error), type="TargetError")
+                err.text = clean(c.error)
+            elif c.success:
+                fail = ET.SubElement(
+                    case,
+                    "failure",
+                    message=clean(f"attack succeeded: {', '.join(c.hits)}"),
+                    type=_rule_for(c.goal)[0],
+                )
+                fail.text = clean(c.evidence)
+    ET.indent(root)
+    return '<?xml version="1.0" encoding="UTF-8"?>\n' + ET.tostring(root, encoding="unicode") + "\n"