diff --git a/agentic_security/models/schemas.py b/agentic_security/models/schemas.py index c4c274a..9996047 100644 --- a/agentic_security/models/schemas.py +++ b/agentic_security/models/schemas.py @@ -20,6 +20,7 @@ class Scan(BaseModel): maxBudget: int datasets: list[dict] = [] optimize: bool = False + enableMultiStepAttack: bool = False class ScanResult(BaseModel): diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index 9b44fab..1d93d17 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -4,14 +4,13 @@ from collections.abc import AsyncGenerator import httpx import pandas as pd +from agentic_security.models.schemas import Scan, ScanResult +from agentic_security.probe_actor.refusal import refusal_heuristic +from agentic_security.probe_data.data import prepare_prompts from loguru import logger from skopt import Optimizer from skopt.space import Real -from agentic_security.models.schemas import ScanResult -from agentic_security.probe_actor.refusal import refusal_heuristic -from agentic_security.probe_data.data import prepare_prompts - async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]: if isinstance(prompts, list): @@ -293,3 +292,31 @@ async def perform_multi_step_scan( logger.exception("Scan failed") yield ScanResult.status_msg(f"Scan failed: {str(e)}") raise e + + +def scan_router( + request_factory, + scan_parameters: Scan, + tools_inbox=None, + stop_event: asyncio.Event = None, +): + + if scan_parameters.enableMultiStepAttack: + return perform_multi_step_scan( + request_factory=request_factory, + max_budget=scan_parameters.maxBudget, + datasets=scan_parameters.datasets, + probe_datasets=scan_parameters.probeDatasets, + tools_inbox=tools_inbox, + optimize=scan_parameters.optimize, + stop_event=stop_event, + ) + else: + return perform_scan( + request_factory=request_factory, + max_budget=scan_parameters.maxBudget, + datasets=scan_parameters.datasets, + tools_inbox=tools_inbox, + optimize=scan_parameters.optimize, + stop_event=stop_event, + ) diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py index 46344e6..4c04455 100644 --- a/agentic_security/routes/scan.py +++ b/agentic_security/routes/scan.py @@ -29,12 +29,10 @@ def streaming_response_generator(scan_parameters: Scan): request_factory = LLMSpec.from_string(scan_parameters.llmSpec) async def _gen(): - async for scan_result in fuzzer.perform_scan( + async for scan_result in fuzzer.scan_router( request_factory=request_factory, - max_budget=scan_parameters.maxBudget, - datasets=scan_parameters.datasets, + scan_parameters=scan_parameters, tools_inbox=get_tools_inbox(), - optimize=scan_parameters.optimize, stop_event=get_stop_event(), ): yield scan_result + "\n" diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html index 02c4ae3..a7e7ed2 100644 --- a/agentic_security/static/index.html +++ b/agentic_security/static/index.html @@ -268,21 +268,6 @@ concurrently. This can significantly reduce the total scan time but may increase resource usage.
- -- When enabled, the scan will attempt multi-step attack - simulations, - increasing accuracy and depth of analysis. -
@@ -304,13 +289,29 @@+ When enabled, the scan will attempt Many-shot jailbreaking + simulations +
+