From f128864db11cd855e8f937355af07ec8221df925 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 19 Oct 2024 15:31:29 +0300
Subject: [PATCH] feat(add stop event):

---
 agentic_security/app.py                | 10 ++++++++++
 agentic_security/probe_actor/fuzzer.py |  8 ++++++++
 agentic_security/static/index.html     | 13 +++++++++++++
 agentic_security/static/main.js        | 10 ++++++++++
 pyproject.toml                         |  2 +-
 5 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/agentic_security/app.py b/agentic_security/app.py
index cbb8c1d..1cf0ef3 100644
--- a/agentic_security/app.py
+++ b/agentic_security/app.py
@@ -33,6 +33,9 @@ app.add_middleware(
 )
 
 tools_inbox = Queue()
+# Global stop event for cancelling scans
+stop_event = Event()  # Added stop_event to cancel the scan
+
 FEATURE_PROXY = False
 
 
@@ -99,6 +102,7 @@ def streaming_response_generator(scan_parameters: Scan):
             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
 
@@ -238,6 +242,12 @@ config.dictConfig(
 )
 
 
+@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:
diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py
index b4fddbb..88b8235 100644
--- a/agentic_security/probe_actor/fuzzer.py
+++ b/agentic_security/probe_actor/fuzzer.py
@@ -1,3 +1,4 @@
+import asyncio
 import os
 from typing import AsyncGenerator
 
@@ -49,6 +50,7 @@ async def perform_scan(
     datasets: list[dict[str, str]] = [],
     tools_inbox=None,
     optimize=False,
+    stop_event: asyncio.Event = None,
 ) -> AsyncGenerator[str, None]:
     if IS_VERCEL:
         yield ScanResult.status_msg(
@@ -81,6 +83,12 @@ async def perform_scan(
         )
         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
+
             processed_prompts += 1
             progress = 100 * processed_prompts / total_prompts if total_prompts else 0
 
diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html
index fc8d325..04c4627 100644
--- a/agentic_security/static/index.html
+++ b/agentic_security/static/index.html
@@ -335,6 +335,7 @@
           </button>
           <button
             @click="startScan"
+            v-if="!scanRunning"
             class="bg-dark-accent-green text-dark-bg rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors flex items-center">
             <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24"
               viewBox="0 0 24 24" fill="none" stroke="currentColor"
@@ -342,6 +343,18 @@
               class="mr-2"><polygon points="5 3 19 12 5 21 5 3"></polygon></svg>
             Run Scan
           </button>
+          <button
+            @click="stopScan"
+            v-if="scanRunning"
+            class="bg-dark-accent-red text-dark-bg rounded-lg px-6 py-3 font-medium hover:bg-opacity-80 transition-colors flex items-center">
+            <!-- Stop Icon -->
+            <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24"
+              viewBox="0 0 24 24" fill="none" stroke="currentColor"
+              stroke-width="2" stroke-linecap="round" stroke-linejoin="round"
+              class="mr-2"><rect x="6" y="6" width="12"
+                height="12"></rect></svg>
+            Stop Scan
+          </button>
         </section>
 
         <!-- Progress Bar -->
diff --git a/agentic_security/static/main.js b/agentic_security/static/main.js
index 3b02fd0..abbea52 100644
--- a/agentic_security/static/main.js
+++ b/agentic_security/static/main.js
@@ -341,6 +341,15 @@ var app = new Vue({
             }
             this.budget = value;
         },
+        stopScan: async function () {
+            this.scanRunning = false;
+            const response = await fetch(`${URL}/stop`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+            });
+        },
         startScan: async function () {
             this.showLLMSpec = false;
             let payload = {
@@ -358,6 +367,7 @@ var app = new Vue({
             });
             this.okMsg = 'Scan started';
             this.mainTable = [];
+            this.scanRunning = true;
             const reader = response.body.getReader();
             let receivedLength = 0; // received that many bytes at the moment
             let chunks = []; // array of received binary chunks (comprises the body)
diff --git a/pyproject.toml b/pyproject.toml
index 6e5b17e..81dc2e7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.2.1"
+version = "0.2.2"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]