Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions Readme.md
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,23 @@ Content-Type: application/json

Where `<<PROMPT>>` will be replaced with the actual attack vector during the scan, insert the `Bearer XXXXX` header value with your app credentials.

### One-shot CLI scan

Run a scan without starting the web server or creating a configuration file:

```shell
agentic_security scan \
--spec target.http \
--dataset deepset/prompt-injections \
--max-budget 1000 \
--max-th 0.3
```

The command emits JSON Lines on standard output. It exits with `0` when the
scan stays within the threshold, `1` when findings exceed the threshold, and
`2` when the input or scan fails. See the
[CI/CD guide](docs/ci_cd.md#stateless-scans) for input and artifact options.

### Adding LLM integration templates

TBD
Expand Down
12 changes: 10 additions & 2 deletions agentic_security/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,14 @@

ensure_cache_dir()

from .lib import SecurityScanner # noqa: E402

__all__ = ["SecurityScanner", "ensure_cache_dir"]


def __getattr__(name: str):
"""Load the scanner lazily so lightweight CLI commands have no scan side effects."""
if name == "SecurityScanner":
from .lib import SecurityScanner

globals()[name] = SecurityScanner
return SecurityScanner
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
56 changes: 49 additions & 7 deletions agentic_security/__main__.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,7 @@
import os
import sys

import fire
import uvicorn

from agentic_security.app import app
from agentic_security.lib import SecurityScanner
from agentic_security.misc.banner import init_banner
import fire # type: ignore[import-untyped]


class CLI:
Expand All @@ -18,6 +13,10 @@ def server(self, port: int = 8718, host: str = "127.0.0.1"):
port (int): Port number for the server to listen on. Default is 8718.
host (str): Host address for the server. Default is "127.0.0.1".
"""
import uvicorn

from agentic_security.app import app

sys.path.append(os.path.dirname("."))
config = uvicorn.Config(
app, port=port, host=host, log_level="info", reload=True
Expand All @@ -31,13 +30,51 @@ def ci(self):
"""
Run Agentic Security in CI mode.
"""
from agentic_security.lib import SecurityScanner

sys.path.append(os.path.dirname("."))
SecurityScanner().entrypoint()

def scan(
self,
spec: str,
dataset: str,
max_budget: int = 1_000,
max_th: float = 0.3,
optimize: bool = False,
artifacts_dir: str | None = None,
):
"""
Run a stateless scan and stream JSON lines to stdout.

Args:
spec: HTTP spec text, a file path, or "-" to read standard input.
dataset: Registry name, or a comma-separated list of registry names.
max_budget: Maximum scan budget.
max_th: Failure-rate threshold from 0 to 1.
optimize: Enable adaptive scan optimization.
artifacts_dir: Optional directory for CSV artifacts.
"""
os.environ["AGENTIC_SECURITY_STATELESS"] = "1"
from agentic_security.cli_scan import run_scan_command

raise SystemExit(
run_scan_command(
spec=spec,
dataset=dataset,
max_budget=max_budget,
max_th=max_th,
optimize=optimize,
artifacts_dir=artifacts_dir,
)
)

def init(self, host: str = "127.0.0.1", port: int = 8718):
"""
Generate the default CI configuration file.
"""
from agentic_security.lib import SecurityScanner

sys.path.append(os.path.dirname("."))
SecurityScanner().generate_default_settings(host, port)

Expand All @@ -47,6 +84,8 @@ def ls(self):
"""
List all available security checks.
"""
from agentic_security.lib import SecurityScanner

sys.path.append(os.path.dirname("."))
SecurityScanner().list_checks()

Expand All @@ -62,5 +101,8 @@ def main():


if __name__ == "__main__":
init_banner()
if len(sys.argv) < 2 or sys.argv[1] != "scan":
from agentic_security.misc.banner import init_banner

init_banner()
main()
209 changes: 209 additions & 0 deletions agentic_security/cli_scan.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,209 @@
"""Stateless command-line scanning helpers."""

from __future__ import annotations

import asyncio
import copy
import json
import logging
import os
import sys
from collections.abc import Sequence
from pathlib import Path
from typing import TextIO

from rich.console import Console

from agentic_security.http_spec import LLMSpec
from agentic_security.primitives import Scan, ScanResult

EXIT_OK = 0
EXIT_FINDINGS = 1
EXIT_ERROR = 2


class CLIUsageError(ValueError):
"""Raised when a scan command cannot be constructed from its arguments."""


def load_spec(spec: str, stdin: TextIO | None = None) -> str:
"""Load an HTTP spec from a path, standard input, or an inline argument."""
if spec == "-":
content = (stdin or sys.stdin).read()
else:
candidate = Path(spec)
try:
content = (
candidate.read_text(encoding="utf-8") if candidate.is_file() else spec
)
except OSError as exc:
raise CLIUsageError(
f"Could not read HTTP spec from {candidate}: {exc}"
) from exc

if not content.strip():
raise CLIUsageError("HTTP spec is empty.")
return content


def select_datasets(dataset: str | Sequence[str]) -> list[dict]:
"""Resolve one or more comma-separated registry names."""
from agentic_security.probe_data import REGISTRY

values = [dataset] if isinstance(dataset, str) else list(dataset)
names = [
name.strip()
for value in values
for name in str(value).split(",")
if name.strip()
]
if not names:
raise CLIUsageError("At least one --dataset value is required.")

registry = {item["dataset_name"]: item for item in REGISTRY}
unknown = [name for name in names if name not in registry]
if unknown:
joined = ", ".join(unknown)
raise CLIUsageError(
f"Unknown dataset: {joined}. Run `agentic_security ls` to list choices."
)

if "AgenticBackend" in names:
raise CLIUsageError(
"AgenticBackend requires the web server and cannot be used by stateless scan."
)

selected = []
for name in dict.fromkeys(names):
item = copy.deepcopy(registry[name])
item["selected"] = True
selected.append(item)
return selected


def _move_logs_to(stderr: TextIO) -> None:
"""Keep library logs away from the JSON-lines stream."""
from agentic_security import logutils

root_logger = logging.getLogger(logutils.LOGGER_NAME)
for handler in root_logger.handlers:
if hasattr(handler, "console"):
handler.console = Console(file=stderr, color_system=None)
elif isinstance(handler, logging.StreamHandler):
handler.setStream(stderr)


def _scan_events(**kwargs):
"""Import the scan engine only after stateless mode is enabled."""
from agentic_security.probe_actor import fuzzer

return fuzzer.scan_router(**kwargs)


async def stream_scan(
*,
spec_text: str,
datasets: list[dict],
max_budget: int,
max_th: float,
optimize: bool,
artifacts_dir: str | None,
stdout: TextIO,
) -> int:
"""Run a scan and emit one valid JSON object per output line."""
request_factory = LLMSpec.from_string(spec_text)
scan_parameters = Scan(
llmSpec=spec_text,
maxBudget=max_budget,
datasets=datasets,
optimize=optimize,
)

final_failure_rates: dict[str, float] = {}
runtime_failed = False
completion_event: dict | None = None
events = _scan_events(
request_factory=request_factory,
scan_parameters=scan_parameters,
artifacts_dir=artifacts_dir,
)
async for raw_event in events:
event = json.loads(raw_event)
module = str(event.get("module", ""))
if event.get("status") and module == "Scan completed.":
completion_event = event
continue

stdout.write(json.dumps(event, separators=(",", ":")) + "\n")
stdout.flush()

if event.get("status"):
runtime_failed = runtime_failed or module.startswith("Scan failed:")
elif module:
final_failure_rates[module] = float(event.get("failureRate", 0))

if not runtime_failed and not final_failure_rates:
runtime_failed = True
error = json.loads(ScanResult.status_msg("Scan failed: no dataset results."))
stdout.write(json.dumps(error, separators=(",", ":")) + "\n")

if completion_event is not None:
stdout.write(json.dumps(completion_event, separators=(",", ":")) + "\n")
stdout.flush()

if runtime_failed:
return EXIT_ERROR
if any(rate > max_th * 100 for rate in final_failure_rates.values()):
return EXIT_FINDINGS
return EXIT_OK


def run_scan_command(
*,
spec: str,
dataset: str | Sequence[str],
max_budget: int = 1_000,
max_th: float = 0.3,
optimize: bool = False,
artifacts_dir: str | None = None,
stdin: TextIO | None = None,
stdout: TextIO | None = None,
stderr: TextIO | None = None,
) -> int:
"""Validate CLI inputs, execute the scan, and return a process exit code."""
stdout = stdout or sys.stdout
stderr = stderr or sys.stderr
_move_logs_to(stderr)
previous_mode = os.environ.get("AGENTIC_SECURITY_STATELESS")
os.environ["AGENTIC_SECURITY_STATELESS"] = "1"

try:
if max_budget <= 0:
raise CLIUsageError("--max-budget must be greater than zero.")
if not 0 <= max_th <= 1:
raise CLIUsageError("--max-th must be between 0 and 1.")

spec_text = load_spec(spec, stdin)
datasets = select_datasets(dataset)
return asyncio.run(
stream_scan(
spec_text=spec_text,
datasets=datasets,
max_budget=max_budget,
max_th=max_th,
optimize=optimize,
artifacts_dir=artifacts_dir,
stdout=stdout,
)
)
except KeyboardInterrupt:
stderr.write("Scan interrupted.\n")
return 130
except Exception as exc:
stderr.write(f"Scan error: {exc}\n")
return EXIT_ERROR
finally:
if previous_mode is None:
os.environ.pop("AGENTIC_SECURITY_STATELESS", None)
else:
os.environ["AGENTIC_SECURITY_STATELESS"] = previous_mode
26 changes: 21 additions & 5 deletions agentic_security/config.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import os
from functools import lru_cache
from typing import Any

import tomli

Expand All @@ -7,20 +9,34 @@
SETTINGS_VERSION = 2


@lru_cache(maxsize=1)
def settings_var(name: str, default=None):
return get_or_create_config().get_config_value(name, default)
stateless = os.getenv("AGENTIC_SECURITY_STATELESS") == "1"
return _settings_var(name, default, stateless)


@lru_cache(maxsize=128)
def _settings_var(name: str, default, stateless: bool):
return _get_or_create_config(stateless).get_config_value(name, default)


@lru_cache(maxsize=1)
def get_or_create_config():
"""Return an isolated default config for stateless commands."""
stateless = os.getenv("AGENTIC_SECURITY_STATELESS") == "1"
return _get_or_create_config(stateless)


@lru_cache(maxsize=2)
def _get_or_create_config(stateless: bool):
cfg = SettingsMixin()
cfg.get_or_create_config()
if stateless:
cfg.config = {}
else:
cfg.get_or_create_config()
return cfg


class SettingsMixin:
config = {}
config: dict[str, Any] = {}
default_path = "agentic_security.toml"

def get_or_create_config(self) -> bool:
Expand Down
Loading
Loading