mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-08-09 03:06:05 +02:00
feat(add stop event):
This commit is contained in:
@@ -33,6 +33,9 @@ app.add_middleware(
|
|||||||
)
|
)
|
||||||
|
|
||||||
tools_inbox = Queue()
|
tools_inbox = Queue()
|
||||||
|
# Global stop event for cancelling scans
|
||||||
|
stop_event = Event() # Added stop_event to cancel the scan
|
||||||
|
|
||||||
FEATURE_PROXY = False
|
FEATURE_PROXY = False
|
||||||
|
|
||||||
|
|
||||||
@@ -99,6 +102,7 @@ def streaming_response_generator(scan_parameters: Scan):
|
|||||||
datasets=scan_parameters.datasets,
|
datasets=scan_parameters.datasets,
|
||||||
tools_inbox=tools_inbox,
|
tools_inbox=tools_inbox,
|
||||||
optimize=scan_parameters.optimize,
|
optimize=scan_parameters.optimize,
|
||||||
|
stop_event=stop_event, # Pass the stop_event to the generator
|
||||||
):
|
):
|
||||||
yield scan_result + "\n" # Adding a newline for separation
|
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):
|
class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
|
||||||
async def dispatch(self, request: Request, call_next):
|
async def dispatch(self, request: Request, call_next):
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
from typing import AsyncGenerator
|
from typing import AsyncGenerator
|
||||||
|
|
||||||
@@ -49,6 +50,7 @@ async def perform_scan(
|
|||||||
datasets: list[dict[str, str]] = [],
|
datasets: list[dict[str, str]] = [],
|
||||||
tools_inbox=None,
|
tools_inbox=None,
|
||||||
optimize=False,
|
optimize=False,
|
||||||
|
stop_event: asyncio.Event = None,
|
||||||
) -> AsyncGenerator[str, None]:
|
) -> AsyncGenerator[str, None]:
|
||||||
if IS_VERCEL:
|
if IS_VERCEL:
|
||||||
yield ScanResult.status_msg(
|
yield ScanResult.status_msg(
|
||||||
@@ -81,6 +83,12 @@ async def perform_scan(
|
|||||||
)
|
)
|
||||||
should_stop_early = False
|
should_stop_early = False
|
||||||
async for prompt in prompt_iter(module.prompts):
|
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
|
processed_prompts += 1
|
||||||
progress = 100 * processed_prompts / total_prompts if total_prompts else 0
|
progress = 100 * processed_prompts / total_prompts if total_prompts else 0
|
||||||
|
|
||||||
|
|||||||
@@ -335,6 +335,7 @@
|
|||||||
</button>
|
</button>
|
||||||
<button
|
<button
|
||||||
@click="startScan"
|
@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">
|
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"
|
<svg xmlns="http://www.w3.org/2000/svg" width="24" height="24"
|
||||||
viewBox="0 0 24 24" fill="none" stroke="currentColor"
|
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>
|
class="mr-2"><polygon points="5 3 19 12 5 21 5 3"></polygon></svg>
|
||||||
Run Scan
|
Run Scan
|
||||||
</button>
|
</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>
|
</section>
|
||||||
|
|
||||||
<!-- Progress Bar -->
|
<!-- Progress Bar -->
|
||||||
|
|||||||
@@ -341,6 +341,15 @@ var app = new Vue({
|
|||||||
}
|
}
|
||||||
this.budget = value;
|
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 () {
|
startScan: async function () {
|
||||||
this.showLLMSpec = false;
|
this.showLLMSpec = false;
|
||||||
let payload = {
|
let payload = {
|
||||||
@@ -358,6 +367,7 @@ var app = new Vue({
|
|||||||
});
|
});
|
||||||
this.okMsg = 'Scan started';
|
this.okMsg = 'Scan started';
|
||||||
this.mainTable = [];
|
this.mainTable = [];
|
||||||
|
this.scanRunning = true;
|
||||||
const reader = response.body.getReader();
|
const reader = response.body.getReader();
|
||||||
let receivedLength = 0; // received that many bytes at the moment
|
let receivedLength = 0; // received that many bytes at the moment
|
||||||
let chunks = []; // array of received binary chunks (comprises the body)
|
let chunks = []; // array of received binary chunks (comprises the body)
|
||||||
|
|||||||
+1
-1
@@ -1,6 +1,6 @@
|
|||||||
[tool.poetry]
|
[tool.poetry]
|
||||||
name = "agentic_security"
|
name = "agentic_security"
|
||||||
version = "0.2.1"
|
version = "0.2.2"
|
||||||
description = "Agentic LLM vulnerability scanner"
|
description = "Agentic LLM vulnerability scanner"
|
||||||
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
||||||
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
|
||||||
|
|||||||
Reference in New Issue
Block a user