36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import asyncio
|
|
|
|
import websockets.exceptions
|
|
from fastapi import WebSocket, WebSocketDisconnect, APIRouter
|
|
|
|
from tools.web_monitor.scan.func import do_websocket_scan
|
|
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_scan(websocket: WebSocket):
|
|
await websocket.accept()
|
|
cur_task = None
|
|
try:
|
|
while True:
|
|
ws_data = await websocket.receive_text()
|
|
if "-Cancel-" in ws_data:
|
|
if cur_task:
|
|
cur_task.cancel()
|
|
print("Cancelling scan...")
|
|
try:
|
|
await cur_task
|
|
except asyncio.CancelledError:
|
|
cur_task = None
|
|
await websocket.send_text("Cancelled")
|
|
else:
|
|
cur_task = asyncio.create_task(do_websocket_scan(websocket, ws_data))
|
|
if cur_task and cur_task.done():
|
|
cur_task = None
|
|
except WebSocketDisconnect:
|
|
print("Websocket disconnected.")
|
|
except websockets.exceptions.ConnectionClosedOK:
|
|
pass
|