feat(add stop event):

This commit is contained in:
Alexander Myasoedov
2024-10-19 15:31:29 +03:00
parent e4c0436636
commit f128864db1
5 changed files with 42 additions and 1 deletions
+10
View File
@@ -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:
+8
View File
@@ -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
+13
View File
@@ -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 -->
+10
View File
@@ -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)
+1 -1
View File
@@ -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>"]