Add scan cancellation feature
- Replace subprocess.run() with Popen for cancellable processes - Add cancel() method to SneakyScanner with process termination - Track running scanners in registry for stop signal delivery - Handle ScanCancelledError to set scan status to 'cancelled' - Add POST /api/scans/<id>/stop endpoint - Add 'cancelled' as valid scan status - Add Stop button to scans list and detail views - Show cancelled status with warning badge in UI
This commit is contained in:
@@ -7,6 +7,7 @@ updating database status and handling errors.
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
@@ -14,13 +15,49 @@ from pathlib import Path
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from src.scanner import SneakyScanner
|
||||
from src.scanner import SneakyScanner, ScanCancelledError
|
||||
from web.models import Scan, ScanProgress
|
||||
from web.services.scan_service import ScanService
|
||||
from web.services.alert_service import AlertService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Registry for tracking running scanners (scan_id -> SneakyScanner instance)
|
||||
_running_scanners = {}
|
||||
_running_scanners_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_running_scanner(scan_id: int):
|
||||
"""Get a running scanner instance by scan ID."""
|
||||
with _running_scanners_lock:
|
||||
return _running_scanners.get(scan_id)
|
||||
|
||||
|
||||
def stop_scan(scan_id: int, db_url: str) -> bool:
|
||||
"""
|
||||
Stop a running scan.
|
||||
|
||||
Args:
|
||||
scan_id: ID of the scan to stop
|
||||
db_url: Database connection URL
|
||||
|
||||
Returns:
|
||||
True if scan was cancelled, False if not found or already stopped
|
||||
"""
|
||||
logger.info(f"Attempting to stop scan {scan_id}")
|
||||
|
||||
# Get the scanner instance
|
||||
scanner = get_running_scanner(scan_id)
|
||||
if not scanner:
|
||||
logger.warning(f"Scanner for scan {scan_id} not found in registry")
|
||||
return False
|
||||
|
||||
# Cancel the scanner
|
||||
scanner.cancel()
|
||||
logger.info(f"Cancellation signal sent to scan {scan_id}")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def create_progress_callback(scan_id: int, session):
|
||||
"""
|
||||
@@ -186,6 +223,11 @@ def execute_scan(scan_id: int, config_id: int, db_url: str = None):
|
||||
# Initialize scanner with database config
|
||||
scanner = SneakyScanner(config_id=config_id)
|
||||
|
||||
# Register scanner in the running registry
|
||||
with _running_scanners_lock:
|
||||
_running_scanners[scan_id] = scanner
|
||||
logger.debug(f"Scan {scan_id}: Registered in running scanners registry")
|
||||
|
||||
# Create progress callback
|
||||
progress_callback = create_progress_callback(scan_id, session)
|
||||
|
||||
@@ -220,6 +262,19 @@ def execute_scan(scan_id: int, config_id: int, db_url: str = None):
|
||||
|
||||
logger.info(f"Scan {scan_id}: Completed successfully")
|
||||
|
||||
except ScanCancelledError:
|
||||
# Scan was cancelled by user
|
||||
logger.info(f"Scan {scan_id}: Cancelled by user")
|
||||
|
||||
scan = session.query(Scan).filter_by(id=scan_id).first()
|
||||
if scan:
|
||||
scan.status = 'cancelled'
|
||||
scan.error_message = 'Scan cancelled by user'
|
||||
scan.completed_at = datetime.utcnow()
|
||||
if scan.started_at:
|
||||
scan.duration = (datetime.utcnow() - scan.started_at).total_seconds()
|
||||
session.commit()
|
||||
|
||||
except FileNotFoundError as e:
|
||||
# Config file not found
|
||||
error_msg = f"Configuration file not found: {str(e)}"
|
||||
@@ -249,6 +304,12 @@ def execute_scan(scan_id: int, config_id: int, db_url: str = None):
|
||||
logger.error(f"Scan {scan_id}: Failed to update error status in database: {str(db_error)}")
|
||||
|
||||
finally:
|
||||
# Unregister scanner from registry
|
||||
with _running_scanners_lock:
|
||||
if scan_id in _running_scanners:
|
||||
del _running_scanners[scan_id]
|
||||
logger.debug(f"Scan {scan_id}: Unregistered from running scanners registry")
|
||||
|
||||
# Always close the session
|
||||
session.close()
|
||||
logger.info(f"Scan {scan_id}: Background job completed, session closed")
|
||||
|
||||
Reference in New Issue
Block a user