mstoeck3/dnscope
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix race condition

Browse Source
This commit is contained in:
overcuriousity 2025-09-14 01:40:17 +02:00
parent 2185177a84
commit b26002eff9
1 changed files with 84 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

134
core/scanner.py
View File

@ -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