mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 06:09:55 +02:00
feat(add stop event):
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 -->
|
||||
|
||||
@@ -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
@@ -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>"]
|
||||
|
||||
Reference in New Issue
Block a user