mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-09-29 03:11:51 +02:00
Merge pull request #340 from rwinkelman/feat/scan-cli-spec-309
feat(cli): add stateless scan command for agent-invocable runs (#309)
This commit is contained in:
@@ -50,6 +50,39 @@ class CLI:
|
||||
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():
|
||||
"""
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
import colorama
|
||||
@@ -199,6 +200,51 @@ class SecurityScanner(SettingsMixin):
|
||||
)
|
||||
)
|
||||
|
||||
@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():
|
||||
|
||||
Reference in New Issue
Block a user