Add real-time scan progress tracking
- Add ScanProgress model and progress fields to Scan model - Implement progress callback in scanner to report phase completion - Update scan_job to write per-IP results to database during execution - Add /api/scans/<id>/progress endpoint for progress polling - Add progress section to scan detail page with live updates - Progress table shows current phase, completion bar, and per-IP results - Poll every 3 seconds during active scans - Sort IPs numerically for proper ordering - Add database migration for new tables/columns
This commit is contained in:
@@ -5,6 +5,7 @@ This module handles the execution of scans in background threads,
|
||||
updating database status and handling errors.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
@@ -14,13 +15,132 @@ from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from src.scanner import SneakyScanner
|
||||
from web.models import Scan
|
||||
from web.models import Scan, ScanProgress
|
||||
from web.services.scan_service import ScanService
|
||||
from web.services.alert_service import AlertService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_progress_callback(scan_id: int, session):
|
||||
"""
|
||||
Create a progress callback function for updating scan progress in database.
|
||||
|
||||
Args:
|
||||
scan_id: ID of the scan record
|
||||
session: Database session
|
||||
|
||||
Returns:
|
||||
Callback function that accepts (phase, ip, data)
|
||||
"""
|
||||
ip_to_site = {}
|
||||
|
||||
def progress_callback(phase: str, ip: str, data: dict):
|
||||
"""Update scan progress in database."""
|
||||
nonlocal ip_to_site
|
||||
|
||||
try:
|
||||
# Get scan record
|
||||
scan = session.query(Scan).filter_by(id=scan_id).first()
|
||||
if not scan:
|
||||
return
|
||||
|
||||
# Handle initialization phase
|
||||
if phase == 'init':
|
||||
scan.total_ips = data.get('total_ips', 0)
|
||||
scan.completed_ips = 0
|
||||
scan.current_phase = 'ping'
|
||||
ip_to_site = data.get('ip_to_site', {})
|
||||
|
||||
# Create progress entries for all IPs
|
||||
for ip_addr, site_name in ip_to_site.items():
|
||||
progress = ScanProgress(
|
||||
scan_id=scan_id,
|
||||
ip_address=ip_addr,
|
||||
site_name=site_name,
|
||||
phase='pending',
|
||||
status='pending'
|
||||
)
|
||||
session.add(progress)
|
||||
|
||||
session.commit()
|
||||
return
|
||||
|
||||
# Update current phase
|
||||
if data.get('status') == 'starting':
|
||||
scan.current_phase = phase
|
||||
scan.completed_ips = 0
|
||||
session.commit()
|
||||
return
|
||||
|
||||
# Handle phase completion with results
|
||||
if data.get('status') == 'completed':
|
||||
results = data.get('results', {})
|
||||
|
||||
if phase == 'ping':
|
||||
# Update progress entries with ping results
|
||||
for ip_addr, ping_result in results.items():
|
||||
progress = session.query(ScanProgress).filter_by(
|
||||
scan_id=scan_id, ip_address=ip_addr
|
||||
).first()
|
||||
if progress:
|
||||
progress.ping_result = ping_result
|
||||
progress.phase = 'ping'
|
||||
progress.status = 'completed'
|
||||
|
||||
scan.completed_ips = len(results)
|
||||
|
||||
elif phase == 'tcp_scan':
|
||||
# Update progress entries with TCP/UDP port results
|
||||
for ip_addr, port_data in results.items():
|
||||
progress = session.query(ScanProgress).filter_by(
|
||||
scan_id=scan_id, ip_address=ip_addr
|
||||
).first()
|
||||
if progress:
|
||||
progress.tcp_ports = json.dumps(port_data.get('tcp_ports', []))
|
||||
progress.udp_ports = json.dumps(port_data.get('udp_ports', []))
|
||||
progress.phase = 'tcp_scan'
|
||||
progress.status = 'completed'
|
||||
|
||||
scan.completed_ips = len(results)
|
||||
|
||||
elif phase == 'service_detection':
|
||||
# Update progress entries with service detection results
|
||||
for ip_addr, services in results.items():
|
||||
progress = session.query(ScanProgress).filter_by(
|
||||
scan_id=scan_id, ip_address=ip_addr
|
||||
).first()
|
||||
if progress:
|
||||
# Simplify service data for storage
|
||||
service_list = []
|
||||
for svc in services:
|
||||
service_list.append({
|
||||
'port': svc.get('port'),
|
||||
'service': svc.get('service', 'unknown'),
|
||||
'product': svc.get('product', ''),
|
||||
'version': svc.get('version', '')
|
||||
})
|
||||
progress.services = json.dumps(service_list)
|
||||
progress.phase = 'service_detection'
|
||||
progress.status = 'completed'
|
||||
|
||||
scan.completed_ips = len(results)
|
||||
|
||||
elif phase == 'http_analysis':
|
||||
# Mark HTTP analysis as complete
|
||||
scan.current_phase = 'completed'
|
||||
scan.completed_ips = scan.total_ips
|
||||
|
||||
session.commit()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Progress callback error for scan {scan_id}: {str(e)}")
|
||||
# Don't re-raise - we don't want to break the scan
|
||||
session.rollback()
|
||||
|
||||
return progress_callback
|
||||
|
||||
|
||||
def execute_scan(scan_id: int, config_id: int, db_url: str = None):
|
||||
"""
|
||||
Execute a scan in the background.
|
||||
@@ -66,10 +186,13 @@ def execute_scan(scan_id: int, config_id: int, db_url: str = None):
|
||||
# Initialize scanner with database config
|
||||
scanner = SneakyScanner(config_id=config_id)
|
||||
|
||||
# Execute scan
|
||||
# Create progress callback
|
||||
progress_callback = create_progress_callback(scan_id, session)
|
||||
|
||||
# Execute scan with progress tracking
|
||||
logger.info(f"Scan {scan_id}: Running scanner...")
|
||||
start_time = datetime.utcnow()
|
||||
report, timestamp = scanner.scan()
|
||||
report, timestamp = scanner.scan(progress_callback=progress_callback)
|
||||
end_time = datetime.utcnow()
|
||||
|
||||
scan_duration = (end_time - start_time).total_seconds()
|
||||
|
||||
Reference in New Issue
Block a user