1506 lines
69 KiB
Python
1506 lines
69 KiB
Python
# dnsrecon-reduced/core/scanner.py
|
|
|
|
import threading
|
|
import traceback
|
|
import os
|
|
import importlib
|
|
import redis
|
|
import time
|
|
import random # Imported for jitter
|
|
from typing import List, Set, Dict, Any, Tuple, Optional
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from collections import defaultdict
|
|
from queue import PriorityQueue
|
|
from datetime import datetime, timezone
|
|
|
|
from core.graph_manager import GraphManager, NodeType
|
|
from core.logger import get_forensic_logger, new_session
|
|
from core.provider_result import ProviderResult
|
|
from utils.helpers import _is_valid_ip, _is_valid_domain
|
|
from utils.export_manager import export_manager
|
|
from providers.base_provider import BaseProvider
|
|
from providers.correlation_provider import CorrelationProvider
|
|
from core.rate_limiter import GlobalRateLimiter
|
|
|
|
class ScanStatus:
|
|
"""Enumeration of scan statuses."""
|
|
IDLE = "idle"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
STOPPED = "stopped"
|
|
|
|
|
|
class Scanner:
|
|
"""
|
|
Main scanning orchestrator for DNSRecon passive reconnaissance.
|
|
UNIFIED: Combines comprehensive features with improved display formatting.
|
|
FIXED: Enhanced threading object initialization to prevent None references.
|
|
"""
|
|
|
|
def __init__(self, session_config=None, socketio=None):
|
|
"""Initialize scanner with session-specific configuration."""
|
|
try:
|
|
# Use provided session config or create default
|
|
if session_config is None:
|
|
from core.session_config import create_session_config
|
|
session_config = create_session_config()
|
|
|
|
# FIXED: Initialize all threading objects first
|
|
self._initialize_threading_objects()
|
|
|
|
# Set socketio (but will be set to None for storage)
|
|
self.socketio = socketio
|
|
|
|
self.config = session_config
|
|
self.graph = GraphManager()
|
|
self.providers = []
|
|
self.status = ScanStatus.IDLE
|
|
self.current_target = None
|
|
self.current_depth = 0
|
|
self.max_depth = 2
|
|
self.scan_thread = None
|
|
self.session_id: Optional[str] = None # Will be set by session manager
|
|
self.initial_targets = set()
|
|
|
|
# Thread-safe processing tracking (from Document 1)
|
|
self.currently_processing = set()
|
|
# Display-friendly processing list (from Document 2)
|
|
self.currently_processing_display = []
|
|
|
|
# Scanning progress tracking
|
|
self.total_indicators_found = 0
|
|
self.indicators_processed = 0
|
|
self.indicators_completed = 0
|
|
self.tasks_re_enqueued = 0
|
|
self.tasks_skipped = 0 # BUGFIX: Initialize tasks_skipped
|
|
self.total_tasks_ever_enqueued = 0
|
|
self.current_indicator = ""
|
|
self.last_task_from_queue = None
|
|
|
|
# Concurrent processing configuration
|
|
self.max_workers = self.config.max_concurrent_requests
|
|
self.executor = None
|
|
|
|
# Initialize collections that will be recreated during unpickling
|
|
self.task_queue = PriorityQueue()
|
|
self.target_retries = defaultdict(int)
|
|
self.scan_failed_due_to_retries = False
|
|
|
|
# Initialize providers with session config
|
|
self._initialize_providers()
|
|
|
|
# Initialize logger
|
|
self.logger = get_forensic_logger()
|
|
|
|
# Initialize global rate limiter
|
|
self.rate_limiter = GlobalRateLimiter(redis.StrictRedis(db=0))
|
|
|
|
except Exception as e:
|
|
print(f"ERROR: Scanner initialization failed: {e}")
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
def _initialize_threading_objects(self):
|
|
"""
|
|
FIXED: Initialize all threading objects with proper error handling.
|
|
This method can be called during both __init__ and __setstate__.
|
|
"""
|
|
self.stop_event = threading.Event()
|
|
self.processing_lock = threading.Lock()
|
|
self.status_logger_stop_event = threading.Event()
|
|
self.status_logger_thread = None
|
|
|
|
def _is_stop_requested(self) -> bool:
|
|
"""
|
|
Check if stop is requested using both local and Redis-based signals.
|
|
This ensures reliable termination across process boundaries.
|
|
FIXED: Added None check for stop_event.
|
|
"""
|
|
# FIXED: Ensure stop_event exists before checking
|
|
if hasattr(self, 'stop_event') and self.stop_event and self.stop_event.is_set():
|
|
return True
|
|
|
|
if self.session_id:
|
|
try:
|
|
from core.session_manager import session_manager
|
|
return session_manager.is_stop_requested(self.session_id)
|
|
except Exception as e:
|
|
# Fall back to local event if it exists
|
|
if hasattr(self, 'stop_event') and self.stop_event:
|
|
return self.stop_event.is_set()
|
|
return False
|
|
|
|
# Final fallback
|
|
if hasattr(self, 'stop_event') and self.stop_event:
|
|
return self.stop_event.is_set()
|
|
return False
|
|
|
|
def _set_stop_signal(self) -> None:
|
|
"""
|
|
Set stop signal both locally and in Redis.
|
|
FIXED: Added None check for stop_event.
|
|
"""
|
|
# FIXED: Ensure stop_event exists before setting
|
|
if hasattr(self, 'stop_event') and self.stop_event:
|
|
self.stop_event.set()
|
|
|
|
if self.session_id:
|
|
try:
|
|
from core.session_manager import session_manager
|
|
session_manager.set_stop_signal(self.session_id)
|
|
except Exception as e:
|
|
pass
|
|
|
|
def __getstate__(self):
|
|
"""Prepare object for pickling by excluding unpicklable attributes."""
|
|
state = self.__dict__.copy()
|
|
|
|
unpicklable_attrs = [
|
|
'stop_event',
|
|
'scan_thread',
|
|
'executor',
|
|
'processing_lock',
|
|
'task_queue',
|
|
'rate_limiter',
|
|
'logger',
|
|
'status_logger_thread',
|
|
'status_logger_stop_event',
|
|
'socketio'
|
|
]
|
|
|
|
for attr in unpicklable_attrs:
|
|
if attr in state:
|
|
del state[attr]
|
|
|
|
if 'providers' in state:
|
|
for provider in state['providers']:
|
|
if hasattr(provider, '_stop_event'):
|
|
provider._stop_event = None
|
|
|
|
return state
|
|
|
|
def __setstate__(self, state):
|
|
"""Restore object after unpickling by reconstructing threading objects."""
|
|
self.__dict__.update(state)
|
|
|
|
# FIXED: Ensure all threading objects are properly initialized
|
|
self._initialize_threading_objects()
|
|
|
|
# Re-initialize other objects
|
|
self.scan_thread = None
|
|
self.executor = None
|
|
self.task_queue = PriorityQueue()
|
|
self.rate_limiter = GlobalRateLimiter(redis.StrictRedis(db=0))
|
|
self.logger = get_forensic_logger()
|
|
|
|
# FIXED: Initialize socketio as None but preserve ability to set it
|
|
if not hasattr(self, 'socketio'):
|
|
self.socketio = None
|
|
|
|
# Initialize missing attributes with defaults
|
|
if not hasattr(self, 'providers') or not self.providers:
|
|
self._initialize_providers()
|
|
|
|
if not hasattr(self, 'currently_processing'):
|
|
self.currently_processing = set()
|
|
|
|
if not hasattr(self, 'currently_processing_display'):
|
|
self.currently_processing_display = []
|
|
|
|
if not hasattr(self, 'target_retries'):
|
|
self.target_retries = defaultdict(int)
|
|
|
|
if not hasattr(self, 'scan_failed_due_to_retries'):
|
|
self.scan_failed_due_to_retries = False
|
|
|
|
if not hasattr(self, 'initial_targets'):
|
|
self.initial_targets = set()
|
|
|
|
# Ensure providers have stop events
|
|
if hasattr(self, 'providers'):
|
|
for provider in self.providers:
|
|
if hasattr(provider, 'set_stop_event') and self.stop_event:
|
|
provider.set_stop_event(self.stop_event)
|
|
|
|
def _ensure_threading_objects_exist(self):
|
|
"""
|
|
FIXED: Utility method to ensure threading objects exist before use.
|
|
Call this before any method that might use threading objects.
|
|
"""
|
|
if not hasattr(self, 'stop_event') or self.stop_event is None:
|
|
print("WARNING: Threading objects not initialized, recreating...")
|
|
self._initialize_threading_objects()
|
|
|
|
if not hasattr(self, 'processing_lock') or self.processing_lock is None:
|
|
self.processing_lock = threading.Lock()
|
|
|
|
if not hasattr(self, 'task_queue') or self.task_queue is None:
|
|
self.task_queue = PriorityQueue()
|
|
|
|
def _initialize_providers(self) -> None:
|
|
"""Initialize all available providers based on session configuration."""
|
|
self.providers = []
|
|
provider_dir = os.path.join(os.path.dirname(__file__), '..', 'providers')
|
|
|
|
print(f"=== INITIALIZING PROVIDERS FROM {provider_dir} ===")
|
|
|
|
for filename in os.listdir(provider_dir):
|
|
if filename.endswith('_provider.py') and not filename.startswith('base'):
|
|
module_name = f"providers.{filename[:-3]}"
|
|
try:
|
|
print(f"Loading provider module: {module_name}")
|
|
module = importlib.import_module(module_name)
|
|
for attribute_name in dir(module):
|
|
attribute = getattr(module, attribute_name)
|
|
if isinstance(attribute, type) and issubclass(attribute, BaseProvider) and attribute is not BaseProvider:
|
|
provider_class = attribute
|
|
# FIXED: Pass the 'name' argument during initialization
|
|
provider = provider_class(name=attribute_name, session_config=self.config)
|
|
provider_name = provider.get_name()
|
|
|
|
print(f" Provider: {provider_name}")
|
|
print(f" Class: {provider_class.__name__}")
|
|
print(f" Config enabled: {self.config.is_provider_enabled(provider_name)}")
|
|
print(f" Requires API key: {provider.requires_api_key()}")
|
|
|
|
if provider.requires_api_key():
|
|
api_key = self.config.get_api_key(provider_name)
|
|
print(f" API key present: {'Yes' if api_key else 'No'}")
|
|
if api_key:
|
|
print(f" API key preview: {api_key[:8]}...")
|
|
|
|
if self.config.is_provider_enabled(provider_name):
|
|
is_available = provider.is_available()
|
|
print(f" Available: {is_available}")
|
|
|
|
if is_available:
|
|
# FIXED: Ensure stop_event exists before setting it
|
|
if hasattr(self, 'stop_event') and self.stop_event:
|
|
provider.set_stop_event(self.stop_event)
|
|
if isinstance(provider, CorrelationProvider):
|
|
provider.set_graph_manager(self.graph)
|
|
self.providers.append(provider)
|
|
print(f" ✓ Added to scanner")
|
|
else:
|
|
print(f" ✗ Not available - skipped")
|
|
else:
|
|
print(f" ✗ Disabled in config - skipped")
|
|
|
|
except Exception as e:
|
|
print(f" ERROR loading {module_name}: {e}")
|
|
traceback.print_exc()
|
|
|
|
print(f"=== PROVIDER INITIALIZATION COMPLETE ===")
|
|
print(f"Active providers: {[p.get_name() for p in self.providers]}")
|
|
print(f"Provider count: {len(self.providers)}")
|
|
print("=" * 50)
|
|
|
|
def _status_logger_thread(self):
|
|
"""Periodically prints a clean, formatted scan status to the terminal."""
|
|
HEADER = "\033[95m"
|
|
CYAN = "\033[96m"
|
|
GREEN = "\033[92m"
|
|
YELLOW = "\033[93m"
|
|
BLUE = "\033[94m"
|
|
ENDC = "\033[0m"
|
|
BOLD = "\033[1m"
|
|
|
|
last_status_str = ""
|
|
|
|
# FIXED: Ensure threading objects exist
|
|
self._ensure_threading_objects_exist()
|
|
|
|
while not (hasattr(self, 'status_logger_stop_event') and
|
|
self.status_logger_stop_event and
|
|
self.status_logger_stop_event.is_set()):
|
|
try:
|
|
# FIXED: Check if processing_lock exists before using
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
in_flight_tasks = list(self.currently_processing)
|
|
self.currently_processing_display = in_flight_tasks.copy()
|
|
else:
|
|
in_flight_tasks = list(getattr(self, 'currently_processing', []))
|
|
|
|
status_str = (
|
|
f"{BOLD}{HEADER}Scan Status: {self.status.upper()}{ENDC} | "
|
|
f"{CYAN}Queued: {self.task_queue.qsize() if hasattr(self, 'task_queue') and self.task_queue else 0}{ENDC} | "
|
|
f"{YELLOW}In-Flight: {len(in_flight_tasks)}{ENDC} | "
|
|
f"{GREEN}Completed: {self.indicators_completed}{ENDC} | "
|
|
f"Skipped: {self.tasks_skipped} | "
|
|
f"Rescheduled: {self.tasks_re_enqueued}"
|
|
)
|
|
|
|
if status_str != last_status_str:
|
|
print(f"\n{'-'*80}")
|
|
print(status_str)
|
|
if self.last_task_from_queue:
|
|
# Unpack the new time-based queue item
|
|
_, p, (pn, ti, d) = self.last_task_from_queue
|
|
print(f"{BLUE}Last task dequeued -> Prio:{p} | Provider:{pn} | Target:'{ti}' | Depth:{d}{ENDC}")
|
|
if in_flight_tasks:
|
|
print(f"{BOLD}{YELLOW}Currently Processing:{ENDC}")
|
|
display_tasks = [f" - {p}: {t}" for p, t in in_flight_tasks[:3]]
|
|
print("\n".join(display_tasks))
|
|
if len(in_flight_tasks) > 3:
|
|
print(f" ... and {len(in_flight_tasks) - 3} more")
|
|
print(f"{'-'*80}")
|
|
last_status_str = status_str
|
|
except Exception:
|
|
pass
|
|
|
|
time.sleep(2)
|
|
|
|
def start_scan(self, target: str, max_depth: int = 2, clear_graph: bool = True, force_rescan_target: Optional[str] = None) -> bool:
|
|
"""
|
|
FIXED: Enhanced start_scan with proper threading object initialization and socketio management.
|
|
"""
|
|
# FIXED: Ensure threading objects exist before proceeding
|
|
self._ensure_threading_objects_exist()
|
|
|
|
if self.scan_thread and self.scan_thread.is_alive():
|
|
self.logger.logger.info("Stopping existing scan before starting new one")
|
|
self._set_stop_signal()
|
|
self.status = ScanStatus.STOPPED
|
|
|
|
# Clean up processing state
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.clear()
|
|
self.currently_processing_display = []
|
|
|
|
# Clear task queue
|
|
if hasattr(self, 'task_queue') and self.task_queue:
|
|
while not self.task_queue.empty():
|
|
try:
|
|
self.task_queue.get_nowait()
|
|
except:
|
|
break
|
|
|
|
# Shutdown executor
|
|
if self.executor:
|
|
try:
|
|
self.executor.shutdown(wait=False, cancel_futures=True)
|
|
except:
|
|
pass
|
|
finally:
|
|
self.executor = None
|
|
|
|
# Wait for scan thread to finish (with timeout)
|
|
self.scan_thread.join(timeout=5.0)
|
|
if self.scan_thread.is_alive():
|
|
self.logger.logger.warning("Previous scan thread did not terminate cleanly")
|
|
|
|
self.status = ScanStatus.IDLE
|
|
|
|
# FIXED: Ensure stop_event exists before clearing
|
|
if hasattr(self, 'stop_event') and self.stop_event:
|
|
self.stop_event.clear()
|
|
|
|
if self.session_id:
|
|
from core.session_manager import session_manager
|
|
session_manager.clear_stop_signal(self.session_id)
|
|
|
|
# FIXED: Restore socketio connection if missing
|
|
if not hasattr(self, 'socketio') or not self.socketio:
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
self.socketio = registered_socketio
|
|
print(f"✓ Restored socketio connection for scan start")
|
|
|
|
# FIXED: Safe cleanup with existence checks
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.clear()
|
|
self.currently_processing_display = []
|
|
|
|
self.task_queue = PriorityQueue()
|
|
self.target_retries.clear()
|
|
self.scan_failed_due_to_retries = False
|
|
self.tasks_skipped = 0
|
|
self.last_task_from_queue = None
|
|
|
|
self._update_session_state()
|
|
|
|
try:
|
|
if not hasattr(self, 'providers') or not self.providers:
|
|
self.logger.logger.error("No providers available for scanning")
|
|
return False
|
|
|
|
available_providers = [p for p in self.providers if p.is_available()]
|
|
if not available_providers:
|
|
self.logger.logger.error("No providers are currently available/configured")
|
|
return False
|
|
|
|
if clear_graph:
|
|
self.graph.clear()
|
|
self.initial_targets.clear()
|
|
|
|
if force_rescan_target and self.graph.graph.has_node(force_rescan_target):
|
|
try:
|
|
node_data = self.graph.graph.nodes[force_rescan_target]
|
|
if 'metadata' in node_data and 'provider_states' in node_data['metadata']:
|
|
node_data['metadata']['provider_states'] = {}
|
|
self.logger.logger.info(f"Cleared provider states for forced rescan of {force_rescan_target}")
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error clearing provider states for {force_rescan_target}: {e}")
|
|
|
|
target = target.lower().strip()
|
|
if not target:
|
|
self.logger.logger.error("Empty target provided")
|
|
return False
|
|
|
|
from utils.helpers import is_valid_target
|
|
if not is_valid_target(target):
|
|
self.logger.logger.error(f"Invalid target format: {target}")
|
|
return False
|
|
|
|
self.current_target = target
|
|
self.initial_targets.add(self.current_target)
|
|
self.max_depth = max(1, min(5, max_depth)) # Clamp depth between 1-5
|
|
self.current_depth = 0
|
|
|
|
self.total_indicators_found = 0
|
|
self.indicators_processed = 0
|
|
self.indicators_completed = 0
|
|
self.tasks_re_enqueued = 0
|
|
self.total_tasks_ever_enqueued = 0
|
|
self.current_indicator = self.current_target
|
|
|
|
self._update_session_state()
|
|
self.logger = new_session()
|
|
|
|
try:
|
|
self.scan_thread = threading.Thread(
|
|
target=self._execute_scan,
|
|
args=(self.current_target, self.max_depth),
|
|
daemon=True,
|
|
name=f"ScanThread-{self.session_id or 'default'}"
|
|
)
|
|
self.scan_thread.start()
|
|
|
|
# FIXED: Ensure status_logger_stop_event exists before clearing
|
|
if hasattr(self, 'status_logger_stop_event') and self.status_logger_stop_event:
|
|
self.status_logger_stop_event.clear()
|
|
|
|
self.status_logger_thread = threading.Thread(
|
|
target=self._status_logger_thread,
|
|
daemon=True,
|
|
name=f"StatusLogger-{self.session_id or 'default'}"
|
|
)
|
|
self.status_logger_thread.start()
|
|
|
|
self.logger.logger.info(f"Scan started successfully for {target} with depth {self.max_depth}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.logger.error(f"Error starting scan threads: {e}")
|
|
self.status = ScanStatus.FAILED
|
|
self._update_session_state()
|
|
return False
|
|
|
|
except Exception as e:
|
|
self.logger.logger.error(f"Error in scan startup: {e}")
|
|
traceback.print_exc()
|
|
self.status = ScanStatus.FAILED
|
|
self._update_session_state()
|
|
return False
|
|
|
|
def _get_priority(self, provider_name):
|
|
if provider_name == 'correlation':
|
|
return 100 # Highest priority number = lowest priority (runs last)
|
|
|
|
rate_limit = self.config.get_rate_limit(provider_name)
|
|
|
|
# Handle edge cases
|
|
if rate_limit <= 0:
|
|
return 90 # Very low priority for invalid/disabled providers
|
|
|
|
if provider_name == 'dns':
|
|
return 1 # DNS is fastest, should run first
|
|
elif provider_name == 'shodan':
|
|
return 3 # Shodan is medium speed, good priority
|
|
elif provider_name == 'crtsh':
|
|
return 5 # crt.sh is slower, lower priority
|
|
else:
|
|
# For any other providers, use rate limit as a guide
|
|
if rate_limit >= 100:
|
|
return 2 # High rate limit = high priority
|
|
elif rate_limit >= 50:
|
|
return 4 # Medium-high rate limit = medium-high priority
|
|
elif rate_limit >= 20:
|
|
return 6 # Medium rate limit = medium priority
|
|
elif rate_limit >= 5:
|
|
return 8 # Low rate limit = low priority
|
|
else:
|
|
return 10 # Very low rate limit = very low priority
|
|
|
|
def _execute_scan(self, target: str, max_depth: int) -> None:
|
|
"""
|
|
FIXED: Enhanced execute_scan with proper threading object handling.
|
|
"""
|
|
# FIXED: Ensure threading objects exist
|
|
self._ensure_threading_objects_exist()
|
|
update_counter = 0 # Track updates for throttling
|
|
last_update_time = time.time()
|
|
self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
|
|
processed_tasks = set() # FIXED: Now includes depth to avoid incorrect skipping
|
|
|
|
is_ip = _is_valid_ip(target)
|
|
initial_providers = self._get_eligible_providers(target, is_ip, False)
|
|
# FIXED: Filter out correlation provider from initial providers
|
|
initial_providers = [p for p in initial_providers if not isinstance(p, CorrelationProvider)]
|
|
|
|
for provider in initial_providers:
|
|
provider_name = provider.get_name()
|
|
priority = self._get_priority(provider_name)
|
|
self.task_queue.put((time.time(), priority, (provider_name, target, 0)))
|
|
self.total_tasks_ever_enqueued += 1
|
|
|
|
try:
|
|
self.status = ScanStatus.RUNNING
|
|
self._update_session_state()
|
|
|
|
enabled_providers = [provider.get_name() for provider in self.providers]
|
|
self.logger.log_scan_start(target, max_depth, enabled_providers)
|
|
|
|
node_type = NodeType.IP if is_ip else NodeType.DOMAIN
|
|
self.graph.add_node(target, node_type)
|
|
self._initialize_provider_states(target)
|
|
consecutive_empty_iterations = 0
|
|
max_empty_iterations = 50 # Allow 5 seconds of empty queue before considering completion
|
|
|
|
# PHASE 1: Run all non-correlation providers
|
|
print(f"\n=== PHASE 1: Running non-correlation providers ===")
|
|
while not self._is_stop_requested():
|
|
queue_empty = self.task_queue.empty()
|
|
|
|
# FIXED: Safe processing lock usage
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
no_active_processing = len(self.currently_processing) == 0
|
|
else:
|
|
no_active_processing = len(getattr(self, 'currently_processing', [])) == 0
|
|
|
|
if queue_empty and no_active_processing:
|
|
consecutive_empty_iterations += 1
|
|
if consecutive_empty_iterations >= max_empty_iterations:
|
|
break # Phase 1 complete
|
|
time.sleep(0.1)
|
|
continue
|
|
else:
|
|
consecutive_empty_iterations = 0
|
|
|
|
# Process tasks (same logic as before, but correlations are filtered out)
|
|
try:
|
|
run_at, priority, (provider_name, target_item, depth) = self.task_queue.get(timeout=0.1)
|
|
|
|
# Skip correlation tasks during Phase 1
|
|
if provider_name == 'correlation':
|
|
continue
|
|
|
|
# Check if task is ready to run
|
|
current_time = time.time()
|
|
if run_at > current_time:
|
|
self.task_queue.put((run_at, priority, (provider_name, target_item, depth)))
|
|
time.sleep(min(0.5, run_at - current_time))
|
|
continue
|
|
|
|
except: # Queue is empty or timeout occurred
|
|
time.sleep(0.1)
|
|
continue
|
|
|
|
self.last_task_from_queue = (run_at, priority, (provider_name, target_item, depth))
|
|
|
|
# Skip if already processed
|
|
task_tuple = (provider_name, target_item, depth)
|
|
if task_tuple in processed_tasks:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
|
|
# Skip if depth exceeded
|
|
if depth > max_depth:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
|
|
# Rate limiting with proper time-based deferral
|
|
if self.rate_limiter.is_rate_limited(provider_name, self.config.get_rate_limit(provider_name), 60):
|
|
defer_until = time.time() + 60
|
|
self.task_queue.put((defer_until, priority, (provider_name, target_item, depth)))
|
|
self.tasks_re_enqueued += 1
|
|
continue
|
|
|
|
# Thread-safe processing state management
|
|
processing_key = (provider_name, target_item)
|
|
|
|
# FIXED: Safe processing lock usage
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
if self._is_stop_requested():
|
|
break
|
|
if processing_key in self.currently_processing:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
self.currently_processing.add(processing_key)
|
|
else:
|
|
if self._is_stop_requested():
|
|
break
|
|
if not hasattr(self, 'currently_processing'):
|
|
self.currently_processing = set()
|
|
if processing_key in self.currently_processing:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
self.currently_processing.add(processing_key)
|
|
|
|
try:
|
|
self.current_depth = depth
|
|
self.current_indicator = target_item
|
|
self._update_session_state()
|
|
|
|
if self._is_stop_requested():
|
|
break
|
|
|
|
provider = next((p for p in self.providers if p.get_name() == provider_name), None)
|
|
|
|
if provider and not isinstance(provider, CorrelationProvider):
|
|
new_targets, _, success = self._process_provider_task(provider, target_item, depth)
|
|
update_counter += 1
|
|
current_time = time.time()
|
|
if (update_counter % 5 == 0) or (current_time - last_update_time > 3.0):
|
|
self._update_session_state()
|
|
last_update_time = current_time
|
|
update_counter = 0
|
|
if self._is_stop_requested():
|
|
break
|
|
|
|
if not success:
|
|
retry_key = (provider_name, target_item, depth)
|
|
self.target_retries[retry_key] += 1
|
|
|
|
if self.target_retries[retry_key] <= self.config.max_retries_per_target:
|
|
retry_count = self.target_retries[retry_key]
|
|
backoff_delay = min(300, (2 ** retry_count) + random.uniform(0, 1))
|
|
retry_at = time.time() + backoff_delay
|
|
self.task_queue.put((retry_at, priority, (provider_name, target_item, depth)))
|
|
self.tasks_re_enqueued += 1
|
|
self.logger.logger.debug(f"Retrying {provider_name}:{target_item} in {backoff_delay:.1f}s (attempt {retry_count})")
|
|
else:
|
|
self.scan_failed_due_to_retries = True
|
|
self._log_target_processing_error(str(task_tuple), f"Max retries ({self.config.max_retries_per_target}) exceeded")
|
|
else:
|
|
processed_tasks.add(task_tuple)
|
|
self.indicators_completed += 1
|
|
|
|
# Enqueue new targets with proper depth tracking
|
|
if not self._is_stop_requested():
|
|
for new_target in new_targets:
|
|
is_ip_new = _is_valid_ip(new_target)
|
|
eligible_providers_new = self._get_eligible_providers(new_target, is_ip_new, False)
|
|
# FIXED: Filter out correlation providers in Phase 1
|
|
eligible_providers_new = [p for p in eligible_providers_new if not isinstance(p, CorrelationProvider)]
|
|
|
|
for p_new in eligible_providers_new:
|
|
p_name_new = p_new.get_name()
|
|
new_depth = depth + 1
|
|
new_task_tuple = (p_name_new, new_target, new_depth)
|
|
|
|
if new_task_tuple not in processed_tasks and new_depth <= max_depth:
|
|
new_priority = self._get_priority(p_name_new)
|
|
self.task_queue.put((time.time(), new_priority, (p_name_new, new_target, new_depth)))
|
|
self.total_tasks_ever_enqueued += 1
|
|
else:
|
|
self.logger.logger.warning(f"Provider {provider_name} not found in active providers")
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
|
|
finally:
|
|
# FIXED: Safe processing lock usage for cleanup
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.discard(processing_key)
|
|
else:
|
|
if hasattr(self, 'currently_processing'):
|
|
self.currently_processing.discard(processing_key)
|
|
|
|
# PHASE 2: Run correlations on all discovered nodes
|
|
if not self._is_stop_requested():
|
|
print(f"\n=== PHASE 2: Running correlation analysis ===")
|
|
self._run_correlation_phase(max_depth, processed_tasks)
|
|
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
self.status = ScanStatus.FAILED
|
|
self.logger.logger.error(f"Scan failed: {e}")
|
|
finally:
|
|
# Comprehensive cleanup (same as before)
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.clear()
|
|
self.currently_processing_display = []
|
|
|
|
while not self.task_queue.empty():
|
|
try:
|
|
self.task_queue.get_nowait()
|
|
except:
|
|
break
|
|
|
|
if self._is_stop_requested():
|
|
self.status = ScanStatus.STOPPED
|
|
elif self.scan_failed_due_to_retries:
|
|
self.status = ScanStatus.FAILED
|
|
else:
|
|
self.status = ScanStatus.COMPLETED
|
|
|
|
# FIXED: Safe stop event handling
|
|
if hasattr(self, 'status_logger_stop_event') and self.status_logger_stop_event:
|
|
self.status_logger_stop_event.set()
|
|
if self.status_logger_thread and self.status_logger_thread.is_alive():
|
|
self.status_logger_thread.join(timeout=2.0)
|
|
|
|
self._update_session_state()
|
|
self.logger.log_scan_complete()
|
|
|
|
if self.executor:
|
|
try:
|
|
self.executor.shutdown(wait=False, cancel_futures=True)
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error shutting down executor: {e}")
|
|
finally:
|
|
self.executor = None
|
|
|
|
def _run_correlation_phase(self, max_depth: int, processed_tasks: set) -> None:
|
|
"""
|
|
PHASE 2: Run correlation analysis on all discovered nodes.
|
|
This ensures correlations run after all other providers have completed.
|
|
"""
|
|
correlation_provider = next((p for p in self.providers if isinstance(p, CorrelationProvider)), None)
|
|
if not correlation_provider:
|
|
print("No correlation provider found - skipping correlation phase")
|
|
return
|
|
|
|
# Get all nodes from the graph for correlation analysis
|
|
all_nodes = list(self.graph.graph.nodes())
|
|
correlation_tasks = []
|
|
|
|
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)
|
|
|
|
# Don't re-process already processed correlation tasks
|
|
if task_tuple not in processed_tasks:
|
|
priority = self._get_priority('correlation')
|
|
self.task_queue.put((time.time(), priority, ('correlation', node_id, correlation_depth)))
|
|
correlation_tasks.append(task_tuple)
|
|
self.total_tasks_ever_enqueued += 1
|
|
|
|
print(f"Enqueued {len(correlation_tasks)} correlation tasks")
|
|
|
|
# Process correlation tasks
|
|
consecutive_empty_iterations = 0
|
|
max_empty_iterations = 20 # Shorter timeout for correlation phase
|
|
|
|
while not self._is_stop_requested() and correlation_tasks:
|
|
queue_empty = self.task_queue.empty()
|
|
|
|
# FIXED: Safe processing check
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
no_active_processing = len(self.currently_processing) == 0
|
|
else:
|
|
no_active_processing = len(getattr(self, 'currently_processing', [])) == 0
|
|
|
|
if queue_empty and no_active_processing:
|
|
consecutive_empty_iterations += 1
|
|
if consecutive_empty_iterations >= max_empty_iterations:
|
|
break
|
|
time.sleep(0.1)
|
|
continue
|
|
else:
|
|
consecutive_empty_iterations = 0
|
|
|
|
try:
|
|
run_at, priority, (provider_name, target_item, depth) = self.task_queue.get(timeout=0.1)
|
|
|
|
# Only process correlation tasks in this phase
|
|
if provider_name != 'correlation':
|
|
continue
|
|
|
|
except:
|
|
time.sleep(0.1)
|
|
continue
|
|
|
|
task_tuple = (provider_name, target_item, depth)
|
|
|
|
# Skip if already processed
|
|
if task_tuple in processed_tasks:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
if task_tuple in correlation_tasks:
|
|
correlation_tasks.remove(task_tuple)
|
|
continue
|
|
|
|
processing_key = (provider_name, target_item)
|
|
|
|
# FIXED: Safe processing management
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
if self._is_stop_requested():
|
|
break
|
|
if processing_key in self.currently_processing:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
self.currently_processing.add(processing_key)
|
|
else:
|
|
if self._is_stop_requested():
|
|
break
|
|
if not hasattr(self, 'currently_processing'):
|
|
self.currently_processing = set()
|
|
if processing_key in self.currently_processing:
|
|
self.tasks_skipped += 1
|
|
self.indicators_completed += 1
|
|
continue
|
|
self.currently_processing.add(processing_key)
|
|
|
|
try:
|
|
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)
|
|
|
|
if success:
|
|
processed_tasks.add(task_tuple)
|
|
self.indicators_completed += 1
|
|
if task_tuple in correlation_tasks:
|
|
correlation_tasks.remove(task_tuple)
|
|
else:
|
|
# For correlations, don't retry - just mark as completed
|
|
self.indicators_completed += 1
|
|
if task_tuple in correlation_tasks:
|
|
correlation_tasks.remove(task_tuple)
|
|
|
|
finally:
|
|
# FIXED: Safe cleanup
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.discard(processing_key)
|
|
else:
|
|
if hasattr(self, 'currently_processing'):
|
|
self.currently_processing.discard(processing_key)
|
|
|
|
print(f"Correlation phase complete. Remaining tasks: {len(correlation_tasks)}")
|
|
|
|
# Rest of the methods remain the same but with similar threading object safety checks...
|
|
# I'll include the key methods that need fixes:
|
|
|
|
def stop_scan(self) -> bool:
|
|
"""Request immediate scan termination with proper cleanup."""
|
|
try:
|
|
self.logger.logger.info("Scan termination requested by user")
|
|
self._set_stop_signal()
|
|
self.status = ScanStatus.STOPPED
|
|
|
|
# FIXED: Safe cleanup
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
self.currently_processing.clear()
|
|
self.currently_processing_display = []
|
|
|
|
self.task_queue = PriorityQueue()
|
|
|
|
if self.executor:
|
|
try:
|
|
self.executor.shutdown(wait=False, cancel_futures=True)
|
|
except Exception:
|
|
pass
|
|
|
|
self._update_session_state()
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.logger.error(f"Error during scan termination: {e}")
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
def get_scan_status(self) -> Dict[str, Any]:
|
|
"""Get current scan status with comprehensive graph data for real-time updates."""
|
|
try:
|
|
# FIXED: Safe processing state access
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
currently_processing_count = len(self.currently_processing)
|
|
currently_processing_list = list(self.currently_processing)
|
|
else:
|
|
currently_processing_count = len(getattr(self, 'currently_processing', []))
|
|
currently_processing_list = list(getattr(self, 'currently_processing', []))
|
|
|
|
# FIXED: Always include complete graph data for real-time updates
|
|
graph_data = self.get_graph_data()
|
|
|
|
return {
|
|
'status': self.status,
|
|
'target_domain': self.current_target,
|
|
'current_depth': self.current_depth,
|
|
'max_depth': self.max_depth,
|
|
'current_indicator': self.current_indicator,
|
|
'indicators_processed': self.indicators_processed,
|
|
'indicators_completed': self.indicators_completed,
|
|
'tasks_re_enqueued': self.tasks_re_enqueued,
|
|
'progress_percentage': self._calculate_progress(),
|
|
'total_tasks_ever_enqueued': self.total_tasks_ever_enqueued,
|
|
'enabled_providers': [provider.get_name() for provider in self.providers],
|
|
'graph': graph_data, # FIXED: Always include complete graph data
|
|
'task_queue_size': self.task_queue.qsize() if hasattr(self, 'task_queue') and self.task_queue else 0,
|
|
'currently_processing_count': currently_processing_count,
|
|
'currently_processing': currently_processing_list[:5],
|
|
'tasks_in_queue': self.task_queue.qsize() if hasattr(self, 'task_queue') and self.task_queue else 0,
|
|
'tasks_completed': self.indicators_completed,
|
|
'tasks_skipped': self.tasks_skipped,
|
|
'tasks_rescheduled': self.tasks_re_enqueued,
|
|
}
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return {
|
|
'status': 'error',
|
|
'message': 'Failed to get status',
|
|
'graph': {'nodes': [], 'edges': [], 'statistics': {'node_count': 0, 'edge_count': 0}}
|
|
}
|
|
|
|
def _update_session_state(self) -> None:
|
|
"""
|
|
FIXED: Update the scanner state in Redis and emit real-time WebSocket updates.
|
|
Enhanced with better error handling and socketio management.
|
|
"""
|
|
if self.session_id:
|
|
try:
|
|
# Get current status for WebSocket emission
|
|
current_status = self.get_scan_status()
|
|
|
|
# FIXED: Emit real-time update via WebSocket with better error handling
|
|
socketio_available = False
|
|
if hasattr(self, 'socketio') and self.socketio:
|
|
try:
|
|
print(f"📡 EMITTING WebSocket update: {current_status.get('status', 'unknown')} - "
|
|
f"Nodes: {len(current_status.get('graph', {}).get('nodes', []))}, "
|
|
f"Edges: {len(current_status.get('graph', {}).get('edges', []))}")
|
|
|
|
self.socketio.emit('scan_update', current_status)
|
|
print("✅ WebSocket update emitted successfully")
|
|
socketio_available = True
|
|
|
|
except Exception as ws_error:
|
|
print(f"⚠️ WebSocket emission failed: {ws_error}")
|
|
# Try to get socketio from session manager
|
|
try:
|
|
from core.session_manager import session_manager
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
print("🔄 Attempting to use registered socketio connection...")
|
|
registered_socketio.emit('scan_update', current_status)
|
|
self.socketio = registered_socketio # Update our reference
|
|
print("✅ WebSocket update emitted via registered connection")
|
|
socketio_available = True
|
|
else:
|
|
print("⚠️ No registered socketio connection found")
|
|
except Exception as fallback_error:
|
|
print(f"⚠️ Fallback socketio emission also failed: {fallback_error}")
|
|
else:
|
|
# Try to restore socketio from session manager
|
|
try:
|
|
from core.session_manager import session_manager
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
print(f"🔄 Restoring socketio connection for session {self.session_id}")
|
|
self.socketio = registered_socketio
|
|
self.socketio.emit('scan_update', current_status)
|
|
print("✅ WebSocket update emitted via restored connection")
|
|
socketio_available = True
|
|
else:
|
|
print(f"⚠️ No socketio connection available for session {self.session_id}")
|
|
except Exception as restore_error:
|
|
print(f"⚠️ Failed to restore socketio connection: {restore_error}")
|
|
|
|
if not socketio_available:
|
|
print(f"⚠️ Real-time updates unavailable for session {self.session_id}")
|
|
|
|
# Update session in Redis for persistence (always do this)
|
|
try:
|
|
from core.session_manager import session_manager
|
|
session_manager.update_session_scanner(self.session_id, self)
|
|
except Exception as redis_error:
|
|
print(f"⚠️ Failed to update session in Redis: {redis_error}")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Error updating session state: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
def _process_provider_task(self, provider: BaseProvider, target: str, depth: int) -> Tuple[Set[str], Set[str], bool]:
|
|
"""
|
|
FIXED: Manages the entire process for a given target and provider with enhanced real-time updates.
|
|
"""
|
|
if self._is_stop_requested():
|
|
return set(), set(), False
|
|
|
|
is_ip = _is_valid_ip(target)
|
|
target_type = NodeType.IP if is_ip else NodeType.DOMAIN
|
|
|
|
self.graph.add_node(target, target_type)
|
|
self._initialize_provider_states(target)
|
|
|
|
new_targets = set()
|
|
provider_successful = True
|
|
|
|
try:
|
|
provider_result = self._execute_provider_query(provider, target, is_ip)
|
|
|
|
if provider_result is None:
|
|
provider_successful = False
|
|
elif not self._is_stop_requested():
|
|
discovered, is_large_entity = self._process_provider_result_unified(
|
|
target, provider, provider_result, depth
|
|
)
|
|
new_targets.update(discovered)
|
|
|
|
# FIXED: Emit real-time update after processing provider result
|
|
if discovered or provider_result.get_relationship_count() > 0:
|
|
# Ensure we have socketio connection for real-time updates
|
|
if self.session_id and (not hasattr(self, 'socketio') or not self.socketio):
|
|
try:
|
|
from core.session_manager import session_manager
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
self.socketio = registered_socketio
|
|
print(f"🔄 Restored socketio connection during provider processing")
|
|
except Exception as restore_error:
|
|
print(f"⚠️ Failed to restore socketio during provider processing: {restore_error}")
|
|
|
|
self._update_session_state()
|
|
|
|
except Exception as e:
|
|
provider_successful = False
|
|
self._log_provider_error(target, provider.get_name(), str(e))
|
|
|
|
return new_targets, set(), provider_successful
|
|
|
|
|
|
|
|
def _execute_provider_query(self, provider: BaseProvider, target: str, is_ip: bool) -> Optional[ProviderResult]:
|
|
"""The "worker" function that directly communicates with the provider to fetch data."""
|
|
provider_name = provider.get_name()
|
|
start_time = datetime.now(timezone.utc)
|
|
|
|
if self._is_stop_requested():
|
|
return None
|
|
|
|
try:
|
|
if is_ip:
|
|
result = provider.query_ip(target)
|
|
else:
|
|
result = provider.query_domain(target)
|
|
|
|
if self._is_stop_requested():
|
|
return None
|
|
|
|
relationship_count = result.get_relationship_count() if result else 0
|
|
self._update_provider_state(target, provider_name, 'success', relationship_count, None, start_time)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
self._update_provider_state(target, provider_name, 'failed', 0, str(e), start_time)
|
|
return None
|
|
|
|
def _create_large_entity_from_result(self, source_node: str, provider_name: str,
|
|
provider_result: ProviderResult, depth: int) -> Tuple[str, Set[str]]:
|
|
"""Creates a large entity node, tags all member nodes, and returns its ID and members."""
|
|
members = {rel.target_node for rel in provider_result.relationships
|
|
if _is_valid_domain(rel.target_node) or _is_valid_ip(rel.target_node)}
|
|
|
|
if not members:
|
|
return "", set()
|
|
|
|
large_entity_id = f"le_{provider_name}_{source_node}"
|
|
|
|
self.graph.add_node(
|
|
node_id=large_entity_id,
|
|
node_type=NodeType.LARGE_ENTITY,
|
|
attributes=[
|
|
{"name": "count", "value": len(members), "type": "statistic"},
|
|
{"name": "source_provider", "value": provider_name, "type": "metadata"},
|
|
{"name": "discovery_depth", "value": depth, "type": "metadata"},
|
|
{"name": "nodes", "value": list(members), "type": "metadata"}
|
|
],
|
|
description=f"A collection of {len(members)} nodes discovered from {source_node} via {provider_name}."
|
|
)
|
|
|
|
for member_id in members:
|
|
node_type = NodeType.IP if _is_valid_ip(member_id) else NodeType.DOMAIN
|
|
self.graph.add_node(
|
|
node_id=member_id,
|
|
node_type=node_type,
|
|
metadata={'large_entity_id': large_entity_id}
|
|
)
|
|
|
|
return large_entity_id, members
|
|
|
|
def extract_node_from_large_entity(self, large_entity_id: str, node_id: str) -> bool:
|
|
"""
|
|
FIXED: Removes a node from a large entity with immediate real-time update.
|
|
"""
|
|
if not self.graph.graph.has_node(node_id):
|
|
return False
|
|
|
|
node_data = self.graph.graph.nodes[node_id]
|
|
metadata = node_data.get('metadata', {})
|
|
|
|
if metadata.get('large_entity_id') == large_entity_id:
|
|
# Remove the large entity tag
|
|
del metadata['large_entity_id']
|
|
self.graph.add_node(node_id, NodeType(node_data['type']), metadata=metadata)
|
|
|
|
# Re-enqueue the node for full processing
|
|
is_ip = _is_valid_ip(node_id)
|
|
eligible_providers = self._get_eligible_providers(node_id, is_ip, False)
|
|
for provider in eligible_providers:
|
|
provider_name = provider.get_name()
|
|
priority = self._get_priority(provider_name)
|
|
depth = 0
|
|
if self.graph.graph.has_node(large_entity_id):
|
|
le_attrs = self.graph.graph.nodes[large_entity_id].get('attributes', [])
|
|
depth_attr = next((a for a in le_attrs if a['name'] == 'discovery_depth'), None)
|
|
if depth_attr:
|
|
depth = depth_attr['value']
|
|
|
|
self.task_queue.put((time.time(), priority, (provider_name, node_id, depth)))
|
|
self.total_tasks_ever_enqueued += 1
|
|
|
|
# FIXED: Emit real-time update after extraction with socketio management
|
|
if self.session_id and (not hasattr(self, 'socketio') or not self.socketio):
|
|
try:
|
|
from core.session_manager import session_manager
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
self.socketio = registered_socketio
|
|
print(f"🔄 Restored socketio for node extraction update")
|
|
except Exception as restore_error:
|
|
print(f"⚠️ Failed to restore socketio for extraction: {restore_error}")
|
|
|
|
self._update_session_state()
|
|
|
|
return True
|
|
|
|
return False
|
|
|
|
def _process_provider_result_unified(self, target: str, provider: BaseProvider,
|
|
provider_result: ProviderResult, current_depth: int) -> Tuple[Set[str], bool]:
|
|
"""
|
|
FIXED: Process a unified ProviderResult object to update the graph with enhanced real-time updates.
|
|
"""
|
|
provider_name = provider.get_name()
|
|
discovered_targets = set()
|
|
large_entity_id = ""
|
|
large_entity_members = set()
|
|
|
|
if self._is_stop_requested():
|
|
return discovered_targets, False
|
|
|
|
eligible_rel_count = sum(
|
|
1 for rel in provider_result.relationships if _is_valid_domain(rel.target_node) or _is_valid_ip(rel.target_node)
|
|
)
|
|
is_large_entity = eligible_rel_count > self.config.large_entity_threshold
|
|
|
|
if is_large_entity:
|
|
large_entity_id, large_entity_members = self._create_large_entity_from_result(
|
|
target, provider_name, provider_result, current_depth
|
|
)
|
|
|
|
# Track if we added anything significant
|
|
nodes_added = 0
|
|
edges_added = 0
|
|
|
|
for i, relationship in enumerate(provider_result.relationships):
|
|
if i % 5 == 0 and self._is_stop_requested():
|
|
break
|
|
|
|
source_node_id = relationship.source_node
|
|
target_node_id = relationship.target_node
|
|
|
|
# Determine visual source and target, substituting with large entity ID if necessary
|
|
visual_source = large_entity_id if source_node_id in large_entity_members else source_node_id
|
|
visual_target = large_entity_id if target_node_id in large_entity_members else target_node_id
|
|
|
|
# Prevent self-loops on the large entity node
|
|
if visual_source == visual_target:
|
|
continue
|
|
|
|
# Determine node types for the actual nodes
|
|
source_type = NodeType.IP if _is_valid_ip(source_node_id) else NodeType.DOMAIN
|
|
if provider_name == 'shodan' and relationship.relationship_type == 'shodan_isp':
|
|
target_type = NodeType.ISP
|
|
elif provider_name == 'crtsh' and relationship.relationship_type == 'crtsh_cert_issuer':
|
|
target_type = NodeType.CA
|
|
elif provider_name == 'correlation':
|
|
target_type = NodeType.CORRELATION_OBJECT
|
|
elif _is_valid_ip(target_node_id):
|
|
target_type = NodeType.IP
|
|
else:
|
|
target_type = NodeType.DOMAIN
|
|
|
|
max_depth_reached = current_depth >= self.max_depth
|
|
|
|
# Add actual nodes to the graph (they might be hidden by the UI)
|
|
if self.graph.add_node(source_node_id, source_type):
|
|
nodes_added += 1
|
|
if self.graph.add_node(target_node_id, target_type, metadata={'max_depth_reached': max_depth_reached}):
|
|
nodes_added += 1
|
|
|
|
# Add the visual edge to the graph
|
|
if self.graph.add_edge(
|
|
visual_source, visual_target,
|
|
relationship.relationship_type,
|
|
relationship.confidence,
|
|
provider_name,
|
|
relationship.raw_data
|
|
):
|
|
edges_added += 1
|
|
|
|
if (_is_valid_domain(target_node_id) or _is_valid_ip(target_node_id)) and not max_depth_reached:
|
|
if target_node_id not in large_entity_members:
|
|
discovered_targets.add(target_node_id)
|
|
|
|
if large_entity_members:
|
|
self.logger.logger.info(f"Enqueuing DNS and Correlation for {len(large_entity_members)} members of {large_entity_id}")
|
|
for member in large_entity_members:
|
|
for provider_name_to_run in ['dns', 'correlation']:
|
|
p_instance = next((p for p in self.providers if p.get_name() == provider_name_to_run), None)
|
|
if p_instance and p_instance.get_eligibility().get('domains' if _is_valid_domain(member) else 'ips'):
|
|
priority = self._get_priority(provider_name_to_run)
|
|
self.task_queue.put((time.time(), priority, (provider_name_to_run, member, current_depth)))
|
|
self.total_tasks_ever_enqueued += 1
|
|
|
|
attributes_by_node = defaultdict(list)
|
|
for attribute in provider_result.attributes:
|
|
attr_dict = {
|
|
"name": attribute.name, "value": attribute.value, "type": attribute.type,
|
|
"provider": attribute.provider, "confidence": attribute.confidence, "metadata": attribute.metadata
|
|
}
|
|
attributes_by_node[attribute.target_node].append(attr_dict)
|
|
|
|
for node_id, node_attributes_list in attributes_by_node.items():
|
|
if not self.graph.graph.has_node(node_id):
|
|
node_type = NodeType.IP if _is_valid_ip(node_id) else NodeType.DOMAIN
|
|
self.graph.add_node(node_id, node_type, attributes=node_attributes_list)
|
|
nodes_added += 1
|
|
else:
|
|
existing_attrs = self.graph.graph.nodes[node_id].get('attributes', [])
|
|
self.graph.graph.nodes[node_id]['attributes'] = existing_attrs + node_attributes_list
|
|
|
|
# FIXED: Emit real-time update if we added anything significant
|
|
if nodes_added > 0 or edges_added > 0:
|
|
print(f"🔄 Added {nodes_added} nodes, {edges_added} edges - triggering real-time update")
|
|
|
|
# Ensure we have socketio connection for immediate update
|
|
if self.session_id and (not hasattr(self, 'socketio') or not self.socketio):
|
|
try:
|
|
from core.session_manager import session_manager
|
|
registered_socketio = session_manager.get_socketio_connection(self.session_id)
|
|
if registered_socketio:
|
|
self.socketio = registered_socketio
|
|
print(f"🔄 Restored socketio for immediate update")
|
|
except Exception as restore_error:
|
|
print(f"⚠️ Failed to restore socketio for immediate update: {restore_error}")
|
|
|
|
self._update_session_state()
|
|
|
|
return discovered_targets, is_large_entity
|
|
|
|
def _initialize_provider_states(self, target: str) -> None:
|
|
"""FIXED: Safer provider state initialization with error handling."""
|
|
try:
|
|
if not self.graph.graph.has_node(target):
|
|
return
|
|
|
|
node_data = self.graph.graph.nodes[target]
|
|
if 'metadata' not in node_data:
|
|
node_data['metadata'] = {}
|
|
if 'provider_states' not in node_data['metadata']:
|
|
node_data['metadata']['provider_states'] = {}
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error initializing provider states for {target}: {e}")
|
|
|
|
def _get_eligible_providers(self, target: str, is_ip: bool, dns_only: bool) -> List:
|
|
"""FIXED: Improved provider eligibility checking with better filtering."""
|
|
if dns_only:
|
|
return [p for p in self.providers if p.get_name() == 'dns']
|
|
|
|
eligible = []
|
|
target_key = 'ips' if is_ip else 'domains'
|
|
|
|
# Check if the target is part of a large entity
|
|
is_in_large_entity = False
|
|
if self.graph.graph.has_node(target):
|
|
metadata = self.graph.graph.nodes[target].get('metadata', {})
|
|
if 'large_entity_id' in metadata:
|
|
is_in_large_entity = True
|
|
|
|
for provider in self.providers:
|
|
try:
|
|
# If in large entity, only allow dns and correlation providers
|
|
if is_in_large_entity and provider.get_name() not in ['dns', 'correlation']:
|
|
continue
|
|
|
|
# Check if provider supports this target type
|
|
if not provider.get_eligibility().get(target_key, False):
|
|
continue
|
|
|
|
# Check if provider is available/configured
|
|
if not provider.is_available():
|
|
continue
|
|
|
|
# Check if we already successfully queried this provider
|
|
if not self._already_queried_provider(target, provider.get_name()):
|
|
eligible.append(provider)
|
|
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error checking provider eligibility {provider.get_name()}: {e}")
|
|
continue
|
|
|
|
return eligible
|
|
|
|
def _already_queried_provider(self, target: str, provider_name: str) -> bool:
|
|
"""FIXED: More robust check for already queried providers with proper error handling."""
|
|
try:
|
|
if not self.graph.graph.has_node(target):
|
|
return False
|
|
|
|
node_data = self.graph.graph.nodes[target]
|
|
provider_states = node_data.get('metadata', {}).get('provider_states', {})
|
|
provider_state = provider_states.get(provider_name)
|
|
|
|
# Only consider it already queried if it was successful
|
|
return (provider_state is not None and
|
|
provider_state.get('status') == 'success' and
|
|
provider_state.get('results_count', 0) > 0)
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error checking provider state for {target}:{provider_name}: {e}")
|
|
return False
|
|
|
|
def _update_provider_state(self, target: str, provider_name: str, status: str,
|
|
results_count: int, error: Optional[str], start_time: datetime) -> None:
|
|
"""FIXED: More robust provider state updates with validation."""
|
|
try:
|
|
if not self.graph.graph.has_node(target):
|
|
self.logger.logger.warning(f"Cannot update provider state: node {target} not found")
|
|
return
|
|
|
|
node_data = self.graph.graph.nodes[target]
|
|
if 'metadata' not in node_data:
|
|
node_data['metadata'] = {}
|
|
if 'provider_states' not in node_data['metadata']:
|
|
node_data['metadata']['provider_states'] = {}
|
|
|
|
# Calculate duration safely
|
|
try:
|
|
duration_ms = (datetime.now(timezone.utc) - start_time).total_seconds() * 1000
|
|
except Exception:
|
|
duration_ms = 0
|
|
|
|
node_data['metadata']['provider_states'][provider_name] = {
|
|
'status': status,
|
|
'timestamp': start_time.isoformat(),
|
|
'results_count': max(0, results_count), # Ensure non-negative
|
|
'error': str(error) if error else None,
|
|
'duration_ms': duration_ms
|
|
}
|
|
|
|
# Update last modified time for forensic integrity
|
|
if hasattr(self, 'last_modified'):
|
|
self.last_modified = datetime.now(timezone.utc).isoformat()
|
|
|
|
except Exception as e:
|
|
self.logger.logger.error(f"Error updating provider state for {target}:{provider_name}: {e}")
|
|
|
|
def _log_target_processing_error(self, target: str, error: str) -> None:
|
|
self.logger.logger.error(f"Target processing failed for {target}: {error}")
|
|
|
|
def _log_provider_error(self, target: str, provider_name: str, error: str) -> None:
|
|
self.logger.logger.error(f"Provider {provider_name} failed for {target}: {error}")
|
|
|
|
def _calculate_progress(self) -> float:
|
|
try:
|
|
if self.total_tasks_ever_enqueued == 0:
|
|
return 0.0
|
|
|
|
# Add small buffer for tasks still in queue to avoid showing 100% too early
|
|
queue_size = max(0, self.task_queue.qsize() if hasattr(self, 'task_queue') and self.task_queue else 0)
|
|
|
|
# FIXED: Safe processing count
|
|
if hasattr(self, 'processing_lock') and self.processing_lock:
|
|
with self.processing_lock:
|
|
active_tasks = len(self.currently_processing)
|
|
else:
|
|
active_tasks = len(getattr(self, 'currently_processing', []))
|
|
|
|
# Adjust total to account for remaining work
|
|
adjusted_total = max(self.total_tasks_ever_enqueued,
|
|
self.indicators_completed + queue_size + active_tasks)
|
|
|
|
if adjusted_total == 0:
|
|
return 100.0
|
|
|
|
progress = (self.indicators_completed / adjusted_total) * 100
|
|
return max(0.0, min(100.0, progress)) # Clamp between 0 and 100
|
|
|
|
except Exception as e:
|
|
self.logger.logger.warning(f"Error calculating progress: {e}")
|
|
return 0.0
|
|
|
|
def get_graph_data(self) -> Dict[str, Any]:
|
|
"""Get current graph data formatted for frontend visualization."""
|
|
graph_data = self.graph.get_graph_data()
|
|
graph_data['initial_targets'] = list(self.initial_targets)
|
|
return graph_data
|
|
|
|
def get_provider_info(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Get comprehensive information about all available providers."""
|
|
info = {}
|
|
provider_dir = os.path.join(os.path.dirname(__file__), '..', 'providers')
|
|
for filename in os.listdir(provider_dir):
|
|
if filename.endswith('_provider.py') and not filename.startswith('base'):
|
|
module_name = f"providers.{filename[:-3]}"
|
|
try:
|
|
module = importlib.import_module(module_name)
|
|
for attribute_name in dir(module):
|
|
attribute = getattr(module, attribute_name)
|
|
if isinstance(attribute, type) and issubclass(attribute, BaseProvider) and attribute is not BaseProvider:
|
|
provider_class = attribute
|
|
temp_provider = provider_class(name=attribute_name, session_config=self.config)
|
|
provider_name = temp_provider.get_name()
|
|
live_provider = next((p for p in self.providers if p.get_name() == provider_name), None)
|
|
info[provider_name] = {
|
|
'display_name': temp_provider.get_display_name(),
|
|
'requires_api_key': temp_provider.requires_api_key(),
|
|
'statistics': live_provider.get_statistics() if live_provider else temp_provider.get_statistics(),
|
|
'enabled': self.config.is_provider_enabled(provider_name),
|
|
'rate_limit': self.config.get_rate_limit(provider_name),
|
|
}
|
|
except Exception:
|
|
traceback.print_exc()
|
|
return info |