From b26002eff9866fc0c7bebe8bffd0e2625d80ee6b Mon Sep 17 00:00:00 2001 From: overcuriousity Date: Sun, 14 Sep 2025 01:40:17 +0200 Subject: [PATCH] fix race condition --- core/scanner.py | 134 ++++++++++++++++++++++++++++++------------------ 1 file changed, 84 insertions(+), 50 deletions(-) diff --git a/core/scanner.py b/core/scanner.py index a366f72..23406e3 100644 --- a/core/scanner.py +++ b/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 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] + + # 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") - 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.") + 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] - # 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.") + 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() + + # Initialize forensic session only for new scans + if clear_graph: + self.logger = new_session() - # Start new forensic session - print(f"Starting new forensic session for scanner {id(self)}...") - 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 + # 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: - self.status = ScanStatus.COMPLETED + 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