This commit is contained in:
overcuriousity 2025-09-13 00:39:00 +02:00
parent c0b820c96c
commit d36fb7d814

View File

@ -5,7 +5,7 @@ import traceback
import time import time
from typing import List, Set, Dict, Any, Tuple from typing import List, Set, Dict, Any, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed, CancelledError, Future from concurrent.futures import ThreadPoolExecutor, as_completed, CancelledError, Future
from collections import defaultdict from collections import defaultdict, deque
from datetime import datetime, timezone from datetime import datetime, timezone
from core.graph_manager import GraphManager, NodeType, RelationshipType from core.graph_manager import GraphManager, NodeType, RelationshipType
@ -270,14 +270,16 @@ class Scanner:
return False return False
def _execute_scan(self, target_domain: str, max_depth: int) -> None: def _execute_scan(self, target_domain: str, max_depth: int) -> None:
"""Execute the reconnaissance scan with frequent state updates for GUI.""" """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}")
self.executor = ThreadPoolExecutor(max_workers=self.max_workers) self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
processed_targets = set() processed_targets = set()
# Initialize the task queue with the starting target and its depth
task_queue = deque([(target_domain, 0)])
try: try:
self.status = ScanStatus.RUNNING self.status = ScanStatus.RUNNING
# Immediate status update for GUI
self._update_session_state() self._update_session_state()
enabled_providers = [provider.get_name() for provider in self.providers] enabled_providers = [provider.get_name() for provider in self.providers]
@ -285,60 +287,44 @@ class Scanner:
self.graph.add_node(target_domain, NodeType.DOMAIN) self.graph.add_node(target_domain, NodeType.DOMAIN)
self._initialize_provider_states(target_domain) self._initialize_provider_states(target_domain)
current_level_targets = {target_domain} while task_queue:
all_discovered_targets = {target_domain}
for depth in range(max_depth + 1):
if self._is_stop_requested(): if self._is_stop_requested():
print(f"Stop requested at depth {depth}") print("Stop requested, terminating scan.")
break break
target, depth = task_queue.popleft()
if target in processed_targets:
continue
if depth > max_depth:
continue
self.current_depth = depth self.current_depth = depth
# Update current depth immediately for GUI self.current_indicator = target
self._update_session_state() self._update_session_state()
targets_to_process = current_level_targets - processed_targets # Process the current target
if not targets_to_process: new_targets = self._query_providers_for_target(target, depth)
print("No new targets to process at this level.") processed_targets.add(target)
break
# Add new, unprocessed targets to the queue
print(f"Processing depth level {depth} with {len(targets_to_process)} new targets") for new_target in new_targets:
self.total_indicators_found += len(targets_to_process) if new_target not in processed_targets:
# Update total indicators for GUI task_queue.append((new_target, depth + 1))
self._update_session_state()
target_results = self._process_targets_sequential_with_stop_checks(
targets_to_process, processed_targets, all_discovered_targets, depth
)
processed_targets.update(targets_to_process)
next_level_targets = set()
for _target, new_targets in target_results:
all_discovered_targets.update(new_targets)
next_level_targets.update(new_targets)
current_level_targets = next_level_targets - processed_targets
# Additional stop check at end of each depth level
if self._is_stop_requested():
print(f"Stop requested after processing depth {depth}")
break
except Exception as e: except Exception as e:
print(f"ERROR: Scan execution failed with error: {e}") print(f"ERROR: Scan execution failed with error: {e}")
traceback.print_exc() traceback.print_exc()
self.status = ScanStatus.FAILED self.status = ScanStatus.FAILED
self._update_session_state() # Update failed status immediately
self.logger.logger.error(f"Scan failed: {e}") self.logger.logger.error(f"Scan failed: {e}")
finally: finally:
if self._is_stop_requested(): if self._is_stop_requested():
self.status = ScanStatus.STOPPED self.status = ScanStatus.STOPPED
else: else:
self.status = ScanStatus.COMPLETED self.status = ScanStatus.COMPLETED
# Final status update for GUI
self._update_session_state() self._update_session_state()
self.logger.log_scan_complete() self.logger.log_scan_complete()
if self.executor: if self.executor:
self.executor.shutdown(wait=False, cancel_futures=True) self.executor.shutdown(wait=False, cancel_futures=True)
@ -348,6 +334,56 @@ class Scanner:
print(f" - Total edges: {stats['basic_metrics']['total_edges']}") print(f" - Total edges: {stats['basic_metrics']['total_edges']}")
print(f" - Targets processed: {len(processed_targets)}") print(f" - Targets processed: {len(processed_targets)}")
def _query_providers_for_target(self, target: str, depth: int) -> Set[str]:
"""Helper method to query providers for a single target."""
is_ip = _is_valid_ip(target)
target_type = NodeType.IP if is_ip else NodeType.DOMAIN
print(f"Querying providers for {target_type.value}: {target} at depth {depth}")
# Early stop check
if self._is_stop_requested():
print(f"Stop requested before querying providers for {target}")
return set()
# Initialize node and provider states
self.graph.add_node(target, target_type)
self._initialize_provider_states(target)
new_targets = set()
target_metadata = defaultdict(lambda: defaultdict(list))
# Determine eligible providers for this target
eligible_providers = self._get_eligible_providers(target, is_ip)
if not eligible_providers:
self._log_no_eligible_providers(target, is_ip)
return new_targets
# Query each eligible provider sequentially with stop checks
for provider in eligible_providers:
if self._is_stop_requested():
print(f"Stop requested while querying providers for {target}")
break
try:
provider_results = self._query_single_provider_forensic(provider, target, is_ip, depth)
if provider_results and not self._is_stop_requested():
discovered_targets = self._process_provider_results_forensic(
target, provider, provider_results, target_metadata, depth
)
new_targets.update(discovered_targets)
except Exception as e:
self._log_provider_error(target, provider.get_name(), str(e))
# Update node metadata
for node_id, metadata_dict in target_metadata.items():
if self.graph.graph.has_node(node_id):
node_is_ip = _is_valid_ip(node_id)
node_type_to_add = NodeType.IP if node_is_ip else NodeType.DOMAIN
self.graph.add_node(node_id, node_type_to_add, metadata=metadata_dict)
return new_targets
def _update_session_state(self) -> None: def _update_session_state(self) -> None:
""" """
Update the scanner state in Redis for GUI updates. Update the scanner state in Redis for GUI updates.