From b2a12a3a6282f839653f40c9ca2399e89b47d603 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Mon, 2 Dec 2024 20:41:08 +0200
Subject: [PATCH] feat(add MSJ part 1):

---
 agentic_security/models/schemas.py     |  1 +
 agentic_security/probe_actor/fuzzer.py | 35 +++++++++++++++++++++++---
 agentic_security/routes/scan.py        |  6 ++---
 agentic_security/static/index.html     | 33 ++++++++++++------------
 agentic_security/static/main.js        |  1 +
 5 files changed, 52 insertions(+), 24 deletions(-)

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.
               </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>
@@ -304,13 +289,29 @@
           </div>
 
           <div v-show="showModules" class="mt-4">
-            <div class="flex justify-between mb-4">
+            <!-- Many-shot jailbreaking Toggle -->
+            <div class="flex items-center justify-between mb-2 mt-10">
+              <h3 class="text-lg font-semibold">Enable Many-shot
+                jailbreaking</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 mb-2">
+              When enabled, the scan will attempt Many-shot jailbreaking
+              simulations
+            </p>
+            <div class="flex justify-between mb-4 mt-4">
               <button @click="selectAllPackages"
                 class="text-dark-accent-green hover:underline">Select
                 All</button>
               <button @click="deselectAllPackages"
                 class="text-gray-400 hover:underline">Deselect All</button>
             </div>
+
             <div class="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 gap-4">
               <div
                 v-for="(package, index) in dataConfig"
diff --git a/agentic_security/static/main.js b/agentic_security/static/main.js
index 8bdc3d6..c4645f7 100644
--- a/agentic_security/static/main.js
+++ b/agentic_security/static/main.js
@@ -419,6 +419,7 @@ var app = new Vue({
                 llmSpec: this.modelSpec,
                 datasets: this.dataConfig,
                 optimize: this.optimize,
+                enableMultiStepAttack: this.enableMultiStepAttack,
             };
             const response = await fetch(`${URL}/scan`, {
                 method: 'POST',