fix race condition
This commit is contained in:
overcuriousity
parent
2185177a84
commit
b26002eff9
132
core/scanner.py
132
core/scanner.py
@ -50,6 +50,7 @@ class Scanner:
|
||||
self.stop_event = threading.Event()
|
||||
self.scan_thread = None
|
||||
self.session_id = None # Will be set by session manager
|
||||
self.current_scan_id = None # NEW: Track current scan ID
|
||||
|
||||
# Scanning progress tracking
|
||||
self.total_indicators_found = 0
|
||||
@ -193,24 +194,43 @@ class Scanner:
|
||||
print(f"=== STARTING SCAN IN SCANNER {id(self)} ===")
|
||||
print(f"Session ID: {self.session_id}")
|
||||
print(f"Initial scanner status: {self.status}")
|
||||
print(f"Clear graph: {clear_graph}")
|
||||
|
||||
# Clean up previous scan thread if needed
|
||||
if self.scan_thread and self.scan_thread.is_alive():
|
||||
print("A previous scan thread is still alive. Sending termination signal and waiting...")
|
||||
self.stop_scan()
|
||||
self.scan_thread.join(10.0)
|
||||
# Generate scan ID based on clear_graph behavior
|
||||
import uuid
|
||||
|
||||
if self.scan_thread.is_alive():
|
||||
print("ERROR: The previous scan thread is unresponsive and could not be stopped.")
|
||||
self.status = ScanStatus.FAILED
|
||||
self._update_session_state()
|
||||
return False
|
||||
print("Previous scan thread terminated successfully.")
|
||||
if clear_graph:
|
||||
# NEW SCAN: Generate new ID and terminate existing scan
|
||||
print("NEW SCAN: Generating new scan ID and terminating existing scan")
|
||||
self.current_scan_id = str(uuid.uuid4())[:8]
|
||||
|
||||
# Reset state for new scan
|
||||
self.status = ScanStatus.IDLE
|
||||
self._update_session_state() # Update GUI immediately
|
||||
print("Scanner state is now clean for a new scan.")
|
||||
# Aggressive cleanup of previous scan
|
||||
if self.scan_thread and self.scan_thread.is_alive():
|
||||
print("Terminating previous scan thread...")
|
||||
self._set_stop_signal()
|
||||
|
||||
if self.executor:
|
||||
self.executor.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
self.scan_thread.join(timeout=8.0)
|
||||
if self.scan_thread.is_alive():
|
||||
print("WARNING: Previous scan thread did not terminate cleanly")
|
||||
|
||||
else:
|
||||
# ADD TO GRAPH: Keep existing scan ID if scan is running, or generate new one
|
||||
if self.status == ScanStatus.RUNNING and self.current_scan_id:
|
||||
print(f"ADD TO GRAPH: Keeping existing scan ID {self.current_scan_id}")
|
||||
# Don't terminate existing scan - we're adding to it
|
||||
else:
|
||||
print("ADD TO GRAPH: No active scan, generating new scan ID")
|
||||
self.current_scan_id = str(uuid.uuid4())[:8]
|
||||
|
||||
print(f"Using scan ID: {self.current_scan_id}")
|
||||
|
||||
# Reset state for new scan (but preserve graph if clear_graph=False)
|
||||
if clear_graph or self.status != ScanStatus.RUNNING:
|
||||
self.status = ScanStatus.IDLE
|
||||
self._update_session_state()
|
||||
|
||||
try:
|
||||
if not hasattr(self, 'providers') or not self.providers:
|
||||
@ -221,32 +241,33 @@ class Scanner:
|
||||
|
||||
if clear_graph:
|
||||
self.graph.clear()
|
||||
|
||||
self.current_target = target_domain.lower().strip()
|
||||
self.max_depth = max_depth
|
||||
self.current_depth = 0
|
||||
|
||||
# Clear both local and Redis stop signals
|
||||
self.stop_event.clear()
|
||||
if self.session_id:
|
||||
from core.session_manager import session_manager
|
||||
session_manager.clear_stop_signal(self.session_id)
|
||||
# Clear stop signals only if starting new scan
|
||||
if clear_graph or self.status != ScanStatus.RUNNING:
|
||||
self.stop_event.clear()
|
||||
if self.session_id:
|
||||
from core.session_manager import session_manager
|
||||
session_manager.clear_stop_signal(self.session_id)
|
||||
|
||||
self.total_indicators_found = 0
|
||||
self.indicators_processed = 0
|
||||
self.current_indicator = self.current_target
|
||||
|
||||
# Update GUI with scan preparation
|
||||
self._update_session_state()
|
||||
|
||||
# Start new forensic session
|
||||
print(f"Starting new forensic session for scanner {id(self)}...")
|
||||
self.logger = new_session()
|
||||
# Initialize forensic session only for new scans
|
||||
if clear_graph:
|
||||
self.logger = new_session()
|
||||
|
||||
# Start scan in separate thread
|
||||
print(f"Starting scan thread for scanner {id(self)}...")
|
||||
# Start scan thread (original behavior allows concurrent threads for "Add to Graph")
|
||||
print(f"Starting scan thread with scan ID {self.current_scan_id}...")
|
||||
self.scan_thread = threading.Thread(
|
||||
target=self._execute_scan,
|
||||
args=(self.current_target, max_depth),
|
||||
args=(self.current_target, max_depth, self.current_scan_id),
|
||||
daemon=True
|
||||
)
|
||||
self.scan_thread.start()
|
||||
@ -258,12 +279,13 @@ class Scanner:
|
||||
print(f"ERROR: Exception in start_scan for scanner {id(self)}: {e}")
|
||||
traceback.print_exc()
|
||||
self.status = ScanStatus.FAILED
|
||||
self._update_session_state() # Update failed status immediately
|
||||
self._update_session_state()
|
||||
return False
|
||||
|
||||
def _execute_scan(self, target_domain: str, max_depth: int) -> None:
|
||||
def _execute_scan(self, target_domain: str, max_depth: int, scan_id: str) -> None:
|
||||
"""Execute the reconnaissance scan using a task queue-based approach."""
|
||||
print(f"_execute_scan started for {target_domain} with depth {max_depth}")
|
||||
print(f"_execute_scan started for {target_domain} with depth {max_depth}, scan ID {scan_id}")
|
||||
|
||||
self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
|
||||
processed_targets = set()
|
||||
|
||||
@ -279,16 +301,18 @@ class Scanner:
|
||||
self._initialize_provider_states(target_domain)
|
||||
|
||||
while task_queue:
|
||||
# Abort if scan ID changed (new scan started)
|
||||
if self.current_scan_id != scan_id:
|
||||
print(f"Scan aborted - ID mismatch (current: {self.current_scan_id}, expected: {scan_id})")
|
||||
break
|
||||
|
||||
if self._is_stop_requested():
|
||||
print("Stop requested, terminating scan.")
|
||||
break
|
||||
|
||||
target, depth, is_large_entity_member = task_queue.popleft()
|
||||
|
||||
if target in processed_targets:
|
||||
continue
|
||||
|
||||
if depth > max_depth:
|
||||
if target in processed_targets or depth > max_depth:
|
||||
continue
|
||||
|
||||
self.current_depth = depth
|
||||
@ -298,14 +322,15 @@ class Scanner:
|
||||
new_targets, large_entity_members = self._query_providers_for_target(target, depth, is_large_entity_member)
|
||||
processed_targets.add(target)
|
||||
|
||||
for new_target in new_targets:
|
||||
if new_target not in processed_targets:
|
||||
task_queue.append((new_target, depth + 1, False))
|
||||
|
||||
for member in large_entity_members:
|
||||
if member not in processed_targets:
|
||||
task_queue.append((member, depth, True))
|
||||
# Only add new targets if scan ID still matches (prevents stale updates)
|
||||
if self.current_scan_id == scan_id:
|
||||
for new_target in new_targets:
|
||||
if new_target not in processed_targets:
|
||||
task_queue.append((new_target, depth + 1, False))
|
||||
|
||||
for member in large_entity_members:
|
||||
if member not in processed_targets:
|
||||
task_queue.append((member, depth, True))
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Scan execution failed with error: {e}")
|
||||
@ -313,13 +338,18 @@ class Scanner:
|
||||
self.status = ScanStatus.FAILED
|
||||
self.logger.logger.error(f"Scan failed: {e}")
|
||||
finally:
|
||||
if self._is_stop_requested():
|
||||
self.status = ScanStatus.STOPPED
|
||||
else:
|
||||
self.status = ScanStatus.COMPLETED
|
||||
# Only update final status if scan ID still matches (prevents stale status updates)
|
||||
if self.current_scan_id == scan_id:
|
||||
if self._is_stop_requested():
|
||||
self.status = ScanStatus.STOPPED
|
||||
else:
|
||||
self.status = ScanStatus.COMPLETED
|
||||
|
||||
self._update_session_state()
|
||||
self.logger.log_scan_complete()
|
||||
else:
|
||||
print(f"Scan completed but ID mismatch - not updating final status")
|
||||
|
||||
self._update_session_state()
|
||||
self.logger.log_scan_complete()
|
||||
if self.executor:
|
||||
self.executor.shutdown(wait=False, cancel_futures=True)
|
||||
stats = self.graph.get_statistics()
|
||||
@ -621,7 +651,6 @@ class Scanner:
|
||||
if target not in attributes[record_type_name]:
|
||||
attributes[record_type_name].append(target)
|
||||
|
||||
|
||||
def _log_target_processing_error(self, target: str, error: str) -> None:
|
||||
"""Log target processing errors for forensic trail."""
|
||||
self.logger.logger.error(f"Target processing failed for {target}: {error}")
|
||||
@ -641,7 +670,12 @@ class Scanner:
|
||||
print("=== INITIATING IMMEDIATE SCAN TERMINATION ===")
|
||||
self.logger.logger.info("Scan termination requested by user")
|
||||
|
||||
# Set both local and Redis stop signals
|
||||
# Invalidate current scan ID to prevent stale updates
|
||||
old_scan_id = self.current_scan_id
|
||||
self.current_scan_id = None
|
||||
print(f"Invalidated scan ID {old_scan_id}")
|
||||
|
||||
# Set stop signals
|
||||
self._set_stop_signal()
|
||||
self.status = ScanStatus.STOPPED
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user