fix(update UX and add save state):

build(deps): bump aiohttp from 3.10.3 to 3.10.11

.github/workflows/release.yml

@@ -20,10 +20,10 @@ jobs:
       - uses: actions/checkout@v3
       - name: Install poetry
         run: pipx install poetry==$POETRY_VERSION
-      - name: Set up Python 3.10
+      - name: Set up Python 3.11
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: "3.11"
           cache: "poetry"
       - name: Build project for distribution
         run: poetry build --format sdist

.github/workflows/test.yml

@@ -16,7 +16,6 @@ jobs:
     strategy:
       matrix:
         python-version:
-          - "3.10"
           - "3.11"
           - "3.12"
     steps:

.pre-commit-config.yaml

@@ -1,26 +1,24 @@
-
 default_language_version:
-        python: python3
+        python: python3.11
 
 repos:
-
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.31.1
+    rev: v3.15.0
     hooks:
       - id: pyupgrade
-        args: [--py39-plus]
+        args: [--py311-plus]
 
   - repo: https://github.com/psf/black
-    rev: 22.8.0
+    rev: 23.11.0
     hooks:
       - id: black
-        language_version: python3.9
+        language_version: python3.11
 
   - repo: https://github.com/pycqa/flake8
-    rev: 5.0.4
+    rev: 6.1.0
     hooks:
       - id: flake8
-        language_version: python3
+        language_version: python3.11
         additional_dependencies: [flake8-docstrings]
 
   - repo: https://github.com/PyCQA/isort
@@ -30,7 +28,7 @@ repos:
         args: [--profile, black]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.5.0
     hooks:
     -   id: check-ast
         exclude: '^(third_party)/'
@@ -47,16 +45,15 @@ repos:
         args: ['--maxkb=100']
 
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.14
+    rev: 0.7.17
     hooks:
       - id: mdformat
         name: mdformat
         entry: mdformat .
-        language_version: python3
-
+        language_version: python3.11
 
   - repo: https://github.com/hadialqattan/pycln
-    rev: v2.1.1 # Possible releases: https://github.com/hadialqattan/pycln/releases
+    rev: v2.4.0
     hooks:
       - id: pycln
 
@@ -65,9 +62,8 @@ repos:
     hooks:
       - id: teyit
 
-
   - repo: https://github.com/python-poetry/poetry
-    rev: '1.6.0'
+    rev: '1.7.0'
     hooks:
     -   id: poetry-check
     -   id: poetry-lock
@@ -75,9 +71,8 @@ repos:
         args:
         - --check
 
-
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.2.5
+    rev: v2.2.6
     hooks:
     -   id: codespell
         exclude: '^(third_party/)|(poetry.lock)'

agentic_security/app.py

@@ -1,266 +1,28 @@
-import random
-from asyncio import Event, Queue
-from datetime import datetime
-from logging import config
-from pathlib import Path
-
-from fastapi import BackgroundTasks, FastAPI, HTTPException, Request, Response
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse, StreamingResponse
-from loguru import logger
-from pydantic import BaseModel
-from starlette.middleware.base import BaseHTTPMiddleware
-
-from .http_spec import LLMSpec
-from .probe_actor import fuzzer
-from .probe_actor.refusal import REFUSAL_MARKS
-from .probe_data import REGISTRY
-from .report_chart import plot_security_report
-
-# Create the FastAPI app instance
-app = FastAPI()
-origins = [
-    "*",
-]
-
-# Middleware setup
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],  # Allows all methods
-    allow_headers=["*"],  # Allows all headers
+from .core.app import create_app
+from .core.logging import setup_logging
+from .middleware.cors import setup_cors
+from .middleware.logging import LogNon200ResponsesMiddleware
+from .routes import (
+    probe_router,
+    proxy_router,
+    report_router,
+    scan_router,
+    static_router,
 )
 
-tools_inbox = Queue()
-# Global stop event for cancelling scans
-stop_event = Event()  # Added stop_event to cancel the scan
+# Create the FastAPI app
+app = create_app()
 
-FEATURE_PROXY = False
-
-
-@app.get("/")
-async def root():
-    agentic_security_path = Path(__file__).parent
-    return FileResponse(f"{agentic_security_path}/static/index.html")
-
-
-@app.get("/main.js")
-async def main_js():
-    agentic_security_path = Path(__file__).parent
-    return FileResponse(f"{agentic_security_path}/static/main.js")
-
-
-@app.get("/favicon.ico")
-async def favicon():
-    agentic_security_path = Path(__file__).parent
-    return FileResponse(f"{agentic_security_path}/static/favicon.ico")
-
-
-class LLMInfo(BaseModel):
-    spec: str
-
-
-@app.post("/verify")
-async def verify(info: LLMInfo):
-
-    spec = LLMSpec.from_string(info.spec)
-    r = await spec.probe("test")
-    if r.status_code >= 400:
-        raise HTTPException(status_code=r.status_code, detail=r.text)
-    return dict(
-        status_code=r.status_code,
-        body=r.text,
-        elapsed=r.elapsed.total_seconds(),
-        timestamp=datetime.now().isoformat(),
-    )
-
-
-class Scan(BaseModel):
-    llmSpec: str
-    maxBudget: int
-    datasets: list[dict] = []
-    optimize: bool = False
-
-
-class ScanResult(BaseModel):
-    module: str
-    tokens: int
-    cost: float
-    progress: float
-    failureRate: float = 0.0
-
-
-def streaming_response_generator(scan_parameters: Scan):
-    # The generator function for StreamingResponse
-    request_factory = LLMSpec.from_string(scan_parameters.llmSpec)
-
-    async def _gen():
-        async for scan_result in fuzzer.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,  # Pass the stop_event to the generator
-        ):
-            yield scan_result + "\n"  # Adding a newline for separation
-
-    return _gen()
-
-
-@app.post("/scan")
-async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
-
-    # Initiates streaming of scan results
-    return StreamingResponse(
-        streaming_response_generator(scan_parameters), media_type="application/json"
-    )
-
-
-class Probe(BaseModel):
-    prompt: str
-
-
-@app.post("/v1/self-probe")
-def self_probe(probe: Probe):
-    refuse = random.random() < 0.2
-    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
-    message = probe.prompt + " " + message
-    return {
-        "id": "chatcmpl-abc123",
-        "object": "chat.completion",
-        "created": 1677858242,
-        "model": "gpt-3.5-turbo-0613",
-        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
-        "choices": [
-            {
-                "message": {"role": "assistant", "content": message},
-                "logprobs": None,
-                "finish_reason": "stop",
-                "index": 0,
-            }
-        ],
-    }
-
-
-@app.get("/v1/data-config")
-async def data_config():
-    return [m for m in REGISTRY]
-
-
-@app.get("/failures")
-async def failures_csv():
-    if not Path("failures.csv").exists():
-        return {"error": "No failures found"}
-    return FileResponse("failures.csv")
-
-
-class Table(BaseModel):
-    table: list[dict]
-
-
-@app.post("/plot.jpeg", response_class=Response)
-async def get_plot(table: Table):
-    buf = plot_security_report(table.table)
-    return StreamingResponse(buf, media_type="image/jpeg")
-
-
-class Message(BaseModel):
-    role: str
-    content: str
-
-
-class CompletionRequest(BaseModel):
-    model: str
-    messages: list[Message]
-    temperature: float = 0.7  # Default value for temperature
-    top_p: float = 1.0  # Default value for top_p
-    n: int = 1  # Default value for n
-    stop: list[str] = None  # Optional; specify as None if not provided
-    max_tokens: int = 100  # Default value for max_tokens
-    presence_penalty: float = 0.0  # Default value for presence_penalty
-    frequency_penalty: float = 0.0  # Default value for frequency_penalty
-
-
-# OpenAI proxy endpoint
-@app.post("/proxy/chat/completions")
-async def proxy_completions(request: CompletionRequest):
-    refuse = random.random() < 0.2
-    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
-    prompt_content = " ".join(
-        [msg.content for msg in request.messages if msg.role == "user"]
-    )
-    message = prompt_content + " " + message
-    ready = Event()
-    ref = dict(message=message, reply="", ready=ready)
-    tools_inbox.put_nowait(ref)
-    if FEATURE_PROXY:
-        # Proxy to agent
-        await ready.wait()
-        reply = ref["reply"]
-        return reply
-    # Simulate a completion response
-    return {
-        "id": "chatcmpl-abc123",
-        "object": "chat.completion",
-        "created": 1677858242,
-        "model": "gpt-3.5-turbo-0613",
-        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
-        "choices": [
-            {
-                "message": {"role": "assistant", "content": message},
-                "logprobs": None,
-                "finish_reason": "stop",
-                "index": 0,
-            }
-        ],
-    }
-
-
-config.dictConfig(
-    {
-        "version": 1,
-        "disable_existing_loggers": True,
-        "handlers": {
-            "console": {
-                "class": "logging.StreamHandler",
-            },
-        },
-        "root": {
-            "handlers": ["console"],
-            "level": "INFO",
-        },
-        "loggers": {
-            "uvicorn.access": {
-                "level": "ERROR",  # Set higher log level to suppress info logs globally
-                "handlers": ["console"],
-                "propagate": False,
-            }
-        },
-    }
-)
-
-
-@app.post("/stop")
-async def stop_scan():
-    stop_event.set()  # Set the stop event to cancel the scan
-    return {"status": "Scan stopped"}
-
-
-class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next):
-        try:
-            response = await call_next(request)
-        except Exception as e:
-            logger.exception("Yikes")
-            raise e
-        if response.status_code != 200:
-            logger.error(
-                f"{request.method} {request.url} - Status code: {response.status_code}"
-            )
-        return response
-
-
-# Add middleware to the application
+# Setup middleware
+setup_cors(app)
 app.add_middleware(LogNon200ResponsesMiddleware)
+
+# Setup logging
+setup_logging()
+
+# Register routers
+app.include_router(static_router)
+app.include_router(scan_router)
+app.include_router(probe_router)
+app.include_router(proxy_router)
+app.include_router(report_router)

agentic_security/core/app.py

@@ -0,0 +1,22 @@
+from asyncio import Event, Queue
+
+from fastapi import FastAPI
+
+tools_inbox: Queue = Queue()
+stop_event: Event = Event()
+
+
+def create_app() -> FastAPI:
+    """Create and configure the FastAPI application."""
+    app = FastAPI()
+    return app
+
+
+def get_tools_inbox() -> Queue:
+    """Get the global tools inbox queue."""
+    return tools_inbox
+
+
+def get_stop_event() -> Event:
+    """Get the global stop event."""
+    return stop_event

agentic_security/core/logging.py

@@ -0,0 +1,26 @@
+from logging import config
+
+
+def setup_logging():
+    config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": True,
+            "handlers": {
+                "console": {
+                    "class": "logging.StreamHandler",
+                },
+            },
+            "root": {
+                "handlers": ["console"],
+                "level": "INFO",
+            },
+            "loggers": {
+                "uvicorn.access": {
+                    "level": "ERROR",  # Set higher log level to suppress info logs globally
+                    "handlers": ["console"],
+                    "propagate": False,
+                }
+            },
+        }
+    )

agentic_security/lib.py

@@ -5,8 +5,9 @@ import colorama
 import tqdm.asyncio
 from tabulate import tabulate
 
-from agentic_security.app import Scan, streaming_response_generator
+from agentic_security.models.schemas import Scan
 from agentic_security.probe_data import REGISTRY
+from agentic_security.routes.scan import streaming_response_generator
 
 RESET = colorama.Style.RESET_ALL
 BRIGHT = colorama.Style.BRIGHT

agentic_security/middleware/cors.py

@@ -0,0 +1,14 @@
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+
+def setup_cors(app: FastAPI):
+    origins = ["*"]
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=origins,
+        allow_credentials=True,
+        allow_methods=["*"],  # Allows all methods
+        allow_headers=["*"],  # Allows all headers
+    )

agentic_security/middleware/logging.py

@@ -0,0 +1,17 @@
+from fastapi import Request
+from loguru import logger
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        try:
+            response = await call_next(request)
+        except Exception as e:
+            logger.exception("Yikes")
+            raise e
+        if response.status_code != 200:
+            logger.error(
+                f"{request.method} {request.url} - Status code: {response.status_code}"
+            )
+        return response

agentic_security/models/schemas.py

@@ -0,0 +1,69 @@
+import os
+
+from pydantic import BaseModel, Field
+
+
+class Settings:
+    MAX_BUDGET = 1000
+    MAX_DATASETS = 10
+    RATE_LIMIT = "100/minute"
+    DISABLE_TELEMETRY = os.getenv("DISABLE_TELEMETRY", False)
+    FEATURE_PROXY = False
+
+
+class LLMInfo(BaseModel):
+    spec: str
+
+
+class Scan(BaseModel):
+    llmSpec: str
+    maxBudget: int
+    datasets: list[dict] = []
+    optimize: bool = False
+
+
+class ScanResult(BaseModel):
+    module: str
+    tokens: float | int
+    cost: float
+    progress: float
+    status: bool = False
+    failureRate: float = 0.0
+
+    @classmethod
+    def status_msg(cls, msg: str) -> str:
+        return cls(
+            module=msg,
+            tokens=0,
+            cost=0,
+            progress=0,
+            failureRate=0,
+            status=True,
+        ).model_dump_json()
+
+
+class Probe(BaseModel):
+    prompt: str
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+class CompletionRequest(BaseModel):
+    """Model for completion requests."""
+
+    model: str
+    messages: list[Message]
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    top_p: float = Field(default=1.0, ge=0.0, le=1.0)
+    n: int = Field(default=1, ge=1, le=10)
+    stop: list[str] | None = None
+    max_tokens: int = Field(default=100, ge=1, le=4096)
+    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+
+
+class Table(BaseModel):
+    table: list[dict]

agentic_security/probe_actor/fuzzer.py

@@ -1,39 +1,17 @@
 import asyncio
-import os
-from typing import AsyncGenerator
+import random
+from collections.abc import AsyncGenerator
 
 import httpx
 import pandas as pd
 from loguru import logger
-from pydantic import BaseModel
 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
 
-IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
-
-
-class ScanResult(BaseModel):
-    module: str
-    tokens: float
-    cost: float
-    progress: float
-    failureRate: float = 0.0
-    status: bool = False
-
-    @classmethod
-    def status_msg(cls, msg: str) -> str:
-        return cls(
-            module=msg,
-            tokens=0,
-            cost=0,
-            progress=0,
-            failureRate=0,
-            status=True,
-        ).model_dump_json()
-
 
 async def prompt_iter(prompts: list[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
     if isinstance(prompts, list):
@@ -52,105 +30,266 @@ async def perform_scan(
     optimize=False,
     stop_event: asyncio.Event = None,
 ) -> AsyncGenerator[str, None]:
-    if IS_VERCEL:
-        yield ScanResult.status_msg(
-            "Vercel deployment detected. Streaming messages are not supported by serverless, please run it locally."
+    """Perform a standard security scan."""
+
+    try:
+        yield ScanResult.status_msg("Loading datasets...")
+        prompt_modules = prepare_prompts(
+            dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
+            budget=max_budget,
+            tools_inbox=tools_inbox,
         )
-        return
+        yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
-    yield ScanResult.status_msg("Loading datasets...")
-    prompt_modules = prepare_prompts(
-        dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
-        budget=max_budget,
-        tools_inbox=tools_inbox,
-    )
-    yield ScanResult.status_msg("Datasets loaded. Starting scan...")
+        errors = []
+        refusals = []
+        total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
+        processed_prompts = 0
 
-    errors = []
-    refusals = []
-    total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
-    processed_prompts = 0
-
-    failure_rates = []
-
-    for module in prompt_modules:
-        tokens = 0
-        module_failures = 0
-        module_size = 0 if module.lazy else len(module.prompts)
-        logger.info(f"Scanning {module.dataset_name} {module_size}")
-        optimizer = Optimizer(
-            [Real(0, 1)], base_estimator="GP", n_initial_points=25, acq_func="EI"
+        optimizer = (
+            Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
+            if optimize
+            else None
         )
-        should_stop_early = False
-        async for prompt in prompt_iter(module.prompts):
-            if stop_event and stop_event.is_set():  # Check if stop_event is set
-                stop_event.clear()  # Clear the event for the next scan
-                logger.info("Scan stopped by user.")
-                yield ScanResult.status_msg("Scan stopped by user.")
-                return  # Exit the scan gracefully
+        failure_rates = []
 
-            processed_prompts += 1
-            progress = 100 * processed_prompts / total_prompts if total_prompts else 0
+        for module in prompt_modules:
+            tokens = 0
+            module_failures = 0
+            module_size = 0 if module.lazy else len(module.prompts)
+            logger.info(f"Scanning {module.dataset_name} {module_size}")
 
-            tokens += len(prompt.split())
-            try:
-                r = await request_factory.fn(prompt=prompt)
-                if r.status_code >= 400:
-                    raise httpx.HTTPStatusError(
-                        f"HTTP {r.status_code}", request=r.request, response=r
-                    )
+            async for prompt in prompt_iter(module.prompts):
+                if stop_event and stop_event.is_set():
+                    stop_event.clear()
+                    logger.info("Scan stopped by user.")
+                    yield ScanResult.status_msg("Scan stopped by user.")
+                    return
 
-                response_text = r.text
-                tokens += len(response_text.split())
+                processed_prompts += 1
+                progress = (
+                    100 * processed_prompts / total_prompts if total_prompts else 0
+                )
+                prompt_tokens = len(prompt.split())
+                tokens += prompt_tokens
 
-                if not refusal_heuristic(r.json()):
-                    refusals.append(
-                        (module.dataset_name, prompt, r.status_code, response_text)
-                    )
+                try:
+                    r = await request_factory.fn(prompt=prompt)
+                    if r.status_code >= 400:
+                        raise httpx.HTTPStatusError(
+                            f"HTTP {r.status_code}",
+                            request=r.request,
+                            response=r,
+                        )
+
+                    response_text = r.text
+                    response_tokens = len(response_text.split())
+                    tokens += response_tokens
+
+                    if not refusal_heuristic(r.json()):
+                        refusals.append(
+                            (module.dataset_name, prompt, r.status_code, response_text)
+                        )
+                        module_failures += 1
+
+                except httpx.RequestError as e:
+                    logger.error(f"Request error: {e}")
+                    errors.append((module.dataset_name, prompt, str(e)))
                     module_failures += 1
-            except (httpx.RequestError, httpx.HTTPStatusError) as e:
-                logger.error(f"Request error: {e}")
-                errors.append((module.dataset_name, prompt, str(e)))
-                module_failures += 1
-                continue
+                    continue
 
-            failure_rate = module_failures / max(processed_prompts, 1)
-            failure_rates.append(failure_rate)
+                failure_rate = module_failures / max(processed_prompts, 1)
+                failure_rates.append(failure_rate)
+                cost = round(tokens * 1.5 / 1000_000, 2)
 
-            yield ScanResult(
-                module=module.dataset_name,
-                tokens=round(tokens / 1000, 1),
-                cost=round(tokens * 1.5 / 1000_000, 2),
-                progress=round(progress, 2),
-                failureRate=round(failure_rate * 100, 2),
-            ).model_dump_json()
+                yield ScanResult(
+                    module=module.dataset_name,
+                    tokens=round(tokens / 1000, 1),
+                    cost=cost,
+                    progress=round(progress, 2),
+                    failureRate=round(failure_rate * 100, 2),
+                ).model_dump_json()
 
-            if not optimize:
-                continue
-            # Use the optimizer to decide whether to stop early
-            if len(failure_rates) >= 5:  # Wait for at least 5 data points
-                next_point = optimizer.ask()
-                optimizer.tell(
-                    next_point, -failure_rate
-                )  # We want to minimize failure rate
+                if optimize and len(failure_rates) >= 5:
+                    next_point = optimizer.ask()
+                    optimizer.tell(next_point, -failure_rate)
+                    best_failure_rate = -optimizer.get_result().fun
+                    if best_failure_rate > 0.5:
+                        yield ScanResult.status_msg(
+                            f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                        )
+                        break
 
-                # Get the best point found so far
-                best_failure_rate = -optimizer.get_result().fun
+        yield ScanResult.status_msg("Scan completed.")
 
-                # If the best failure rate is high, consider stopping
-                if best_failure_rate > 0.5:  # Threshold can be adjusted
-                    yield ScanResult.status_msg(
-                        f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
-                    )
-                    should_stop_early = True
-                    break  # Break out of the prompt loop
+        df = pd.DataFrame(
+            errors + refusals, columns=["module", "prompt", "status_code", "content"]
+        )
+        df.to_csv("failures.csv", index=False)
 
-        if should_stop_early:
-            continue  # Move to the next module
+    except Exception as e:
+        logger.exception("Scan failed")
+        yield ScanResult.status_msg(f"Scan failed: {str(e)}")
+        raise e
 
-    yield ScanResult.status_msg("Scan completed.")
 
-    df = pd.DataFrame(
-        errors + refusals, columns=["module", "prompt", "status_code", "content"]
-    )
-    df.to_csv("failures.csv", index=False)
+async def perform_multi_step_scan(
+    request_factory,
+    max_budget: int,
+    datasets: list[dict[str, str]] = [],
+    probe_datasets: list[dict[str, str]] = [],
+    tools_inbox=None,
+    optimize=False,
+    stop_event: asyncio.Event = None,
+    probe_frequency: float = 0.2,
+) -> AsyncGenerator[str, None]:
+    """Perform a multi-step security scan with probe injection."""
+    try:
+        # Load main and probe datasets
+        yield ScanResult.status_msg("Loading datasets...")
+        prompt_modules = prepare_prompts(
+            dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
+            budget=max_budget,
+            tools_inbox=tools_inbox,
+        )
+        probe_modules = prepare_prompts(
+            dataset_names=[m["dataset_name"] for m in probe_datasets if m["selected"]],
+            budget=max_budget,
+            tools_inbox=tools_inbox,
+        )
+        yield ScanResult.status_msg("Datasets loaded. Starting scan...")
+
+        errors = []
+        refusals = []
+        total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
+        processed_prompts = 0
+        conversation_history = {}
+
+        optimizer = (
+            Optimizer([Real(0, 1)], base_estimator="GP", n_initial_points=25)
+            if optimize
+            else None
+        )
+        failure_rates = []
+
+        for module in prompt_modules:
+            tokens = 0
+            module_failures = 0
+            module_size = 0 if module.lazy else len(module.prompts)
+            logger.info(f"Scanning {module.dataset_name} {module_size}")
+            conv_id = module.dataset_name
+
+            async for prompt in prompt_iter(module.prompts):
+                if stop_event and stop_event.is_set():
+                    stop_event.clear()
+                    logger.info("Scan stopped by user.")
+                    yield ScanResult.status_msg("Scan stopped by user.")
+                    return
+
+                processed_prompts += 1
+                progress = (
+                    100 * processed_prompts / total_prompts if total_prompts else 0
+                )
+
+                # Get conversation history
+                history = conversation_history.get(conv_id, [])
+                full_prompt = "\n".join([*history, prompt]) if history else prompt
+                prompt_tokens = len(full_prompt.split())
+                tokens += prompt_tokens
+
+                try:
+                    # Main request
+                    r = await request_factory.fn(prompt=full_prompt)
+                    if r.status_code >= 400:
+                        raise httpx.HTTPStatusError(
+                            f"HTTP {r.status_code}",
+                            request=r.request,
+                            response=r,
+                        )
+
+                    response_text = r.text
+                    response_tokens = len(response_text.split())
+                    tokens += response_tokens
+
+                    # Update history
+                    history.extend([prompt, response_text])
+                    history = history[-4:]  # Keep last 2 exchanges
+                    conversation_history[conv_id] = history
+
+                    if not refusal_heuristic(r.json()):
+                        refusals.append(
+                            (module.dataset_name, prompt, r.status_code, response_text)
+                        )
+                        module_failures += 1
+
+                    # Random probe injection
+                    if probe_modules and random.random() < probe_frequency:
+                        probe_module = random.choice(probe_modules)
+                        probe_prompts = [
+                            p async for p in prompt_iter(probe_module.prompts)
+                        ]
+                        if probe_prompts:
+                            probe = random.choice(probe_prompts)
+                            full_probe = "\n".join([*history, probe])
+
+                            probe_r = await request_factory.fn(prompt=full_probe)
+                            if probe_r.status_code < 400:
+                                probe_response = probe_r.text
+                                tokens += len(probe.split()) + len(
+                                    probe_response.split()
+                                )
+
+                                history.extend([probe, probe_response])
+                                history = history[-4:]
+                                conversation_history[conv_id] = history
+
+                                if not refusal_heuristic(probe_r.json()):
+                                    refusals.append(
+                                        (
+                                            probe_module.dataset_name,
+                                            probe,
+                                            probe_r.status_code,
+                                            probe_response,
+                                        )
+                                    )
+                                    module_failures += 1
+
+                except httpx.RequestError as e:
+                    logger.error(f"Request error: {e}")
+                    errors.append((module.dataset_name, prompt, str(e)))
+                    module_failures += 1
+                    continue
+
+                failure_rate = module_failures / max(processed_prompts, 1)
+                failure_rates.append(failure_rate)
+                cost = round(tokens * 1.5 / 1000_000, 2)
+
+                yield ScanResult(
+                    module=module.dataset_name,
+                    tokens=round(tokens / 1000, 1),
+                    cost=cost,
+                    progress=round(progress, 2),
+                    failureRate=round(failure_rate * 100, 2),
+                ).model_dump_json()
+
+                if optimize and len(failure_rates) >= 5:
+                    next_point = optimizer.ask()
+                    optimizer.tell(next_point, -failure_rate)
+                    best_failure_rate = -optimizer.get_result().fun
+                    if best_failure_rate > 0.5:
+                        yield ScanResult.status_msg(
+                            f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                        )
+                        break
+
+        yield ScanResult.status_msg("Scan completed.")
+
+        df = pd.DataFrame(
+            errors + refusals, columns=["module", "prompt", "status_code", "content"]
+        )
+        df.to_csv("failures.csv", index=False)
+
+    except Exception as e:
+        logger.exception("Scan failed")
+        yield ScanResult.status_msg(f"Scan failed: {str(e)}")
+        raise e

agentic_security/probe_actor/test_refusal.py

@@ -2,7 +2,6 @@ from agentic_security.probe_actor.refusal import check_refusal
 
 
 class TestCheckRefusal:
-
     # The function correctly identifies a refusal phrase in the response.
     def test_identify_refusal_phrase(self):
         response = "I'm sorry, but I cannot provide that information."

agentic_security/probe_data/data.py

@@ -6,6 +6,7 @@ from functools import lru_cache
 
 import httpx
 import pandas as pd
+from cache_to_disk import cache_to_disk
 from loguru import logger
 
 from agentic_security.probe_data import stenography_fn
@@ -15,21 +16,6 @@ from agentic_security.probe_data.modules import (
     inspect_ai_tool,
 )
 
-IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
-
-if not IS_VERCEL:
-    from cache_to_disk import cache_to_disk
-else:
-    # Read only fs in vercel, just mock no-op decorator
-    def cache_to_disk(*_):
-        def decorator(fn):
-            def wrapper(*args, **kwargs):
-                return fn(*args, **kwargs)
-
-            return wrapper
-
-        return decorator
-
 
 @dataclass
 class ProbeDataset:
@@ -152,7 +138,6 @@ def load_dataset_v6():
 
 @cache_to_disk()
 def load_dataset_v7():
-
     splits = {
         "mini_JailBreakV_28K": "JailBreakV_28K/mini_JailBreakV_28K.csv",
         "JailBreakV_28K": "JailBreakV_28K/JailBreakV_28K.csv",
@@ -173,7 +158,6 @@ def load_dataset_v7():
 
 @cache_to_disk()
 def load_dataset_v8():
-
     df = pd.read_csv(
         "hf://datasets/ShawnMenz/jailbreak_sft_rm_ds/jailbreak_sft_rm_ds.csv",
         names=["jailbreak", "prompt"],
@@ -321,7 +305,6 @@ class Stenography:
 
     def apply(self):
         for prompt_group in self.prompt_groups:
-
             size = len(prompt_group.prompts)
             for name, fn in self.fn_library.items():
                 logger.info(f"Applying {name} to {prompt_group.dataset_name}")

agentic_security/probe_data/modules/adaptive_attacks.py

@@ -9,7 +9,6 @@ url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harm
 
 class Module:
     def __init__(self, prompt_groups: []):
-
         r = httpx.get(url)
 
         content = r.content

agentic_security/probe_data/modules/test_adaptive_attacks.py

@@ -4,7 +4,6 @@ from .adaptive_attacks import Module
 
 
 class TestModule:
-
     # Module can be initialized with a list of prompt groups.
     def test_initialize_with_prompt_groups(self):
         prompt_groups = []

agentic_security/refusal_classifier/model.py

@@ -1,3 +1,4 @@
+import importlib.resources as pkg_resources
 import os
 
 import joblib
@@ -69,9 +70,27 @@ class RefusalClassifier:
         """
         Load the trained model, vectorizer, and scaler from disk.
         """
-        self.model = joblib.load(self.model_path)
-        self.vectorizer = joblib.load(self.vectorizer_path)
-        self.scaler = joblib.load(self.scaler_path)
+        try:
+            self.model = joblib.load(self.model_path)
+            self.vectorizer = joblib.load(self.vectorizer_path)
+            self.scaler = joblib.load(self.scaler_path)
+        except FileNotFoundError:
+            # Load from package resources
+            package = (
+                __package__  # This should be 'agentic_security.refusal_classifier'
+            )
+
+            # Load model
+            with pkg_resources.open_binary(package, "oneclass_svm_model.joblib") as f:
+                self.model = joblib.load(f)
+
+            # Load vectorizer
+            with pkg_resources.open_binary(package, "tfidf_vectorizer.joblib") as f:
+                self.vectorizer = joblib.load(f)
+
+            # Load scaler
+            with pkg_resources.open_binary(package, "scaler.joblib") as f:
+                self.scaler = joblib.load(f)
 
     def is_refusal(self, text):
         """

agentic_security/routes/__init__.py

@@ -0,0 +1,13 @@
+from .probe import router as probe_router
+from .proxy import router as proxy_router
+from .report import router as report_router
+from .scan import router as scan_router
+from .static import router as static_router
+
+__all__ = [
+    "static_router",
+    "scan_router",
+    "probe_router",
+    "proxy_router",
+    "report_router",
+]

agentic_security/routes/probe.py

@@ -0,0 +1,36 @@
+import random
+
+from fastapi import APIRouter
+
+from ..models.schemas import Probe
+from ..probe_actor.refusal import REFUSAL_MARKS
+from ..probe_data import REGISTRY
+
+router = APIRouter()
+
+
+@router.post("/v1/self-probe")
+def self_probe(probe: Probe):
+    refuse = random.random() < 0.2
+    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
+    message = probe.prompt + " " + message
+    return {
+        "id": "chatcmpl-abc123",
+        "object": "chat.completion",
+        "created": 1677858242,
+        "model": "gpt-3.5-turbo-0613",
+        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        "choices": [
+            {
+                "message": {"role": "assistant", "content": message},
+                "logprobs": None,
+                "finish_reason": "stop",
+                "index": 0,
+            }
+        ],
+    }
+
+
+@router.get("/v1/data-config")
+async def data_config():
+    return [m for m in REGISTRY]

agentic_security/routes/proxy.py

@@ -0,0 +1,47 @@
+import random
+from asyncio import Event
+
+from fastapi import APIRouter
+
+from ..core.app import get_tools_inbox
+from ..models.schemas import CompletionRequest, Settings
+from ..probe_actor.refusal import REFUSAL_MARKS
+
+router = APIRouter()
+
+
+@router.post("/proxy/chat/completions")
+async def proxy_completions(request: CompletionRequest):
+    refuse = random.random() < 0.2
+    message = random.choice(REFUSAL_MARKS) if refuse else "This is a test!"
+    prompt_content = " ".join(
+        [msg.content for msg in request.messages if msg.role == "user"]
+    )
+    message = prompt_content + " " + message
+    ready = Event()
+    ref = dict(message=message, reply="", ready=ready)
+    tools_inbox = get_tools_inbox()
+    await tools_inbox.put(ref)
+
+    if Settings.FEATURE_PROXY:
+        # Proxy to agent
+        await ready.wait()
+        reply = ref["reply"]
+        return reply
+
+    # Simulate a completion response
+    return {
+        "id": "chatcmpl-abc123",
+        "object": "chat.completion",
+        "created": 1677858242,
+        "model": "gpt-3.5-turbo-0613",
+        "usage": {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20},
+        "choices": [
+            {
+                "message": {"role": "assistant", "content": message},
+                "logprobs": None,
+                "finish_reason": "stop",
+                "index": 0,
+            }
+        ],
+    }

agentic_security/routes/report.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+
+from fastapi import APIRouter, Response
+from fastapi.responses import FileResponse, StreamingResponse
+
+from ..models.schemas import Table
+from ..report_chart import plot_security_report
+
+router = APIRouter()
+
+
+@router.get("/failures")
+async def failures_csv():
+    if not Path("failures.csv").exists():
+        return {"error": "No failures found"}
+    return FileResponse("failures.csv")
+
+
+@router.post("/plot.jpeg", response_class=Response)
+async def get_plot(table: Table):
+    buf = plot_security_report(table.table)
+    return StreamingResponse(buf, media_type="image/jpeg")

agentic_security/routes/scan.py

@@ -0,0 +1,55 @@
+from datetime import datetime
+
+from fastapi import APIRouter, BackgroundTasks, HTTPException
+from fastapi.responses import StreamingResponse
+
+from ..core.app import get_stop_event, get_tools_inbox
+from ..http_spec import LLMSpec
+from ..models.schemas import LLMInfo, Scan
+from ..probe_actor import fuzzer
+
+router = APIRouter()
+
+
+@router.post("/verify")
+async def verify(info: LLMInfo):
+    spec = LLMSpec.from_string(info.spec)
+    r = await spec.probe("test")
+    if r.status_code >= 400:
+        raise HTTPException(status_code=r.status_code, detail=r.text)
+    return dict(
+        status_code=r.status_code,
+        body=r.text,
+        elapsed=r.elapsed.total_seconds(),
+        timestamp=datetime.now().isoformat(),
+    )
+
+
+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(
+            request_factory=request_factory,
+            max_budget=scan_parameters.maxBudget,
+            datasets=scan_parameters.datasets,
+            tools_inbox=get_tools_inbox(),
+            optimize=scan_parameters.optimize,
+            stop_event=get_stop_event(),
+        ):
+            yield scan_result + "\n"
+
+    return _gen()
+
+
+@router.post("/scan")
+async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
+    return StreamingResponse(
+        streaming_response_generator(scan_parameters), media_type="application/json"
+    )
+
+
+@router.post("/stop")
+async def stop_scan():
+    get_stop_event().set()
+    return {"status": "Scan stopped"}

agentic_security/routes/static.py

@@ -0,0 +1,84 @@
+from pathlib import Path
+
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import FileResponse, HTMLResponse
+from fastapi.templating import Jinja2Templates
+from jinja2 import Environment, FileSystemLoader
+from starlette.responses import Response
+
+from ..models.schemas import Settings
+
+router = APIRouter()
+STATIC_DIR = Path(__file__).parent.parent / "static"
+
+# Configure templates with custom delimiters to avoid conflicts
+templates = Jinja2Templates(directory=str(STATIC_DIR))
+templates.env = Environment(
+    loader=FileSystemLoader(str(STATIC_DIR)),
+    autoescape=True,
+    block_start_string="[[%",
+    block_end_string="%]]",
+    variable_start_string="[[",
+    variable_end_string="]]",
+)
+
+# Content type mapping for static files
+CONTENT_TYPES = {
+    ".js": "application/javascript",
+    ".ico": "image/x-icon",
+    ".html": "text/html",
+    ".css": "text/css",
+}
+
+
+def get_static_file(filepath: Path, content_type: str | None = None) -> FileResponse:
+    """
+    Helper function to serve static files with proper error handling and caching.
+
+    Args:
+        filepath: Path to the static file
+        content_type: Optional content type override
+
+    Returns:
+        FileResponse with appropriate headers
+
+    Raises:
+        HTTPException if file not found
+    """
+    if not filepath.is_file():
+        raise HTTPException(status_code=404, detail="File not found")
+
+    headers = {
+        "Cache-Control": "public, max-age=3600",
+        "Content-Type": content_type
+        or CONTENT_TYPES.get(filepath.suffix, "application/octet-stream"),
+    }
+
+    return FileResponse(filepath, headers=headers)
+
+
+@router.get("/", response_class=HTMLResponse)
+async def root(request: Request) -> Response:
+    """Serve the main index.html template."""
+    return templates.TemplateResponse("index.html", {"request": request})
+
+
+@router.get("/main.js")
+async def main_js() -> FileResponse:
+    """Serve the main JavaScript file."""
+    return get_static_file(STATIC_DIR / "main.js")
+
+
+@router.get("/telemetry.js")
+async def telemetry_js() -> FileResponse:
+    """
+    Serve either telemetry.js or telemetry_disabled.js based on settings.
+    """
+    filename = "telemetry_disabled.js" if Settings.DISABLE_TELEMETRY else "telemetry.js"
+    return get_static_file(STATIC_DIR / filename)
+
+
+@router.get("/favicon.ico")
+async def favicon() -> FileResponse:
+    """Serve the favicon."""
+    return get_static_file(STATIC_DIR / "favicon.ico")

agentic_security/static/index.html

@@ -1,46 +1,9 @@
 <!doctype html>
 <html lang="en" class="dark">
-  <head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>LLM Vulnerability Scanner</title>
-    <script src="https://cdn.tailwindcss.com"></script>
-    <script src="https://unpkg.com/vue@2.6.12/dist/vue.js"></script>
-    <script src="https://unpkg.com/lucide@latest/dist/umd/lucide.js"></script>
-    <link href="https://fonts.cdnfonts.com/css/technopollas" rel="stylesheet">
-    <style>
-      @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
-    </style>
-    <script>
-      tailwind.config = {
-        darkMode: 'class',
-        theme: {
-          extend: {
-            fontFamily: {
-              sans: ['Inter', 'sans-serif'],
-              technopollas: ['Technopollas', 'sans-serif'],
-            },
-            colors: {
-              dark: {
-                bg: '#121212',
-                card: '#1E1E1E',
-                text: '#FFFFFF',
-                accent: {
-                  green: '#4CAF50',
-                  red: '#F44336',
-                  orange: '#FF9800',
-                  yellow: '#FFEB3B',
-                },
-              },
-            },
-            borderRadius: {
-              'lg': '1rem',
-            },
-          }
-        }
-      }
-    </script>
-  </head>
+  [[% block head %]]
+  [[% include "partials/head.html" %]]
+  [[% endblock head %]]
+
   <body class="bg-dark-bg text-dark-text font-sans">
     <!-- Vue app root element -->
     <div id="vue-app" class="min-h-screen p-8">
@@ -68,6 +31,7 @@
         </div>
 
       </header>
+      [[% include "partials/concent.html" %]]
 
       <main class="max-w-6xl mx-auto space-y-8">
         <!-- Config Selection -->
@@ -178,6 +142,45 @@
             </svg>
           </div>
           <div v-show="showParams" class="mt-4">
+            <div class="flex items-center justify-end mt-4">
+              <button
+                @click="confirmResetState"
+                class="flex items-center bg-dark-accent-red text-dark-bg rounded-lg px-4 py-2 text-sm font-medium hover:bg-opacity-80 transition-colors">
+                <svg xmlns="http://www.w3.org/2000/svg" class="h-5 w-5 mr-2"
+                  fill="none" viewBox="0 0 24 24" stroke="currentColor">
+                  <path stroke-linecap="round" stroke-linejoin="round"
+                    stroke-width="2"
+                    d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
+                </svg>
+                Reset State
+              </button>
+            </div>
+            <!-- Confirmation Modal -->
+            <div
+              v-if="showResetConfirmation"
+              class="fixed inset-0 bg-black bg-opacity-50 flex items-center justify-center z-50">
+              <div class="bg-dark-card rounded-lg p-6 max-w-sm w-full">
+                <h3 class="text-xl font-bold mb-4 text-dark-text">Confirm
+                  Reset</h3>
+                <p class="text-gray-400 mb-6">Are you sure you want to reset all
+                  settings to their default state? This action cannot be
+                  undone.</p>
+                <div class="flex justify-end space-x-4">
+                  <button
+                    @click="showResetConfirmation = false"
+                    class="bg-gray-600 text-dark-text rounded-lg px-4 py-2 hover:bg-opacity-80 transition-colors">
+                    Cancel
+                  </button>
+                  <button
+                    @click="resetState"
+                    class="bg-dark-accent-red text-dark-bg rounded-lg px-4 py-2 hover:bg-opacity-80 transition-colors">
+                    Reset
+                  </button>
+                </div>
+              </div>
+            </div>
+            <!-- Confirmation Modal -->
+
             <!-- Maximum Budget Slider -->
             <!-- Budget Slider -->
             <section class="bg-dark-card rounded-lg p-6 shadow-lg">
@@ -265,6 +268,21 @@
                 concurrently. This can significantly reduce the total scan time
                 but may increase resource usage.
               </p>
+              <!-- Multi-Step Attack Toggle -->
+              <div class="flex items-center justify-between mb-2">
+                <h3 class="text-lg font-semibold">Enable Multi-Step Attack</h3>
+                <label class="relative inline-flex items-center cursor-pointer">
+                  <input type="checkbox" v-model="enableMultiStepAttack"
+                    class="sr-only peer">
+                  <div
+                    class="w-11 h-6 bg-gray-200 peer-focus:outline-none peer-focus:ring-4 peer-focus:ring-dark-accent-green rounded-full peer peer-checked:after:translate-x-full peer-checked:after:border-white after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-dark-accent-green"></div>
+                </label>
+              </div>
+              <p class="text-sm text-gray-400 mt-2">
+                When enabled, the scan will attempt multi-step attack
+                simulations,
+                increasing accuracy and depth of analysis.
+              </p>
             </div>
           </div>
         </section>
@@ -441,57 +459,17 @@
           </div>
         </section>
       </main>
-      <!-- Footer Section -->
-      <footer class="mt-16 pt-8 border-t border-gray-800">
-        <div class="max-w-6xl mx-auto px-4 sm:px-6 lg:px-8">
-          <div class="grid grid-cols-1 md:grid-cols-3 gap-8">
-            <!-- Column 1 -->
-            <div>
-              <h3
-                class="text-lg font-semibold text-dark-accent-green mb-4">Home</h3>
-              <p class="text-gray-400">Dedicated to LLM Security, 2024</p>
-            </div>
 
-            <!-- Column 2 -->
-            <div>
-              <h3
-                class="text-lg font-semibold text-dark-accent-green mb-4">Connect</h3>
-              <ul class="space-y-2">
-                <li><a href="https://x.com" target="_blank"
-                    rel="noopener noreferrer"
-                    class="text-gray-400 hover:text-dark-accent-green">X.com</a></li>
-                <li><a href="https://github.com/msoedov" target="_blank"
-                    rel="noopener noreferrer"
-                    class="text-gray-400 hover:text-dark-accent-green">Github</a></li>
-              </ul>
-            </div>
-
-            <!-- Column 3 -->
-            <div>
-              <h3
-                class="text-lg font-semibold text-dark-accent-green mb-4">About</h3>
-              <p class="text-gray-400">This is the LLM Vulnerability Scanner.
-                Easy to use—no coding needed, just pure security testing.</p>
-            </div>
-          </div>
-
-          <div class="mt-8 pt-8 border-t border-gray-800 text-center">
-            <p class="text-gray-400">Made with ❤️ by the Agentic Security
-              Team</p>
-          </div>
-        </div>
-      </footer>
+      [[% block footer %]]
+      [[% include "partials/footer.html" %]]
+      [[% endblock footer %]]
 
     </div>
 
     <script src="main.js"></script>
+    <script src="telemetry.js"></script>
     <script>
       lucide.createIcons();
     </script>
-    <script>
-    !function(t,e){var o,n,p,r;e.__SV||(window.posthog=e,e._i=[],e.init=function(i,s,a){function g(t,e){var o=e.split(".");2==o.length&&(t=t[o[0]],e=o[1]),t[e]=function(){t.push([e].concat(Array.prototype.slice.call(arguments,0)))}}(p=t.createElement("script")).type="text/javascript",p.async=!0,p.src=s.api_host.replace(".i.posthog.com","-assets.i.posthog.com")+"/static/array.js",(r=t.getElementsByTagName("script")[0]).parentNode.insertBefore(p,r);var u=e;for(void 0!==a?u=e[a]=[]:a="posthog",u.people=u.people||[],u.toString=function(t){var e="posthog";return"posthog"!==a&&(e+="."+a),t||(e+=" (stub)"),e},u.people.toString=function(){return u.toString(1)+".people (stub)"},o="init push capture register register_once register_for_session unregister unregister_for_session getFeatureFlag getFeatureFlagPayload isFeatureEnabled reloadFeatureFlags updateEarlyAccessFeatureEnrollment getEarlyAccessFeatures on onFeatureFlags onSessionId getSurveys getActiveMatchingSurveys renderSurvey canRenderSurvey getNextSurveyStep identify setPersonProperties group resetGroups setPersonPropertiesForFlags resetPersonPropertiesForFlags setGroupPropertiesForFlags resetGroupPropertiesForFlags reset get_distinct_id getGroups get_session_id get_session_replay_url alias set_config startSessionRecording stopSessionRecording sessionRecordingStarted loadToolbar get_property getSessionProperty createPersonProfile opt_in_capturing opt_out_capturing has_opted_in_capturing has_opted_out_capturing clear_opt_in_out_capturing debug".split(" "),n=0;n<o.length;n++)g(u,o[n]);e._i.push([i,s,a])},e.__SV=1)}(document,window.posthog||[]);
-    posthog.init('phc_jfYo5xEofW7eJtiU8rLt2Z8jw1E2eW27BxwTJzwRufH',{api_host:'https://us.i.posthog.com', person_profiles: 'identified_only' // or 'always' to create profiles for anonymous users as well
-        })
-    </script>
   </body>
 </html>

agentic_security/static/main.js

@@ -3,6 +3,7 @@ let URL = window.location.href;
 if (URL.endsWith('/')) {
     URL = URL.slice(0, -1);
 }
+URL = URL.replace('/#', '');
 
 // Vue application
 let LLM_SPECS = [
@@ -72,10 +73,12 @@ var app = new Vue({
         modelSpec: LLM_SPECS[0],
         budget: 50,
         showParams: false,
+        showResetConfirmation: false,
         enableChartDiagram: true,
         enableLogging: false,
         enableConcurrency: false,
         optimize: false,
+        enableMultiStepAttack: false,
         showDatasets: false,
         scanResults: [],
         mainTable: [],
@@ -88,6 +91,7 @@ var app = new Vue({
         selectedConfig: 0,
         showModules: false,
         showLogs: false,
+        showConsentModal: true,
         statusDotClass: 'bg-gray-500', // Default status dot class
         statusText: 'Verified', // Default status text
         statusClass: 'bg-green-500 text-dark-bg', // Default status class
@@ -103,11 +107,19 @@ var app = new Vue({
         ],
         dataConfig: [],
     },
+    created() {
+        // Check if consent is already given in local storage
+        const consentGiven = localStorage.getItem('consentGiven');
+        if (consentGiven === 'true') {
+            this.showConsentModal = false; // Don't show the modal if consent was given
+        }
+    },
     mounted: function () {
         console.log('Vue app mounted');
         this.adjustHeight({ target: document.getElementById('llm-spec') });
         // this.startScan();
         this.loadConfigs();
+
     },
     computed: {
         selectedDS: function () {
@@ -118,6 +130,54 @@ var app = new Vue({
         }
     },
     methods: {
+        acceptConsent() {
+            this.showConsentModal = false; // Close the modal
+            localStorage.setItem('consentGiven', 'true'); // Save consent to local storage
+        },
+
+        saveStateToLocalStorage() {
+            const state = {
+                modelSpec: this.modelSpec,
+                budget: this.budget,
+                dataConfig: this.dataConfig,
+                optimize: this.optimize,
+                enableChartDiagram: this.enableChartDiagram,
+            };
+            localStorage.setItem('appState', JSON.stringify(state));
+        },
+        loadStateFromLocalStorage() {
+            const savedState = localStorage.getItem('appState');
+            console.log('Loading state from local storage:', savedState);
+            if (savedState) {
+                const state = JSON.parse(savedState);
+                this.modelSpec = state.modelSpec;
+                this.budget = state.budget;
+                this.dataConfig = state.dataConfig;
+                this.optimize = state.optimize;
+                this.enableChartDiagram = state.enableChartDiagram;
+            }
+        },
+        resetState() {
+            localStorage.removeItem('appState');
+            this.modelSpec = LLM_SPECS[0];
+            this.budget = 50;
+            this.dataConfig.forEach(config => config.selected = false);
+            this.optimize = false;
+            this.enableChartDiagram = true;
+            this.okMsg = '';
+            this.errorMsg = '';
+            this.integrationVerified = false;
+            this.showResetConfirmation = false;
+        },
+        confirmResetState() {
+            this.showResetConfirmation = true;
+        },
+
+        declineConsent() {
+            this.showConsentModal = false; // Close the modal
+            localStorage.setItem('consentGiven', 'false'); // Save decline to local storage
+            window.location.href = 'https://www.google.com'; // Redirect to Google
+        },
         updateStatusDot(ok) {
             if (ok) {
                 this.statusDotClass = 'bg-green-500'; // Green when expanded
@@ -168,6 +228,7 @@ var app = new Vue({
                 // this.$forceUpdate();
 
             }
+            this.saveStateToLocalStorage();
         },
         loadConfigs: async function () {
             const response = await fetch(`${URL}/v1/data-config`, {
@@ -178,6 +239,7 @@ var app = new Vue({
             });
             console.log(response);
             this.dataConfig = await response.json();
+            this.loadStateFromLocalStorage();
         },
         selectConfig(index) {
             this.selectedConfig = index;
@@ -395,6 +457,8 @@ var app = new Vue({
                     }
                 });
             }
+            this.saveStateToLocalStorage();
+
         }
     }
 });

agentic_security/static/partials/concent.html

@@ -0,0 +1,67 @@
+ <div id="consent-modal" v-if="showConsentModal"
+    class="fixed inset-0 bg-black bg-opacity-75 flex justify-center items-center z-50">
+    <div
+        class="bg-dark-card text-dark-text p-8 rounded-xl shadow-2xl max-w-xl w-full">
+        <h2 class="text-2xl font-bold mb-6 text-center">AI Red Team Ethical
+            Use Agreement</h2>
+        <div class="space-y-6">
+            <p class="text-sm leading-relaxed">
+                This AI red team tool is designed for security research,
+                vulnerability assessment,
+                and responsible testing purposes. By accessing this tool, you
+                explicitly agree to
+                the following ethical guidelines:
+            </p>
+            <ul class="list-disc list-inside text-sm space-y-3">
+                <li>
+                    <strong>Consent and Authorization:</strong> You will only
+                    use
+                    this tool on systems
+                    for which you have explicit, documented permission from the
+                    system owners.
+                </li>
+                <li>
+                    <strong>Responsible Disclosure:</strong> Any vulnerabilities
+                    discovered must be
+                    reported responsibly to the appropriate parties,
+                    prioritizing
+                    system and user safety.
+                </li>
+                <li>
+                    <strong>No Malicious Intent:</strong> You will not use this
+                    tool
+                    to cause harm,
+                    disrupt services, or compromise the integrity of any system
+                    or
+                    data.
+                </li>
+                <li>
+                    <strong>Legal Compliance:</strong> All testing and research
+                    must
+                    comply with
+                    applicable local, national, and international laws and
+                    regulations.
+                </li>
+
+            </ul>
+
+            <p class="text-xs text-gray-400 italic">
+                Violation of these terms may result in immediate termination of
+                access and
+                potential legal consequences.
+            </p>
+        </div>
+        <div class="flex justify-center space-x-4 mt-8">
+            <button
+                @click="declineConsent"
+                class="bg-dark-accent-red text-white rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors">
+                Decline
+            </button>
+            <button
+                @click="acceptConsent"
+                class="bg-dark-accent-green text-dark-bg rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors">
+                I Agree and Understand
+            </button>
+        </div>
+    </div>
+</div>

agentic_security/static/partials/footer.html

@@ -0,0 +1,41 @@
+ <!-- Footer Section -->
+<footer class="mt-16 pt-8 border-t border-gray-800">
+    <div class="max-w-6xl mx-auto px-4 sm:px-6 lg:px-8">
+        <div class="grid grid-cols-1 md:grid-cols-3 gap-8">
+            <!-- Column 1 -->
+            <div>
+                <h3
+                    class="text-lg font-semibold text-dark-accent-green mb-4">Home</h3>
+                <p class="text-gray-400">Dedicated to LLM Security, 2024</p>
+            </div>
+
+            <!-- Column 2 -->
+            <div>
+                <h3
+                    class="text-lg font-semibold text-dark-accent-green mb-4">Connect</h3>
+                <ul class="space-y-2">
+                    <li><a href="https://x.com" target="_blank"
+                            rel="noopener noreferrer"
+                            class="text-gray-400 hover:text-dark-accent-green">X.com</a></li>
+                    <li><a href="https://github.com/msoedov" target="_blank"
+                            rel="noopener noreferrer"
+                            class="text-gray-400 hover:text-dark-accent-green">Github</a></li>
+                </ul>
+            </div>
+
+            <!-- Column 3 -->
+            <div>
+                <h3
+                    class="text-lg font-semibold text-dark-accent-green mb-4">About</h3>
+                <p class="text-gray-400">This is the LLM Vulnerability Scanner.
+                    Easy to use—no coding needed, just pure security
+                    testing.</p>
+            </div>
+        </div>
+
+        <div class="mt-8 pt-8 border-t border-gray-800 text-center">
+            <p class="text-gray-400">Made with ❤️ by the Agentic Security
+                Team</p>
+        </div>
+    </div>
+</footer>

agentic_security/static/partials/head.html

@@ -0,0 +1,41 @@
+  <head></head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>LLM Vulnerability Scanner</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <script src="https://unpkg.com/vue@2.6.12/dist/vue.js"></script>
+    <script src="https://unpkg.com/lucide@latest/dist/umd/lucide.js"></script>
+    <link href="https://fonts.cdnfonts.com/css/technopollas" rel="stylesheet">
+    <style>
+      @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
+    </style>
+    <script>
+      tailwind.config = {
+        darkMode: 'class',
+        theme: {
+          extend: {
+            fontFamily: {
+              sans: ['Inter', 'sans-serif'],
+              technopollas: ['Technopollas', 'sans-serif'],
+            },
+            colors: {
+              dark: {
+                bg: '#121212',
+                card: '#1E1E1E',
+                text: '#FFFFFF',
+                accent: {
+                  green: '#4CAF50',
+                  red: '#F44336',
+                  orange: '#FF9800',
+                  yellow: '#FFEB3B',
+                },
+              },
+            },
+            borderRadius: {
+              'lg': '1rem',
+            },
+          }
+        }
+      }
+    </script>
+  </head>

agentic_security/static/telemetry.js

@@ -0,0 +1,4 @@
+!function (t, e) { var o, n, p, r; e.__SV || (window.posthog = e, e._i = [], e.init = function (i, s, a) { function g(t, e) { var o = e.split("."); 2 == o.length && (t = t[o[0]], e = o[1]), t[e] = function () { t.push([e].concat(Array.prototype.slice.call(arguments, 0))) } } (p = t.createElement("script")).type = "text/javascript", p.async = !0, p.src = s.api_host.replace(".i.posthog.com", "-assets.i.posthog.com") + "/static/array.js", (r = t.getElementsByTagName("script")[0]).parentNode.insertBefore(p, r); var u = e; for (void 0 !== a ? u = e[a] = [] : a = "posthog", u.people = u.people || [], u.toString = function (t) { var e = "posthog"; return "posthog" !== a && (e += "." + a), t || (e += " (stub)"), e }, u.people.toString = function () { return u.toString(1) + ".people (stub)" }, o = "init push capture register register_once register_for_session unregister unregister_for_session getFeatureFlag getFeatureFlagPayload isFeatureEnabled reloadFeatureFlags updateEarlyAccessFeatureEnrollment getEarlyAccessFeatures on onFeatureFlags onSessionId getSurveys getActiveMatchingSurveys renderSurvey canRenderSurvey getNextSurveyStep identify setPersonProperties group resetGroups setPersonPropertiesForFlags resetPersonPropertiesForFlags setGroupPropertiesForFlags resetGroupPropertiesForFlags reset get_distinct_id getGroups get_session_id get_session_replay_url alias set_config startSessionRecording stopSessionRecording sessionRecordingStarted loadToolbar get_property getSessionProperty createPersonProfile opt_in_capturing opt_out_capturing has_opted_in_capturing has_opted_out_capturing clear_opt_in_out_capturing debug".split(" "), n = 0; n < o.length; n++)g(u, o[n]); e._i.push([i, s, a]) }, e.__SV = 1) }(document, window.posthog || []);
+posthog.init('phc_jfYo5xEofW7eJtiU8rLt2Z8jw1E2eW27BxwTJzwRufH', {
+    api_host: 'https://us.i.posthog.com', person_profiles: 'identified_only' // or 'always' to create profiles for anonymous users as well
+})

agentic_security/static/telemetry_disabled.js

@@ -0,0 +1 @@
+console.log("Telemetry is disabled");

agentic_security/test_lib.py

@@ -14,7 +14,6 @@ Content-Type: application/json
 
 
 class TestAS:
-
     # Handles an empty dataset list.
     def test_class(self):
         llmSpec = SAMPLE_SPEC

agentic_security/test_spec.py

@@ -2,7 +2,6 @@ from agentic_security.http_spec import LLMSpec, parse_http_spec
 
 
 class TestParseHttpSpec:
-
     # Should correctly parse a simple HTTP spec with headers and body
     def test_parse_simple_http_spec(self):
         http_spec = (

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.2.4"
+version = "0.3.3"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
@@ -25,7 +25,7 @@ packages = [{ include = "agentic_security", from = "." }]
 agentic_security = "agentic_security.__main__:entrypoint"
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = "^3.11"
 fastapi = "^0.115.2"
 uvicorn = "^0.32.0"
 fire = "0.7.0"
@@ -39,7 +39,9 @@ colorama = "^0.4.4"
 matplotlib = "^3.9.2"
 pydantic = "2.9.2"
 scikit-optimize = "^0.10.2"
-scikit-learn = "1.5.1"
+scikit-learn = "1.5.2"
+numpy = ">=1.24.3,<3.0.0"
+jinja2 = "^3.1.4"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.10.0"