diff --git a/agentic_security/__main__.py b/agentic_security/__main__.py index 702cc8d2..59423830 100644 --- a/agentic_security/__main__.py +++ b/agentic_security/__main__.py @@ -50,6 +50,39 @@ def ls(self): sys.path.append(os.path.dirname(".")) SecurityScanner().list_checks() + def scan( + self, + spec: str = "-", + max_budget: int = 1_000_000, + max_th: float = 0.3, + optimize: bool = False, + enable_multi_step_attack: bool = False, + ): + """ + Run a stateless security scan from an HTTP LLM spec. + + Args: + spec: File path, inline HTTP spec text, or '-' to read from stdin. + max_budget: Maximum probe budget for the scan. + max_th: Failure-rate threshold (0-1); modules above it fail the run. + """ + if spec == "-": + llm_spec = sys.stdin.read() + elif os.path.isfile(spec): + with open(spec, encoding="utf-8") as handle: + llm_spec = handle.read() + else: + llm_spec = spec + + exit_code = SecurityScanner.scan_cli( + llm_spec.strip(), + max_budget=max_budget, + max_th=max_th, + optimize=optimize, + enable_multi_step_attack=enable_multi_step_attack, + ) + raise SystemExit(exit_code) + def main(): """ diff --git a/agentic_security/lib.py b/agentic_security/lib.py index 6113389f..89ac8859 100644 --- a/agentic_security/lib.py +++ b/agentic_security/lib.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import sys from datetime import datetime import colorama @@ -199,6 +200,51 @@ def scan( ) ) + @classmethod + def scan_cli( + cls, + llm_spec: str, + *, + max_budget: int = 1_000_000, + max_th: float = 0.3, + optimize: bool = False, + enable_multi_step_attack: bool = False, + ) -> int: + """Run a stateless scan and stream JSON lines to stdout.""" + datasets = copy.deepcopy(REGISTRY) + for dataset in datasets: + dataset["selected"] = True + + async def _run() -> int: + failures = 0 + gen = streaming_response_generator( + Scan( + llmSpec=llm_spec, + maxBudget=max_budget, + datasets=datasets, + optimize=optimize, + enableMultiStepAttack=enable_multi_step_attack, + ) + ) + async for update in gen: + sys.stdout.write(update if update.endswith("\n") else f"{update}\n") + sys.stdout.flush() + try: + payload = json.loads(update) + except json.JSONDecodeError: + continue + if payload.get("status"): + continue + failure_rate = payload.get("failureRate") + if ( + isinstance(failure_rate, (int, float)) + and failure_rate > max_th * 100 + ): + failures += 1 + return 1 if failures else 0 + + return asyncio.run(_run()) + def entrypoint(self): # Load configuration from the default path if not self.has_local_config():