core/scanner.py

@@ -664,6 +664,9 @@ class Scanner:
         print(f"Enqueueing correlation tasks for {len(all_nodes)} nodes")
         
         for node_id in all_nodes:
+            if self._is_stop_requested():
+                break
+                
             # Determine appropriate depth for correlation (use 0 for simplicity)
             correlation_depth = 0
             task_tuple = ('correlation', node_id, correlation_depth)
@@ -681,7 +684,7 @@ class Scanner:
         consecutive_empty_iterations = 0
         max_empty_iterations = 20  # Shorter timeout for correlation phase
         
-        while correlation_tasks:
+        while not self._is_stop_requested() and correlation_tasks:
             queue_empty = self.task_queue.empty()
             with self.processing_lock:
                 no_active_processing = len(self.currently_processing) == 0
@@ -717,6 +720,8 @@ class Scanner:
                 continue
 
             with self.processing_lock:
+                if self._is_stop_requested():
+                    break
                 processing_key = (provider_name, target_item)
                 if processing_key in self.currently_processing:
                     self.tasks_skipped += 1
@@ -728,6 +733,9 @@ class Scanner:
                 self.current_indicator = target_item
                 self._update_session_state()
                 
+                if self._is_stop_requested():
+                    break
+                
                 # Process correlation task
                 new_targets, _, success = self._process_provider_task(correlation_provider, target_item, depth)
                 
@@ -754,7 +762,7 @@ class Scanner:
         Manages the entire process for a given target and provider.
         This version is generalized to handle all relationships dynamically.
         """
-        if self._is_stop_requested() and not isinstance(provider, CorrelationProvider):
+        if self._is_stop_requested():
             return set(), set(), False
     
         is_ip = _is_valid_ip(target)
@@ -771,8 +779,7 @@ class Scanner:
     
             if provider_result is None:
                 provider_successful = False
-            # Allow correlation provider to process results even if scan is stopped
-            elif not self._is_stop_requested() or isinstance(provider, CorrelationProvider):
+            elif not self._is_stop_requested():
                 # Pass all relationships to be processed
                 discovered, is_large_entity = self._process_provider_result_unified(
                     target, provider, provider_result, depth
@@ -792,7 +799,7 @@ class Scanner:
         provider_name = provider.get_name()
         start_time = datetime.now(timezone.utc)
         
-        if self._is_stop_requested() and not isinstance(provider, CorrelationProvider):
+        if self._is_stop_requested():
             return None
         
         try:
@@ -801,7 +808,7 @@ class Scanner:
             else:
                 result = provider.query_domain(target)
             
-            if self._is_stop_requested() and not isinstance(provider, CorrelationProvider):
+            if self._is_stop_requested():
                 return None
             
             relationship_count = result.get_relationship_count() if result else 0
@@ -895,8 +902,7 @@ class Scanner:
         large_entity_id = ""
         large_entity_members = set()
 
-        # Stop processing for non-correlation providers if requested
-        if self._is_stop_requested() and not isinstance(provider, CorrelationProvider):
+        if self._is_stop_requested():
             return discovered_targets, False
 
         eligible_rel_count = sum(
@@ -910,8 +916,7 @@ class Scanner:
             )
 
         for i, relationship in enumerate(provider_result.relationships):
-            # Stop processing for non-correlation providers if requested
-            if i % 5 == 0 and self._is_stop_requested() and not isinstance(provider, CorrelationProvider):
+            if i % 5 == 0 and self._is_stop_requested():
                 break
 
             source_node_id = relationship.source_node