Compare commits
2 Commits
434d1f4803
...
try-fix
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4378146d0c | ||
|
|
b26002eff9 |
34
.env.example
34
.env.example
@@ -1,34 +0,0 @@
|
||||
# ===============================================
|
||||
# DNSRecon Environment Variables
|
||||
# ===============================================
|
||||
# Copy this file to .env and fill in your values.
|
||||
|
||||
# --- API Keys ---
|
||||
# Add your Shodan API key for the Shodan provider to be enabled.
|
||||
SHODAN_API_KEY=
|
||||
|
||||
# --- Flask & Session Settings ---
|
||||
# A strong, random secret key is crucial for session security.
|
||||
FLASK_SECRET_KEY=your-very-secret-and-random-key-here
|
||||
FLASK_HOST=127.0.0.1
|
||||
FLASK_PORT=5000
|
||||
FLASK_DEBUG=True
|
||||
# How long a user's session in the browser lasts (in hours).
|
||||
FLASK_PERMANENT_SESSION_LIFETIME_HOURS=2
|
||||
# How long inactive scanner data is stored in Redis (in minutes).
|
||||
SESSION_TIMEOUT_MINUTES=60
|
||||
|
||||
|
||||
# --- Application Core Settings ---
|
||||
# The default number of levels to recurse when scanning.
|
||||
DEFAULT_RECURSION_DEPTH=2
|
||||
# Default timeout for provider API requests in seconds.
|
||||
DEFAULT_TIMEOUT=30
|
||||
# The number of concurrent provider requests to make.
|
||||
MAX_CONCURRENT_REQUESTS=5
|
||||
# The number of results from a provider that triggers the "large entity" grouping.
|
||||
LARGE_ENTITY_THRESHOLD=100
|
||||
# The number of times to retry a target if a provider fails.
|
||||
MAX_RETRIES_PER_TARGET=3
|
||||
# How long cached provider responses are stored (in hours).
|
||||
CACHE_EXPIRY_HOURS=12
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -169,3 +169,4 @@ cython_debug/
|
||||
#.idea/
|
||||
|
||||
dump.rdb
|
||||
.vscode
|
||||
343
app.py
343
app.py
@@ -1,8 +1,6 @@
|
||||
# dnsrecon-reduced/app.py
|
||||
|
||||
"""
|
||||
Flask application entry point for DNSRecon web interface.
|
||||
Provides REST API endpoints and serves the web interface with user session support.
|
||||
Enhanced with user session management and task-based completion model.
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -11,42 +9,81 @@ from flask import Flask, render_template, request, jsonify, send_file, session
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import io
|
||||
|
||||
from core.session_manager import session_manager
|
||||
from core.session_manager import session_manager, UserIdentifier
|
||||
from config import config
|
||||
|
||||
|
||||
app = Flask(__name__)
|
||||
# Use centralized configuration for Flask settings
|
||||
app.config['SECRET_KEY'] = config.flask_secret_key
|
||||
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=config.flask_permanent_session_lifetime_hours)
|
||||
app.config['SECRET_KEY'] = 'dnsrecon-dev-key-change-in-production'
|
||||
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=2) # 2 hour session lifetime
|
||||
|
||||
|
||||
def get_user_scanner():
|
||||
"""
|
||||
Retrieves the scanner for the current session, or creates a new
|
||||
session and scanner if one doesn't exist.
|
||||
Enhanced user scanner retrieval with user identification and session consolidation.
|
||||
Implements single session per user with seamless consolidation.
|
||||
"""
|
||||
# Get current Flask session info for debugging
|
||||
current_flask_session_id = session.get('dnsrecon_session_id')
|
||||
print("=== ENHANCED GET_USER_SCANNER ===")
|
||||
|
||||
# Try to get existing session
|
||||
if current_flask_session_id:
|
||||
existing_scanner = session_manager.get_session(current_flask_session_id)
|
||||
if existing_scanner:
|
||||
return current_flask_session_id, existing_scanner
|
||||
|
||||
# Create new session if none exists
|
||||
print("Creating new session as none was found...")
|
||||
new_session_id = session_manager.create_session()
|
||||
new_scanner = session_manager.get_session(new_session_id)
|
||||
|
||||
if not new_scanner:
|
||||
raise Exception("Failed to create new scanner session")
|
||||
|
||||
# Store in Flask session
|
||||
session['dnsrecon_session_id'] = new_session_id
|
||||
session.permanent = True
|
||||
|
||||
return new_session_id, new_scanner
|
||||
try:
|
||||
# Extract user identification from request
|
||||
client_ip, user_agent = UserIdentifier.extract_request_info(request)
|
||||
user_fingerprint = UserIdentifier.generate_user_fingerprint(client_ip, user_agent)
|
||||
|
||||
print(f"User fingerprint: {user_fingerprint}")
|
||||
print(f"Client IP: {client_ip}")
|
||||
print(f"User Agent: {user_agent[:50]}...")
|
||||
|
||||
# Get current Flask session info for debugging
|
||||
current_flask_session_id = session.get('dnsrecon_session_id')
|
||||
print(f"Flask session ID: {current_flask_session_id}")
|
||||
|
||||
# Try to get existing session first
|
||||
if current_flask_session_id:
|
||||
existing_scanner = session_manager.get_session(current_flask_session_id)
|
||||
if existing_scanner:
|
||||
# Verify session belongs to current user
|
||||
session_info = session_manager.get_session_info(current_flask_session_id)
|
||||
if session_info.get('user_fingerprint') == user_fingerprint:
|
||||
print(f"Found valid existing session {current_flask_session_id} for user {user_fingerprint}")
|
||||
existing_scanner.session_id = current_flask_session_id
|
||||
return current_flask_session_id, existing_scanner
|
||||
else:
|
||||
print(f"Session {current_flask_session_id} belongs to different user, will create new session")
|
||||
else:
|
||||
print(f"Session {current_flask_session_id} not found in Redis, will create new session")
|
||||
|
||||
# Create or replace user session (this handles consolidation automatically)
|
||||
new_session_id = session_manager.create_or_replace_user_session(client_ip, user_agent)
|
||||
new_scanner = session_manager.get_session(new_session_id)
|
||||
|
||||
if not new_scanner:
|
||||
print(f"ERROR: Failed to retrieve newly created session {new_session_id}")
|
||||
raise Exception("Failed to create new scanner session")
|
||||
|
||||
# Store in Flask session for browser persistence
|
||||
session['dnsrecon_session_id'] = new_session_id
|
||||
session.permanent = True
|
||||
|
||||
# Ensure session ID is set on scanner
|
||||
new_scanner.session_id = new_session_id
|
||||
|
||||
# Get session info for user feedback
|
||||
session_info = session_manager.get_session_info(new_session_id)
|
||||
|
||||
print(f"Session created/consolidated successfully")
|
||||
print(f" - Session ID: {new_session_id}")
|
||||
print(f" - User: {user_fingerprint}")
|
||||
print(f" - Scanner status: {new_scanner.status}")
|
||||
print(f" - Session age: {session_info.get('session_age_minutes', 0)} minutes")
|
||||
|
||||
return new_session_id, new_scanner
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_user_scanner: {e}")
|
||||
traceback.print_exc()
|
||||
raise
|
||||
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
@@ -57,15 +94,21 @@ def index():
|
||||
@app.route('/api/scan/start', methods=['POST'])
|
||||
def start_scan():
|
||||
"""
|
||||
Start a new reconnaissance scan. Creates a new isolated scanner if
|
||||
clear_graph is true, otherwise adds to the existing one.
|
||||
Start a new reconnaissance scan with enhanced user session management.
|
||||
"""
|
||||
print("=== API: /api/scan/start called ===")
|
||||
|
||||
try:
|
||||
print("Getting JSON data from request...")
|
||||
data = request.get_json()
|
||||
print(f"Request data: {data}")
|
||||
|
||||
if not data or 'target_domain' not in data:
|
||||
return jsonify({'success': False, 'error': 'Missing target_domain in request'}), 400
|
||||
print("ERROR: Missing target_domain in request")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Missing target_domain in request'
|
||||
}), 400
|
||||
|
||||
target_domain = data['target_domain'].strip()
|
||||
max_depth = data.get('max_depth', config.default_recursion_depth)
|
||||
@@ -75,50 +118,86 @@ def start_scan():
|
||||
|
||||
# Validation
|
||||
if not target_domain:
|
||||
return jsonify({'success': False, 'error': 'Target domain cannot be empty'}), 400
|
||||
if not isinstance(max_depth, int) or not 1 <= max_depth <= 5:
|
||||
return jsonify({'success': False, 'error': 'Max depth must be an integer between 1 and 5'}), 400
|
||||
print("ERROR: Target domain cannot be empty")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Target domain cannot be empty'
|
||||
}), 400
|
||||
|
||||
user_session_id, scanner = None, None
|
||||
|
||||
if clear_graph:
|
||||
print("Clear graph requested: Creating a new, isolated scanner session.")
|
||||
old_session_id = session.get('dnsrecon_session_id')
|
||||
if old_session_id:
|
||||
session_manager.terminate_session(old_session_id)
|
||||
|
||||
user_session_id = session_manager.create_session()
|
||||
session['dnsrecon_session_id'] = user_session_id
|
||||
session.permanent = True
|
||||
scanner = session_manager.get_session(user_session_id)
|
||||
else:
|
||||
print("Adding to existing graph: Reusing the current scanner session.")
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if not scanner:
|
||||
return jsonify({'success': False, 'error': 'Failed to get or create a scanner instance.'}), 500
|
||||
if not isinstance(max_depth, int) or max_depth < 1 or max_depth > 5:
|
||||
print(f"ERROR: Invalid max_depth: {max_depth}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Max depth must be an integer between 1 and 5'
|
||||
}), 400
|
||||
|
||||
print(f"Using scanner {id(scanner)} in session {user_session_id}")
|
||||
print("Validation passed, getting user scanner...")
|
||||
|
||||
# Get user-specific scanner with enhanced session management
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
# Ensure session ID is properly set
|
||||
if not scanner.session_id:
|
||||
scanner.session_id = user_session_id
|
||||
|
||||
print(f"Using session: {user_session_id}")
|
||||
print(f"Scanner object ID: {id(scanner)}")
|
||||
|
||||
# Start scan
|
||||
print(f"Calling start_scan on scanner {id(scanner)}...")
|
||||
success = scanner.start_scan(target_domain, max_depth, clear_graph=clear_graph)
|
||||
|
||||
# Immediately update session state regardless of success
|
||||
session_manager.update_session_scanner(user_session_id, scanner)
|
||||
|
||||
if success:
|
||||
scan_session_id = scanner.logger.session_id
|
||||
print(f"Scan started successfully with scan session ID: {scan_session_id}")
|
||||
|
||||
# Get session info for user feedback
|
||||
session_info = session_manager.get_session_info(user_session_id)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': 'Scan started successfully',
|
||||
'scan_id': scanner.logger.session_id,
|
||||
'scan_id': scan_session_id,
|
||||
'user_session_id': user_session_id,
|
||||
'scanner_status': scanner.status,
|
||||
'session_info': {
|
||||
'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
|
||||
'session_age_minutes': session_info.get('session_age_minutes', 0),
|
||||
'consolidated': session_info.get('session_age_minutes', 0) > 0
|
||||
},
|
||||
'debug_info': {
|
||||
'scanner_object_id': id(scanner),
|
||||
'scanner_status': scanner.status
|
||||
}
|
||||
})
|
||||
else:
|
||||
print("ERROR: Scanner returned False")
|
||||
|
||||
# Provide more detailed error information
|
||||
error_details = {
|
||||
'scanner_status': scanner.status,
|
||||
'scanner_object_id': id(scanner),
|
||||
'session_id': user_session_id,
|
||||
'providers_count': len(scanner.providers) if hasattr(scanner, 'providers') else 0
|
||||
}
|
||||
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Failed to start scan (scanner status: {scanner.status})',
|
||||
'debug_info': error_details
|
||||
}), 409
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in start_scan endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({'success': False, 'error': f'Internal server error: {str(e)}'}), 500
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/scan/stop', methods=['POST'])
|
||||
def stop_scan():
|
||||
@@ -173,7 +252,7 @@ def stop_scan():
|
||||
|
||||
@app.route('/api/scan/status', methods=['GET'])
|
||||
def get_scan_status():
|
||||
"""Get current scan status with error handling."""
|
||||
"""Get current scan status with enhanced session information."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
@@ -204,6 +283,15 @@ def get_scan_status():
|
||||
status = scanner.get_scan_status()
|
||||
status['user_session_id'] = user_session_id
|
||||
|
||||
# Add enhanced session information
|
||||
session_info = session_manager.get_session_info(user_session_id)
|
||||
status['session_info'] = {
|
||||
'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
|
||||
'session_age_minutes': session_info.get('session_age_minutes', 0),
|
||||
'client_ip': session_info.get('client_ip', 'unknown'),
|
||||
'last_activity': session_info.get('last_activity')
|
||||
}
|
||||
|
||||
# Additional debug info
|
||||
status['debug_info'] = {
|
||||
'scanner_object_id': id(scanner),
|
||||
@@ -232,7 +320,6 @@ def get_scan_status():
|
||||
}), 500
|
||||
|
||||
|
||||
|
||||
@app.route('/api/graph', methods=['GET'])
|
||||
def get_graph_data():
|
||||
"""Get current graph data with error handling."""
|
||||
@@ -278,7 +365,6 @@ def get_graph_data():
|
||||
}), 500
|
||||
|
||||
|
||||
|
||||
@app.route('/api/export', methods=['GET'])
|
||||
def export_results():
|
||||
"""Export complete scan results as downloadable JSON for the user session."""
|
||||
@@ -289,17 +375,22 @@ def export_results():
|
||||
# Get complete results
|
||||
results = scanner.export_results()
|
||||
|
||||
# Add session information to export
|
||||
# Add enhanced session information to export
|
||||
session_info = session_manager.get_session_info(user_session_id)
|
||||
results['export_metadata'] = {
|
||||
'user_session_id': user_session_id,
|
||||
'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
|
||||
'client_ip': session_info.get('client_ip', 'unknown'),
|
||||
'session_age_minutes': session_info.get('session_age_minutes', 0),
|
||||
'export_timestamp': datetime.now(timezone.utc).isoformat(),
|
||||
'export_type': 'user_session_results'
|
||||
}
|
||||
|
||||
# Create filename with timestamp
|
||||
# Create filename with user fingerprint
|
||||
timestamp = datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')
|
||||
target = scanner.current_target or 'unknown'
|
||||
filename = f"dnsrecon_{target}_{timestamp}_{user_session_id[:8]}.json"
|
||||
user_fp = session_info.get('user_fingerprint', 'unknown')[:8]
|
||||
filename = f"dnsrecon_{target}_{timestamp}_{user_fp}.json"
|
||||
|
||||
# Create in-memory file
|
||||
json_data = json.dumps(results, indent=2, ensure_ascii=False)
|
||||
@@ -324,18 +415,12 @@ def export_results():
|
||||
@app.route('/api/providers', methods=['GET'])
|
||||
def get_providers():
|
||||
"""Get information about available providers for the user session."""
|
||||
print("=== API: /api/providers called ===")
|
||||
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if scanner:
|
||||
completed_tasks = scanner.indicators_completed
|
||||
enqueued_tasks = len(scanner.task_queue)
|
||||
print(f"DEBUG: Tasks - Completed: {completed_tasks}, Enqueued: {enqueued_tasks}")
|
||||
else:
|
||||
print("DEBUG: No active scanner session found.")
|
||||
|
||||
provider_info = scanner.get_provider_info()
|
||||
|
||||
return jsonify({
|
||||
@@ -410,6 +495,122 @@ def set_api_keys():
|
||||
'error': f'Internal server error: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/session/info', methods=['GET'])
|
||||
def get_session_info():
|
||||
"""Get enhanced information about the current user session."""
|
||||
try:
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
session_info = session_manager.get_session_info(user_session_id)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'session_info': session_info
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_session_info endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/session/terminate', methods=['POST'])
|
||||
def terminate_session():
|
||||
"""Terminate the current user session."""
|
||||
try:
|
||||
user_session_id = session.get('dnsrecon_session_id')
|
||||
|
||||
if user_session_id:
|
||||
success = session_manager.terminate_session(user_session_id)
|
||||
# Clear Flask session
|
||||
session.pop('dnsrecon_session_id', None)
|
||||
|
||||
return jsonify({
|
||||
'success': success,
|
||||
'message': 'Session terminated' if success else 'Session not found'
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'No active session to terminate'
|
||||
}), 400
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in terminate_session endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/admin/sessions', methods=['GET'])
|
||||
def list_sessions():
|
||||
"""Admin endpoint to list all active sessions with enhanced information."""
|
||||
try:
|
||||
sessions = session_manager.list_active_sessions()
|
||||
stats = session_manager.get_statistics()
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'sessions': sessions,
|
||||
'statistics': stats
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in list_sessions endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/health', methods=['GET'])
|
||||
def health_check():
|
||||
"""Health check endpoint with enhanced session statistics."""
|
||||
try:
|
||||
# Get session stats
|
||||
session_stats = session_manager.get_statistics()
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'status': 'healthy',
|
||||
'timestamp': datetime.now(timezone.utc).isoformat(),
|
||||
'version': '2.0.0-enhanced',
|
||||
'phase': 'enhanced_architecture',
|
||||
'features': {
|
||||
'multi_provider': True,
|
||||
'concurrent_processing': True,
|
||||
'real_time_updates': True,
|
||||
'api_key_management': True,
|
||||
'visualization': True,
|
||||
'retry_logic': True,
|
||||
'user_sessions': True,
|
||||
'session_isolation': True,
|
||||
'global_provider_caching': True,
|
||||
'single_session_per_user': True,
|
||||
'session_consolidation': True,
|
||||
'task_completion_model': True
|
||||
},
|
||||
'session_statistics': session_stats,
|
||||
'cache_info': {
|
||||
'global_provider_cache': True,
|
||||
'cache_location': '.cache/<provider_name>/',
|
||||
'cache_expiry_hours': 12
|
||||
}
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in health_check endpoint: {e}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Health check failed: {str(e)}'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.errorhandler(404)
|
||||
def not_found(error):
|
||||
"""Handle 404 errors."""
|
||||
@@ -431,7 +632,7 @@ def internal_error(error):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("Starting DNSRecon Flask application with user session support...")
|
||||
print("Starting DNSRecon Flask application with enhanced user session support...")
|
||||
|
||||
# Load configuration from environment
|
||||
config.load_from_env()
|
||||
|
||||
125
config.py
125
config.py
@@ -5,97 +5,110 @@ Handles API key storage, rate limiting, and default settings.
|
||||
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env file
|
||||
load_dotenv()
|
||||
|
||||
class Config:
|
||||
"""Configuration manager for DNSRecon application."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize configuration with default values."""
|
||||
self.api_keys: Dict[str, Optional[str]] = {}
|
||||
self.api_keys: Dict[str, Optional[str]] = {
|
||||
'shodan': None
|
||||
}
|
||||
|
||||
# --- General Settings ---
|
||||
# Default settings
|
||||
self.default_recursion_depth = 2
|
||||
self.default_timeout = 15
|
||||
self.default_timeout = 10
|
||||
self.max_concurrent_requests = 5
|
||||
self.large_entity_threshold = 100
|
||||
self.max_retries_per_target = 3
|
||||
self.cache_expiry_hours = 12
|
||||
|
||||
# --- Provider Caching Settings ---
|
||||
self.cache_timeout_hours = 6 # Provider-specific cache timeout
|
||||
|
||||
# --- Rate Limiting (requests per minute) ---
|
||||
# Rate limiting settings (requests per minute)
|
||||
self.rate_limits = {
|
||||
'crtsh': 30,
|
||||
'shodan': 60,
|
||||
'dns': 100
|
||||
'crtsh': 60, # Free service, be respectful
|
||||
'shodan': 60, # API dependent
|
||||
'dns': 100 # Local DNS queries
|
||||
}
|
||||
|
||||
# --- Provider Settings ---
|
||||
# Provider settings
|
||||
self.enabled_providers = {
|
||||
'crtsh': True,
|
||||
'dns': True,
|
||||
'shodan': False
|
||||
'crtsh': True, # Always enabled (free)
|
||||
'dns': True, # Always enabled (free)
|
||||
'shodan': False # Requires API key
|
||||
}
|
||||
|
||||
# --- Logging ---
|
||||
# Logging configuration
|
||||
self.log_level = 'INFO'
|
||||
self.log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
|
||||
# --- Flask & Session Settings ---
|
||||
# Flask configuration
|
||||
self.flask_host = '127.0.0.1'
|
||||
self.flask_port = 5000
|
||||
self.flask_debug = True
|
||||
self.flask_secret_key = 'default-secret-key-change-me'
|
||||
self.flask_permanent_session_lifetime_hours = 2
|
||||
self.session_timeout_minutes = 60
|
||||
|
||||
# Load environment variables to override defaults
|
||||
self.load_from_env()
|
||||
def set_api_key(self, provider: str, api_key: str) -> bool:
|
||||
"""
|
||||
Set API key for a provider.
|
||||
|
||||
def load_from_env(self):
|
||||
"""Load configuration from environment variables."""
|
||||
self.set_api_key('shodan', os.getenv('SHODAN_API_KEY'))
|
||||
|
||||
# Override settings from environment
|
||||
self.default_recursion_depth = int(os.getenv('DEFAULT_RECURSION_DEPTH', self.default_recursion_depth))
|
||||
self.default_timeout = int(os.getenv('DEFAULT_TIMEOUT', self.default_timeout))
|
||||
self.max_concurrent_requests = int(os.getenv('MAX_CONCURRENT_REQUESTS', self.max_concurrent_requests))
|
||||
self.large_entity_threshold = int(os.getenv('LARGE_ENTITY_THRESHOLD', self.large_entity_threshold))
|
||||
self.max_retries_per_target = int(os.getenv('MAX_RETRIES_PER_TARGET', self.max_retries_per_target))
|
||||
self.cache_expiry_hours = int(os.getenv('CACHE_EXPIRY_HOURS', self.cache_expiry_hours))
|
||||
self.cache_timeout_hours = int(os.getenv('CACHE_TIMEOUT_HOURS', self.cache_timeout_hours))
|
||||
|
||||
# Override Flask and session settings
|
||||
self.flask_host = os.getenv('FLASK_HOST', self.flask_host)
|
||||
self.flask_port = int(os.getenv('FLASK_PORT', self.flask_port))
|
||||
self.flask_debug = os.getenv('FLASK_DEBUG', str(self.flask_debug)).lower() == 'true'
|
||||
self.flask_secret_key = os.getenv('FLASK_SECRET_KEY', self.flask_secret_key)
|
||||
self.flask_permanent_session_lifetime_hours = int(os.getenv('FLASK_PERMANENT_SESSION_LIFETIME_HOURS', self.flask_permanent_session_lifetime_hours))
|
||||
self.session_timeout_minutes = int(os.getenv('SESSION_TIMEOUT_MINUTES', self.session_timeout_minutes))
|
||||
|
||||
def set_api_key(self, provider: str, api_key: Optional[str]) -> bool:
|
||||
"""Set API key for a provider."""
|
||||
self.api_keys[provider] = api_key
|
||||
if api_key:
|
||||
self.enabled_providers[provider] = True
|
||||
return True
|
||||
Args:
|
||||
provider: Provider name (shodan, etc)
|
||||
api_key: API key string
|
||||
|
||||
Returns:
|
||||
bool: True if key was set successfully
|
||||
"""
|
||||
if provider in self.api_keys:
|
||||
self.api_keys[provider] = api_key
|
||||
self.enabled_providers[provider] = True if api_key else False
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_api_key(self, provider: str) -> Optional[str]:
|
||||
"""Get API key for a provider."""
|
||||
"""
|
||||
Get API key for a provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
API key or None if not set
|
||||
"""
|
||||
return self.api_keys.get(provider)
|
||||
|
||||
def is_provider_enabled(self, provider: str) -> bool:
|
||||
"""Check if a provider is enabled."""
|
||||
"""
|
||||
Check if a provider is enabled.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
bool: True if provider is enabled
|
||||
"""
|
||||
return self.enabled_providers.get(provider, False)
|
||||
|
||||
def get_rate_limit(self, provider: str) -> int:
|
||||
"""Get rate limit for a provider."""
|
||||
"""
|
||||
Get rate limit for a provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
Rate limit in requests per minute
|
||||
"""
|
||||
return self.rate_limits.get(provider, 60)
|
||||
|
||||
def load_from_env(self):
|
||||
"""Load configuration from environment variables."""
|
||||
if os.getenv('SHODAN_API_KEY'):
|
||||
self.set_api_key('shodan', os.getenv('SHODAN_API_KEY'))
|
||||
|
||||
# Override default settings from environment
|
||||
self.default_recursion_depth = int(os.getenv('DEFAULT_RECURSION_DEPTH', '2'))
|
||||
self.flask_debug = os.getenv('FLASK_DEBUG', 'True').lower() == 'true'
|
||||
self.default_timeout = 30
|
||||
self.max_concurrent_requests = 5
|
||||
|
||||
|
||||
# Global configuration instance
|
||||
config = Config()
|
||||
@@ -8,6 +8,7 @@ from .scanner import Scanner, ScanStatus
|
||||
from .logger import ForensicLogger, get_forensic_logger, new_session
|
||||
from .session_manager import session_manager
|
||||
from .session_config import SessionConfig, create_session_config
|
||||
from .task_manager import TaskManager, TaskType, ReconTask
|
||||
|
||||
__all__ = [
|
||||
'GraphManager',
|
||||
@@ -19,7 +20,10 @@ __all__ = [
|
||||
'new_session',
|
||||
'session_manager',
|
||||
'SessionConfig',
|
||||
'create_session_config'
|
||||
'create_session_config',
|
||||
'TaskManager',
|
||||
'TaskType',
|
||||
'ReconTask'
|
||||
]
|
||||
|
||||
__version__ = "1.0.0-phase2"
|
||||
@@ -1,5 +1,3 @@
|
||||
# core/graph_manager.py
|
||||
|
||||
"""
|
||||
Graph data model for DNSRecon using NetworkX.
|
||||
Manages in-memory graph storage with confidence scoring and forensic metadata.
|
||||
@@ -52,23 +50,21 @@ class GraphManager:
|
||||
self.__dict__.update(state)
|
||||
self.date_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}')
|
||||
|
||||
def _update_correlation_index(self, node_id: str, data: Any, path: List[str] = [], parent_attr: str = ""):
|
||||
"""Recursively traverse metadata and add hashable values to the index with better path tracking."""
|
||||
def _update_correlation_index(self, node_id: str, data: Any, path: List[str] = None):
|
||||
"""Recursively traverse metadata and add hashable values to the index."""
|
||||
if path is None:
|
||||
path = []
|
||||
|
||||
if isinstance(data, dict):
|
||||
for key, value in data.items():
|
||||
self._update_correlation_index(node_id, value, path + [key], key)
|
||||
self._update_correlation_index(node_id, value, path + [key])
|
||||
elif isinstance(data, list):
|
||||
for i, item in enumerate(data):
|
||||
# Instead of just using [i], include the parent attribute context
|
||||
list_path_component = f"[{i}]" if not parent_attr else f"{parent_attr}[{i}]"
|
||||
self._update_correlation_index(node_id, item, path + [list_path_component], parent_attr)
|
||||
self._update_correlation_index(node_id, item, path + [f"[{i}]"])
|
||||
else:
|
||||
self._add_to_correlation_index(node_id, data, ".".join(path), parent_attr)
|
||||
self._add_to_correlation_index(node_id, data, ".".join(path))
|
||||
|
||||
def _add_to_correlation_index(self, node_id: str, value: Any, path_str: str, parent_attr: str = ""):
|
||||
def _add_to_correlation_index(self, node_id: str, value: Any, path_str: str):
|
||||
"""Add a hashable value to the correlation index, filtering out noise."""
|
||||
if not isinstance(value, (str, int, float, bool)) or value is None:
|
||||
return
|
||||
@@ -84,8 +80,8 @@ class GraphManager:
|
||||
return
|
||||
if len(value) < 4 or value.lower() in ['true', 'false', 'unknown', 'none', 'crt.sh']:
|
||||
return
|
||||
elif isinstance(value, int) and (abs(value) < 1024 or abs(value) > 65535):
|
||||
return # Ignore small integers and common port numbers
|
||||
elif isinstance(value, int) and abs(value) < 9999:
|
||||
return # Ignore small integers
|
||||
elif isinstance(value, bool):
|
||||
return # Ignore boolean values
|
||||
|
||||
@@ -94,47 +90,10 @@ class GraphManager:
|
||||
self.correlation_index[value] = {}
|
||||
if node_id not in self.correlation_index[value]:
|
||||
self.correlation_index[value][node_id] = []
|
||||
|
||||
# Store both the full path and the parent attribute for better edge labeling
|
||||
correlation_entry = {
|
||||
'path': path_str,
|
||||
'parent_attr': parent_attr,
|
||||
'meaningful_attr': self._extract_meaningful_attribute(path_str, parent_attr)
|
||||
}
|
||||
|
||||
if correlation_entry not in self.correlation_index[value][node_id]:
|
||||
self.correlation_index[value][node_id].append(correlation_entry)
|
||||
if path_str not in self.correlation_index[value][node_id]:
|
||||
self.correlation_index[value][node_id].append(path_str)
|
||||
|
||||
def _extract_meaningful_attribute(self, path_str: str, parent_attr: str = "") -> str:
|
||||
"""Extract the most meaningful attribute name from a path string."""
|
||||
if not path_str:
|
||||
return "unknown"
|
||||
|
||||
path_parts = path_str.split('.')
|
||||
|
||||
# Look for the last non-array-index part
|
||||
for part in reversed(path_parts):
|
||||
# Skip array indices like [0], [1], etc.
|
||||
if not (part.startswith('[') and part.endswith(']') and part[1:-1].isdigit()):
|
||||
# Clean up compound names like "hostnames[0]" to just "hostnames"
|
||||
clean_part = re.sub(r'\[\d+\]$', '', part)
|
||||
if clean_part:
|
||||
return clean_part
|
||||
|
||||
# Fallback to parent attribute if available
|
||||
if parent_attr:
|
||||
return parent_attr
|
||||
|
||||
# Last resort - use the first meaningful part
|
||||
for part in path_parts:
|
||||
if not (part.startswith('[') and part.endswith(']') and part[1:-1].isdigit()):
|
||||
clean_part = re.sub(r'\[\d+\]$', '', part)
|
||||
if clean_part:
|
||||
return clean_part
|
||||
|
||||
return "correlation"
|
||||
|
||||
def _check_for_correlations(self, new_node_id: str, data: Any, path: List[str] = [], parent_attr: str = "") -> List[Dict]:
|
||||
def _check_for_correlations(self, new_node_id: str, data: Any, path: List[str] = None) -> List[Dict]:
|
||||
"""Recursively traverse metadata to find correlations with existing data."""
|
||||
if path is None:
|
||||
path = []
|
||||
@@ -144,11 +103,10 @@ class GraphManager:
|
||||
for key, value in data.items():
|
||||
if key == 'source': # Avoid correlating on the provider name
|
||||
continue
|
||||
all_correlations.extend(self._check_for_correlations(new_node_id, value, path + [key], key))
|
||||
all_correlations.extend(self._check_for_correlations(new_node_id, value, path + [key]))
|
||||
elif isinstance(data, list):
|
||||
for i, item in enumerate(data):
|
||||
list_path_component = f"[{i}]" if not parent_attr else f"{parent_attr}[{i}]"
|
||||
all_correlations.extend(self._check_for_correlations(new_node_id, item, path + [list_path_component], parent_attr))
|
||||
all_correlations.extend(self._check_for_correlations(new_node_id, item, path + [f"[{i}]"]))
|
||||
else:
|
||||
value = data
|
||||
if value in self.correlation_index:
|
||||
@@ -159,31 +117,11 @@ class GraphManager:
|
||||
if len(unique_nodes) < 2:
|
||||
return all_correlations # Correlation must involve at least two distinct nodes
|
||||
|
||||
new_source = {
|
||||
'node_id': new_node_id,
|
||||
'path': ".".join(path),
|
||||
'parent_attr': parent_attr,
|
||||
'meaningful_attr': self._extract_meaningful_attribute(".".join(path), parent_attr)
|
||||
}
|
||||
new_source = {'node_id': new_node_id, 'path': ".".join(path)}
|
||||
all_sources = [new_source]
|
||||
|
||||
for node_id, path_entries in existing_nodes_with_paths.items():
|
||||
for entry in path_entries:
|
||||
if isinstance(entry, dict):
|
||||
all_sources.append({
|
||||
'node_id': node_id,
|
||||
'path': entry['path'],
|
||||
'parent_attr': entry.get('parent_attr', ''),
|
||||
'meaningful_attr': entry.get('meaningful_attr', self._extract_meaningful_attribute(entry['path'], entry.get('parent_attr', '')))
|
||||
})
|
||||
else:
|
||||
# Handle legacy string-only entries
|
||||
all_sources.append({
|
||||
'node_id': node_id,
|
||||
'path': str(entry),
|
||||
'parent_attr': '',
|
||||
'meaningful_attr': self._extract_meaningful_attribute(str(entry))
|
||||
})
|
||||
for node_id, paths in existing_nodes_with_paths.items():
|
||||
for p_str in paths:
|
||||
all_sources.append({'node_id': node_id, 'path': p_str})
|
||||
|
||||
all_correlations.append({
|
||||
'value': value,
|
||||
@@ -225,7 +163,8 @@ class GraphManager:
|
||||
# Skip creating correlation node - would be redundant
|
||||
continue
|
||||
|
||||
eligible_nodes = set(corr['nodes'])
|
||||
# STEP 2: Filter out node pairs that already have direct edges
|
||||
eligible_nodes = self._filter_nodes_without_direct_edges(set(corr['nodes']))
|
||||
|
||||
if len(eligible_nodes) < 2:
|
||||
# Need at least 2 nodes to create a correlation
|
||||
@@ -245,12 +184,11 @@ class GraphManager:
|
||||
metadata={'values': [value], 'sources': corr['sources'],
|
||||
'correlated_nodes': list(eligible_nodes)})
|
||||
|
||||
# Create edges from eligible nodes to this correlation node with better labeling
|
||||
# Create edges from eligible nodes to this correlation node
|
||||
for c_node_id in eligible_nodes:
|
||||
if self.graph.has_node(c_node_id):
|
||||
# Find the best attribute name for this node
|
||||
meaningful_attr = self._find_best_attribute_name_for_node(c_node_id, corr['sources'])
|
||||
relationship_type = f"c_{meaningful_attr}"
|
||||
attribute = corr['sources'][0]['path'].split('.')[-1]
|
||||
relationship_type = f"c_{attribute}"
|
||||
self.add_edge(c_node_id, correlation_node_id, relationship_type, confidence_score=0.9)
|
||||
|
||||
self._update_correlation_index(node_id, attributes)
|
||||
@@ -258,34 +196,27 @@ class GraphManager:
|
||||
self.last_modified = datetime.now(timezone.utc).isoformat()
|
||||
return is_new_node
|
||||
|
||||
def _find_best_attribute_name_for_node(self, node_id: str, sources: List[Dict]) -> str:
|
||||
"""Find the best attribute name for a correlation edge by looking at the sources."""
|
||||
node_sources = [s for s in sources if s['node_id'] == node_id]
|
||||
def _filter_nodes_without_direct_edges(self, node_set: set) -> set:
|
||||
"""
|
||||
Filter out nodes that already have direct edges between them.
|
||||
Returns set of nodes that should be included in correlation.
|
||||
"""
|
||||
nodes_list = list(node_set)
|
||||
eligible_nodes = set(node_set) # Start with all nodes
|
||||
|
||||
if not node_sources:
|
||||
return "correlation"
|
||||
# Check all pairs of nodes
|
||||
for i in range(len(nodes_list)):
|
||||
for j in range(i + 1, len(nodes_list)):
|
||||
node_a = nodes_list[i]
|
||||
node_b = nodes_list[j]
|
||||
|
||||
# Check if direct edge exists in either direction
|
||||
if self._has_direct_edge_bidirectional(node_a, node_b):
|
||||
# Remove both nodes from eligible set since they're already connected
|
||||
eligible_nodes.discard(node_a)
|
||||
eligible_nodes.discard(node_b)
|
||||
|
||||
# Use the meaningful_attr if available
|
||||
for source in node_sources:
|
||||
meaningful_attr = source.get('meaningful_attr')
|
||||
if meaningful_attr and meaningful_attr != "unknown":
|
||||
return meaningful_attr
|
||||
|
||||
# Fallback to parent_attr
|
||||
for source in node_sources:
|
||||
parent_attr = source.get('parent_attr')
|
||||
if parent_attr:
|
||||
return parent_attr
|
||||
|
||||
# Last resort - extract from path
|
||||
for source in node_sources:
|
||||
path = source.get('path', '')
|
||||
if path:
|
||||
extracted = self._extract_meaningful_attribute(path)
|
||||
if extracted != "unknown":
|
||||
return extracted
|
||||
|
||||
return "correlation"
|
||||
return eligible_nodes
|
||||
|
||||
def _has_direct_edge_bidirectional(self, node_a: str, node_b: str) -> bool:
|
||||
"""
|
||||
@@ -359,7 +290,7 @@ class GraphManager:
|
||||
# Create set of unique sources based on (node_id, path) tuples
|
||||
source_set = set()
|
||||
for source in existing_sources + new_sources:
|
||||
source_tuple = (source['node_id'], source.get('path', ''))
|
||||
source_tuple = (source['node_id'], source['path'])
|
||||
source_set.add(source_tuple)
|
||||
|
||||
# Convert back to list of dictionaries
|
||||
@@ -490,14 +421,10 @@ class GraphManager:
|
||||
def _get_confidence_distribution(self) -> Dict[str, int]:
|
||||
"""Get distribution of edge confidence scores."""
|
||||
distribution = {'high': 0, 'medium': 0, 'low': 0}
|
||||
for _, _, data in self.graph.edges(data=True):
|
||||
confidence = data.get('confidence_score', 0)
|
||||
if confidence >= 0.8:
|
||||
distribution['high'] += 1
|
||||
elif confidence >= 0.6:
|
||||
distribution['medium'] += 1
|
||||
else:
|
||||
distribution['low'] += 1
|
||||
for _, _, confidence in self.graph.edges(data='confidence_score', default=0):
|
||||
if confidence >= 0.8: distribution['high'] += 1
|
||||
elif confidence >= 0.6: distribution['medium'] += 1
|
||||
else: distribution['low'] += 1
|
||||
return distribution
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
@@ -512,10 +439,9 @@ class GraphManager:
|
||||
# Calculate distributions
|
||||
for node_type in NodeType:
|
||||
stats['node_type_distribution'][node_type.value] = self.get_nodes_by_type(node_type).__len__()
|
||||
for _, _, data in self.graph.edges(data=True):
|
||||
rel_type = data.get('relationship_type', 'unknown')
|
||||
for _, _, rel_type in self.graph.edges(data='relationship_type', default='unknown'):
|
||||
stats['relationship_type_distribution'][rel_type] = stats['relationship_type_distribution'].get(rel_type, 0) + 1
|
||||
provider = data.get('source_provider', 'unknown')
|
||||
for _, _, provider in self.graph.edges(data='source_provider', default='unknown'):
|
||||
stats['provider_distribution'][provider] = stats['provider_distribution'].get(provider, 0) + 1
|
||||
return stats
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ class ForensicLogger:
|
||||
Maintains detailed audit trail of all reconnaissance activities.
|
||||
"""
|
||||
|
||||
def __init__(self, session_id: str = ""):
|
||||
def __init__(self, session_id: str = None):
|
||||
"""
|
||||
Initialize forensic logger.
|
||||
|
||||
@@ -203,6 +203,8 @@ class ForensicLogger:
|
||||
self.session_metadata['target_domains'] = list(self.session_metadata['target_domains'])
|
||||
|
||||
self.logger.info(f"Scan Complete - Session: {self.session_id}")
|
||||
self.logger.info(f"Total API Requests: {self.session_metadata['total_requests']}")
|
||||
self.logger.info(f"Total Relationships: {self.session_metadata['total_relationships']}")
|
||||
|
||||
def export_audit_trail(self) -> Dict[str, Any]:
|
||||
"""
|
||||
|
||||
1105
core/scanner.py
1105
core/scanner.py
File diff suppressed because it is too large
Load Diff
@@ -1,20 +1,372 @@
|
||||
"""
|
||||
Per-session configuration management for DNSRecon.
|
||||
Provides isolated configuration instances for each user session.
|
||||
Enhanced per-session configuration management for DNSRecon.
|
||||
Provides isolated configuration instances for each user session while supporting global caching.
|
||||
"""
|
||||
|
||||
from config import Config
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
|
||||
class SessionConfig(Config):
|
||||
|
||||
class SessionConfig:
|
||||
"""
|
||||
Session-specific configuration that inherits from global config
|
||||
but maintains isolated API keys and provider settings.
|
||||
Enhanced session-specific configuration that inherits from global config
|
||||
but maintains isolated API keys and provider settings while supporting global caching.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize session config with global defaults."""
|
||||
super().__init__()
|
||||
"""Initialize enhanced session config with global cache support."""
|
||||
# Copy all attributes from global config
|
||||
self.api_keys: Dict[str, Optional[str]] = {
|
||||
'shodan': None
|
||||
}
|
||||
|
||||
# Default settings (copied from global config)
|
||||
self.default_recursion_depth = 2
|
||||
self.default_timeout = 30
|
||||
self.max_concurrent_requests = 5
|
||||
self.large_entity_threshold = 100
|
||||
|
||||
# Enhanced rate limiting settings (per session)
|
||||
self.rate_limits = {
|
||||
'crtsh': 60,
|
||||
'shodan': 60,
|
||||
'dns': 100
|
||||
}
|
||||
|
||||
# Enhanced provider settings (per session)
|
||||
self.enabled_providers = {
|
||||
'crtsh': True,
|
||||
'dns': True,
|
||||
'shodan': False
|
||||
}
|
||||
|
||||
# Task-based execution settings
|
||||
self.task_retry_settings = {
|
||||
'max_retries': 3,
|
||||
'base_backoff_seconds': 1.0,
|
||||
'max_backoff_seconds': 60.0,
|
||||
'retry_on_rate_limit': True,
|
||||
'retry_on_connection_error': True,
|
||||
'retry_on_timeout': True
|
||||
}
|
||||
|
||||
# Cache settings (global across all sessions)
|
||||
self.cache_settings = {
|
||||
'enabled': True,
|
||||
'expiry_hours': 12,
|
||||
'cache_base_dir': '.cache',
|
||||
'per_provider_directories': True,
|
||||
'thread_safe_operations': True
|
||||
}
|
||||
|
||||
# Logging configuration
|
||||
self.log_level = 'INFO'
|
||||
self.log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
|
||||
# Flask configuration (shared)
|
||||
self.flask_host = '127.0.0.1'
|
||||
self.flask_port = 5000
|
||||
self.flask_debug = True
|
||||
|
||||
# Session isolation settings
|
||||
self.session_isolation = {
|
||||
'enforce_single_session_per_user': True,
|
||||
'consolidate_session_data_on_replacement': True,
|
||||
'user_fingerprinting_enabled': True,
|
||||
'session_timeout_minutes': 60
|
||||
}
|
||||
|
||||
# Circuit breaker settings for provider reliability
|
||||
self.circuit_breaker = {
|
||||
'enabled': True,
|
||||
'failure_threshold': 5, # Failures before opening circuit
|
||||
'recovery_timeout_seconds': 300, # 5 minutes before trying again
|
||||
'half_open_max_calls': 3 # Test calls when recovering
|
||||
}
|
||||
|
||||
def set_api_key(self, provider: str, api_key: str) -> bool:
|
||||
"""
|
||||
Set API key for a provider in this session.
|
||||
|
||||
Args:
|
||||
provider: Provider name (shodan, etc)
|
||||
api_key: API key string (empty string to clear)
|
||||
|
||||
Returns:
|
||||
bool: True if key was set successfully
|
||||
"""
|
||||
if provider in self.api_keys:
|
||||
# Handle clearing of API keys
|
||||
if api_key and api_key.strip():
|
||||
self.api_keys[provider] = api_key.strip()
|
||||
self.enabled_providers[provider] = True
|
||||
else:
|
||||
self.api_keys[provider] = None
|
||||
self.enabled_providers[provider] = False
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_api_key(self, provider: str) -> Optional[str]:
|
||||
"""
|
||||
Get API key for a provider in this session.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
API key or None if not set
|
||||
"""
|
||||
return self.api_keys.get(provider)
|
||||
|
||||
def is_provider_enabled(self, provider: str) -> bool:
|
||||
"""
|
||||
Check if a provider is enabled in this session.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
bool: True if provider is enabled
|
||||
"""
|
||||
return self.enabled_providers.get(provider, False)
|
||||
|
||||
def get_rate_limit(self, provider: str) -> int:
|
||||
"""
|
||||
Get rate limit for a provider in this session.
|
||||
|
||||
Args:
|
||||
provider: Provider name
|
||||
|
||||
Returns:
|
||||
Rate limit in requests per minute
|
||||
"""
|
||||
return self.rate_limits.get(provider, 60)
|
||||
|
||||
def get_task_retry_config(self) -> Dict[str, any]:
|
||||
"""
|
||||
Get task retry configuration for this session.
|
||||
|
||||
Returns:
|
||||
Dictionary with retry settings
|
||||
"""
|
||||
return self.task_retry_settings.copy()
|
||||
|
||||
def get_cache_config(self) -> Dict[str, any]:
|
||||
"""
|
||||
Get cache configuration (global settings).
|
||||
|
||||
Returns:
|
||||
Dictionary with cache settings
|
||||
"""
|
||||
return self.cache_settings.copy()
|
||||
|
||||
def is_circuit_breaker_enabled(self) -> bool:
|
||||
"""Check if circuit breaker is enabled for provider reliability."""
|
||||
return self.circuit_breaker.get('enabled', True)
|
||||
|
||||
def get_circuit_breaker_config(self) -> Dict[str, any]:
|
||||
"""Get circuit breaker configuration."""
|
||||
return self.circuit_breaker.copy()
|
||||
|
||||
def update_provider_settings(self, provider_updates: Dict[str, Dict[str, any]]) -> bool:
|
||||
"""
|
||||
Update provider-specific settings in bulk.
|
||||
|
||||
Args:
|
||||
provider_updates: Dictionary of provider -> settings updates
|
||||
|
||||
Returns:
|
||||
bool: True if updates were applied successfully
|
||||
"""
|
||||
try:
|
||||
for provider_name, updates in provider_updates.items():
|
||||
# Update rate limits
|
||||
if 'rate_limit' in updates:
|
||||
self.rate_limits[provider_name] = updates['rate_limit']
|
||||
|
||||
# Update enabled status
|
||||
if 'enabled' in updates:
|
||||
self.enabled_providers[provider_name] = updates['enabled']
|
||||
|
||||
# Update API key
|
||||
if 'api_key' in updates:
|
||||
self.set_api_key(provider_name, updates['api_key'])
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error updating provider settings: {e}")
|
||||
return False
|
||||
|
||||
def validate_configuration(self) -> Dict[str, any]:
|
||||
"""
|
||||
Validate the current configuration and return validation results.
|
||||
|
||||
Returns:
|
||||
Dictionary with validation results and any issues found
|
||||
"""
|
||||
validation_result = {
|
||||
'valid': True,
|
||||
'warnings': [],
|
||||
'errors': [],
|
||||
'provider_status': {}
|
||||
}
|
||||
|
||||
# Validate provider configurations
|
||||
for provider_name, enabled in self.enabled_providers.items():
|
||||
provider_status = {
|
||||
'enabled': enabled,
|
||||
'has_api_key': bool(self.api_keys.get(provider_name)),
|
||||
'rate_limit': self.rate_limits.get(provider_name, 60)
|
||||
}
|
||||
|
||||
# Check for potential issues
|
||||
if enabled and provider_name in ['shodan'] and not provider_status['has_api_key']:
|
||||
validation_result['warnings'].append(
|
||||
f"Provider '{provider_name}' is enabled but missing API key"
|
||||
)
|
||||
|
||||
validation_result['provider_status'][provider_name] = provider_status
|
||||
|
||||
# Validate task settings
|
||||
if self.task_retry_settings['max_retries'] > 10:
|
||||
validation_result['warnings'].append(
|
||||
f"High retry count ({self.task_retry_settings['max_retries']}) may cause long delays"
|
||||
)
|
||||
|
||||
# Validate concurrent settings
|
||||
if self.max_concurrent_requests > 10:
|
||||
validation_result['warnings'].append(
|
||||
f"High concurrency ({self.max_concurrent_requests}) may overwhelm providers"
|
||||
)
|
||||
|
||||
# Validate cache settings
|
||||
if not os.path.exists(self.cache_settings['cache_base_dir']):
|
||||
try:
|
||||
os.makedirs(self.cache_settings['cache_base_dir'], exist_ok=True)
|
||||
except Exception as e:
|
||||
validation_result['errors'].append(f"Cannot create cache directory: {e}")
|
||||
validation_result['valid'] = False
|
||||
|
||||
return validation_result
|
||||
|
||||
def load_from_env(self):
|
||||
"""Load configuration from environment variables with enhanced validation."""
|
||||
# Load API keys from environment
|
||||
if os.getenv('SHODAN_API_KEY') and not self.api_keys['shodan']:
|
||||
self.set_api_key('shodan', os.getenv('SHODAN_API_KEY'))
|
||||
print("Loaded Shodan API key from environment")
|
||||
|
||||
# Override default settings from environment
|
||||
self.default_recursion_depth = int(os.getenv('DEFAULT_RECURSION_DEPTH', '2'))
|
||||
self.default_timeout = int(os.getenv('DEFAULT_TIMEOUT', '30'))
|
||||
self.max_concurrent_requests = int(os.getenv('MAX_CONCURRENT_REQUESTS', '5'))
|
||||
|
||||
# Load task retry settings from environment
|
||||
if os.getenv('TASK_MAX_RETRIES'):
|
||||
self.task_retry_settings['max_retries'] = int(os.getenv('TASK_MAX_RETRIES'))
|
||||
|
||||
if os.getenv('TASK_BASE_BACKOFF'):
|
||||
self.task_retry_settings['base_backoff_seconds'] = float(os.getenv('TASK_BASE_BACKOFF'))
|
||||
|
||||
# Load cache settings from environment
|
||||
if os.getenv('CACHE_EXPIRY_HOURS'):
|
||||
self.cache_settings['expiry_hours'] = int(os.getenv('CACHE_EXPIRY_HOURS'))
|
||||
|
||||
if os.getenv('CACHE_DISABLED'):
|
||||
self.cache_settings['enabled'] = os.getenv('CACHE_DISABLED').lower() != 'true'
|
||||
|
||||
# Load circuit breaker settings
|
||||
if os.getenv('CIRCUIT_BREAKER_DISABLED'):
|
||||
self.circuit_breaker['enabled'] = os.getenv('CIRCUIT_BREAKER_DISABLED').lower() != 'true'
|
||||
|
||||
# Flask settings
|
||||
self.flask_debug = os.getenv('FLASK_DEBUG', 'True').lower() == 'true'
|
||||
|
||||
print("Enhanced configuration loaded from environment")
|
||||
|
||||
def export_config_summary(self) -> Dict[str, any]:
|
||||
"""
|
||||
Export a summary of the current configuration for debugging/logging.
|
||||
|
||||
Returns:
|
||||
Dictionary with configuration summary (API keys redacted)
|
||||
"""
|
||||
return {
|
||||
'providers': {
|
||||
provider: {
|
||||
'enabled': self.enabled_providers.get(provider, False),
|
||||
'has_api_key': bool(self.api_keys.get(provider)),
|
||||
'rate_limit': self.rate_limits.get(provider, 60)
|
||||
}
|
||||
for provider in self.enabled_providers.keys()
|
||||
},
|
||||
'task_settings': {
|
||||
'max_retries': self.task_retry_settings['max_retries'],
|
||||
'max_concurrent_requests': self.max_concurrent_requests,
|
||||
'large_entity_threshold': self.large_entity_threshold
|
||||
},
|
||||
'cache_settings': {
|
||||
'enabled': self.cache_settings['enabled'],
|
||||
'expiry_hours': self.cache_settings['expiry_hours'],
|
||||
'base_directory': self.cache_settings['cache_base_dir']
|
||||
},
|
||||
'session_settings': {
|
||||
'isolation_enabled': self.session_isolation['enforce_single_session_per_user'],
|
||||
'consolidation_enabled': self.session_isolation['consolidate_session_data_on_replacement'],
|
||||
'timeout_minutes': self.session_isolation['session_timeout_minutes']
|
||||
},
|
||||
'circuit_breaker': {
|
||||
'enabled': self.circuit_breaker['enabled'],
|
||||
'failure_threshold': self.circuit_breaker['failure_threshold'],
|
||||
'recovery_timeout': self.circuit_breaker['recovery_timeout_seconds']
|
||||
}
|
||||
}
|
||||
|
||||
def create_session_config() -> 'SessionConfig':
|
||||
"""Create a new session configuration instance."""
|
||||
return SessionConfig()
|
||||
|
||||
def create_session_config() -> SessionConfig:
|
||||
"""
|
||||
Create a new enhanced session configuration instance.
|
||||
|
||||
Returns:
|
||||
Configured SessionConfig instance
|
||||
"""
|
||||
session_config = SessionConfig()
|
||||
session_config.load_from_env()
|
||||
|
||||
# Validate configuration and log any issues
|
||||
validation = session_config.validate_configuration()
|
||||
if validation['warnings']:
|
||||
print("Configuration warnings:")
|
||||
for warning in validation['warnings']:
|
||||
print(f" WARNING: {warning}")
|
||||
|
||||
if validation['errors']:
|
||||
print("Configuration errors:")
|
||||
for error in validation['errors']:
|
||||
print(f" ERROR: {error}")
|
||||
|
||||
if not validation['valid']:
|
||||
raise ValueError("Configuration validation failed - see errors above")
|
||||
|
||||
print(f"Enhanced session configuration created successfully")
|
||||
return session_config
|
||||
|
||||
|
||||
def create_test_config() -> SessionConfig:
|
||||
"""
|
||||
Create a test configuration with safe defaults for testing.
|
||||
|
||||
Returns:
|
||||
Test-safe SessionConfig instance
|
||||
"""
|
||||
test_config = SessionConfig()
|
||||
|
||||
# Override settings for testing
|
||||
test_config.max_concurrent_requests = 2
|
||||
test_config.task_retry_settings['max_retries'] = 1
|
||||
test_config.task_retry_settings['base_backoff_seconds'] = 0.1
|
||||
test_config.cache_settings['expiry_hours'] = 1
|
||||
test_config.session_isolation['session_timeout_minutes'] = 10
|
||||
|
||||
print("Test configuration created")
|
||||
return test_config
|
||||
@@ -5,41 +5,153 @@ import time
|
||||
import uuid
|
||||
import redis
|
||||
import pickle
|
||||
from typing import Dict, Optional, Any, List
|
||||
import hashlib
|
||||
from typing import Dict, Optional, Any, List, Tuple
|
||||
|
||||
from core.scanner import Scanner
|
||||
from config import config
|
||||
|
||||
# WARNING: Using pickle can be a security risk if the data source is not trusted.
|
||||
# In this case, we are only serializing/deserializing our own trusted Scanner objects,
|
||||
# which is generally safe. Do not unpickle data from untrusted sources.
|
||||
|
||||
class UserIdentifier:
|
||||
"""Handles user identification for session management."""
|
||||
|
||||
@staticmethod
|
||||
def generate_user_fingerprint(client_ip: str, user_agent: str) -> str:
|
||||
"""
|
||||
Generate a unique fingerprint for a user based on IP and User-Agent.
|
||||
|
||||
Args:
|
||||
client_ip: Client IP address
|
||||
user_agent: User-Agent header value
|
||||
|
||||
Returns:
|
||||
Unique user fingerprint hash
|
||||
"""
|
||||
# Create deterministic user identifier
|
||||
user_data = f"{client_ip}:{user_agent[:100]}" # Limit UA to 100 chars
|
||||
fingerprint = hashlib.sha256(user_data.encode()).hexdigest()[:16] # 16 char fingerprint
|
||||
return f"user_{fingerprint}"
|
||||
|
||||
@staticmethod
|
||||
def extract_request_info(request) -> Tuple[str, str]:
|
||||
"""
|
||||
Extract client IP and User-Agent from Flask request.
|
||||
|
||||
Args:
|
||||
request: Flask request object
|
||||
|
||||
Returns:
|
||||
Tuple of (client_ip, user_agent)
|
||||
"""
|
||||
# Handle proxy headers for real IP
|
||||
client_ip = request.headers.get('X-Forwarded-For', '').split(',')[0].strip()
|
||||
if not client_ip:
|
||||
client_ip = request.headers.get('X-Real-IP', '')
|
||||
if not client_ip:
|
||||
client_ip = request.remote_addr or 'unknown'
|
||||
|
||||
user_agent = request.headers.get('User-Agent', 'unknown')
|
||||
|
||||
return client_ip, user_agent
|
||||
|
||||
|
||||
class SessionConsolidator:
|
||||
"""Handles consolidation of session data when replacing sessions."""
|
||||
|
||||
@staticmethod
|
||||
def consolidate_scanner_data(old_scanner: 'Scanner', new_scanner: 'Scanner') -> 'Scanner':
|
||||
"""
|
||||
Consolidate useful data from old scanner into new scanner.
|
||||
|
||||
Args:
|
||||
old_scanner: Scanner from terminated session
|
||||
new_scanner: New scanner instance
|
||||
|
||||
Returns:
|
||||
Enhanced new scanner with consolidated data
|
||||
"""
|
||||
try:
|
||||
# Consolidate graph data if old scanner has valuable data
|
||||
if old_scanner and hasattr(old_scanner, 'graph') and old_scanner.graph:
|
||||
old_stats = old_scanner.graph.get_statistics()
|
||||
if old_stats['basic_metrics']['total_nodes'] > 0:
|
||||
print(f"Consolidating graph data: {old_stats['basic_metrics']['total_nodes']} nodes, {old_stats['basic_metrics']['total_edges']} edges")
|
||||
|
||||
# Transfer nodes and edges to new scanner's graph
|
||||
for node_id, node_data in old_scanner.graph.graph.nodes(data=True):
|
||||
# Add node to new graph with all attributes
|
||||
new_scanner.graph.graph.add_node(node_id, **node_data)
|
||||
|
||||
for source, target, edge_data in old_scanner.graph.graph.edges(data=True):
|
||||
# Add edge to new graph with all attributes
|
||||
new_scanner.graph.graph.add_edge(source, target, **edge_data)
|
||||
|
||||
# Update correlation index
|
||||
if hasattr(old_scanner.graph, 'correlation_index'):
|
||||
new_scanner.graph.correlation_index = old_scanner.graph.correlation_index.copy()
|
||||
|
||||
# Update timestamps
|
||||
new_scanner.graph.creation_time = old_scanner.graph.creation_time
|
||||
new_scanner.graph.last_modified = old_scanner.graph.last_modified
|
||||
|
||||
# Consolidate provider statistics
|
||||
if old_scanner and hasattr(old_scanner, 'providers') and old_scanner.providers:
|
||||
for old_provider in old_scanner.providers:
|
||||
# Find matching provider in new scanner
|
||||
matching_new_provider = None
|
||||
for new_provider in new_scanner.providers:
|
||||
if new_provider.get_name() == old_provider.get_name():
|
||||
matching_new_provider = new_provider
|
||||
break
|
||||
|
||||
if matching_new_provider:
|
||||
# Transfer cumulative statistics
|
||||
matching_new_provider.total_requests += old_provider.total_requests
|
||||
matching_new_provider.successful_requests += old_provider.successful_requests
|
||||
matching_new_provider.failed_requests += old_provider.failed_requests
|
||||
matching_new_provider.total_relationships_found += old_provider.total_relationships_found
|
||||
|
||||
# Transfer cache statistics if available
|
||||
if hasattr(old_provider, 'cache_hits'):
|
||||
matching_new_provider.cache_hits += getattr(old_provider, 'cache_hits', 0)
|
||||
matching_new_provider.cache_misses += getattr(old_provider, 'cache_misses', 0)
|
||||
|
||||
print(f"Consolidated {old_provider.get_name()} provider stats: {old_provider.total_requests} requests")
|
||||
|
||||
return new_scanner
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: Error during session consolidation: {e}")
|
||||
return new_scanner
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""
|
||||
Manages multiple scanner instances for concurrent user sessions using Redis.
|
||||
Manages single scanner session per user using Redis with user identification.
|
||||
Enforces one active session per user for consistent state management.
|
||||
"""
|
||||
|
||||
def __init__(self, session_timeout_minutes: int = 0):
|
||||
def __init__(self, session_timeout_minutes: int = 60):
|
||||
"""
|
||||
Initialize session manager with a Redis backend.
|
||||
Initialize session manager with Redis backend and user tracking.
|
||||
"""
|
||||
if session_timeout_minutes is None:
|
||||
session_timeout_minutes = config.session_timeout_minutes
|
||||
|
||||
self.redis_client = redis.StrictRedis(db=0, decode_responses=False)
|
||||
self.session_timeout = session_timeout_minutes * 60 # Convert to seconds
|
||||
self.lock = threading.Lock() # Lock for local operations, Redis handles atomic ops
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# User identification helper
|
||||
self.user_identifier = UserIdentifier()
|
||||
self.consolidator = SessionConsolidator()
|
||||
|
||||
# Start cleanup thread
|
||||
self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
|
||||
self.cleanup_thread.start()
|
||||
|
||||
print(f"SessionManager initialized with Redis backend and {session_timeout_minutes}min timeout")
|
||||
print(f"SessionManager initialized with Redis backend, user tracking, and {session_timeout_minutes}min timeout")
|
||||
|
||||
def __getstate__(self):
|
||||
"""Prepare SessionManager for pickling."""
|
||||
state = self.__dict__.copy()
|
||||
# Exclude unpickleable attributes - Redis client and threading objects
|
||||
# Exclude unpickleable attributes
|
||||
unpicklable_attrs = ['lock', 'cleanup_thread', 'redis_client']
|
||||
for attr in unpicklable_attrs:
|
||||
if attr in state:
|
||||
@@ -57,67 +169,108 @@ class SessionManager:
|
||||
self.cleanup_thread.start()
|
||||
|
||||
def _get_session_key(self, session_id: str) -> str:
|
||||
"""Generates the Redis key for a session."""
|
||||
"""Generate Redis key for a session."""
|
||||
return f"dnsrecon:session:{session_id}"
|
||||
|
||||
def _get_user_session_key(self, user_fingerprint: str) -> str:
|
||||
"""Generate Redis key for user -> session mapping."""
|
||||
return f"dnsrecon:user:{user_fingerprint}"
|
||||
|
||||
def _get_stop_signal_key(self, session_id: str) -> str:
|
||||
"""Generates the Redis key for a session's stop signal."""
|
||||
"""Generate Redis key for session stop signal."""
|
||||
return f"dnsrecon:stop:{session_id}"
|
||||
|
||||
def create_session(self) -> str:
|
||||
def create_or_replace_user_session(self, client_ip: str, user_agent: str) -> str:
|
||||
"""
|
||||
Create a new user session and store it in Redis.
|
||||
Create new session for user, replacing any existing session.
|
||||
Consolidates data from previous session if it exists.
|
||||
|
||||
Args:
|
||||
client_ip: Client IP address
|
||||
user_agent: User-Agent header
|
||||
|
||||
Returns:
|
||||
New session ID
|
||||
"""
|
||||
session_id = str(uuid.uuid4())
|
||||
print(f"=== CREATING SESSION {session_id} IN REDIS ===")
|
||||
user_fingerprint = self.user_identifier.generate_user_fingerprint(client_ip, user_agent)
|
||||
new_session_id = str(uuid.uuid4())
|
||||
|
||||
print(f"=== CREATING/REPLACING SESSION FOR USER {user_fingerprint} ===")
|
||||
|
||||
try:
|
||||
# Check for existing user session
|
||||
existing_session_id = self._get_user_current_session(user_fingerprint)
|
||||
old_scanner = None
|
||||
|
||||
if existing_session_id:
|
||||
print(f"Found existing session {existing_session_id} for user {user_fingerprint}")
|
||||
# Get old scanner data for consolidation
|
||||
old_scanner = self.get_session(existing_session_id)
|
||||
# Terminate old session
|
||||
self._terminate_session_internal(existing_session_id, cleanup_user_mapping=False)
|
||||
print(f"Terminated old session {existing_session_id}")
|
||||
|
||||
# Create new session config and scanner
|
||||
from core.session_config import create_session_config
|
||||
session_config = create_session_config()
|
||||
scanner_instance = Scanner(session_config=session_config)
|
||||
new_scanner = Scanner(session_config=session_config)
|
||||
|
||||
# Set the session ID on the scanner for cross-process stop signal management
|
||||
scanner_instance.session_id = session_id
|
||||
# Set session ID on scanner for cross-process operations
|
||||
new_scanner.session_id = new_session_id
|
||||
|
||||
# Consolidate data from old session if available
|
||||
if old_scanner:
|
||||
new_scanner = self.consolidator.consolidate_scanner_data(old_scanner, new_scanner)
|
||||
print(f"Consolidated data from previous session")
|
||||
|
||||
# Create session data
|
||||
session_data = {
|
||||
'scanner': scanner_instance,
|
||||
'scanner': new_scanner,
|
||||
'config': session_config,
|
||||
'created_at': time.time(),
|
||||
'last_activity': time.time(),
|
||||
'status': 'active'
|
||||
'status': 'active',
|
||||
'user_fingerprint': user_fingerprint,
|
||||
'client_ip': client_ip,
|
||||
'user_agent': user_agent[:200] # Truncate for storage
|
||||
}
|
||||
|
||||
# Serialize the entire session data dictionary using pickle
|
||||
# Store session in Redis
|
||||
session_key = self._get_session_key(new_session_id)
|
||||
serialized_data = pickle.dumps(session_data)
|
||||
|
||||
# Store in Redis
|
||||
session_key = self._get_session_key(session_id)
|
||||
self.redis_client.setex(session_key, self.session_timeout, serialized_data)
|
||||
|
||||
# Initialize stop signal as False
|
||||
stop_key = self._get_stop_signal_key(session_id)
|
||||
# Update user -> session mapping
|
||||
user_session_key = self._get_user_session_key(user_fingerprint)
|
||||
self.redis_client.setex(user_session_key, self.session_timeout, new_session_id.encode('utf-8'))
|
||||
|
||||
# Initialize stop signal
|
||||
stop_key = self._get_stop_signal_key(new_session_id)
|
||||
self.redis_client.setex(stop_key, self.session_timeout, b'0')
|
||||
|
||||
print(f"Session {session_id} stored in Redis with stop signal initialized")
|
||||
return session_id
|
||||
print(f"Created new session {new_session_id} for user {user_fingerprint}")
|
||||
return new_session_id
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to create session {session_id}: {e}")
|
||||
print(f"ERROR: Failed to create session for user {user_fingerprint}: {e}")
|
||||
raise
|
||||
|
||||
def _get_user_current_session(self, user_fingerprint: str) -> Optional[str]:
|
||||
"""Get current session ID for a user."""
|
||||
try:
|
||||
user_session_key = self._get_user_session_key(user_fingerprint)
|
||||
session_id_bytes = self.redis_client.get(user_session_key)
|
||||
if session_id_bytes:
|
||||
return session_id_bytes.decode('utf-8')
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error getting user session: {e}")
|
||||
return None
|
||||
|
||||
def set_stop_signal(self, session_id: str) -> bool:
|
||||
"""
|
||||
Set the stop signal for a session (cross-process safe).
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
bool: True if signal was set successfully
|
||||
"""
|
||||
"""Set stop signal for session (cross-process safe)."""
|
||||
try:
|
||||
stop_key = self._get_stop_signal_key(session_id)
|
||||
# Set stop signal to '1' with the same TTL as the session
|
||||
self.redis_client.setex(stop_key, self.session_timeout, b'1')
|
||||
print(f"Stop signal set for session {session_id}")
|
||||
return True
|
||||
@@ -126,15 +279,7 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def is_stop_requested(self, session_id: str) -> bool:
|
||||
"""
|
||||
Check if stop is requested for a session (cross-process safe).
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
bool: True if stop is requested
|
||||
"""
|
||||
"""Check if stop is requested for session (cross-process safe)."""
|
||||
try:
|
||||
stop_key = self._get_stop_signal_key(session_id)
|
||||
value = self.redis_client.get(stop_key)
|
||||
@@ -144,15 +289,7 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def clear_stop_signal(self, session_id: str) -> bool:
|
||||
"""
|
||||
Clear the stop signal for a session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
bool: True if signal was cleared successfully
|
||||
"""
|
||||
"""Clear stop signal for session."""
|
||||
try:
|
||||
stop_key = self._get_stop_signal_key(session_id)
|
||||
self.redis_client.setex(stop_key, self.session_timeout, b'0')
|
||||
@@ -163,13 +300,13 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def _get_session_data(self, session_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Retrieves and deserializes session data from Redis."""
|
||||
"""Retrieve and deserialize session data from Redis."""
|
||||
try:
|
||||
session_key = self._get_session_key(session_id)
|
||||
serialized_data = self.redis_client.get(session_key)
|
||||
if serialized_data:
|
||||
session_data = pickle.loads(serialized_data)
|
||||
# Ensure the scanner has the correct session ID for stop signal checking
|
||||
# Ensure scanner has correct session ID
|
||||
if 'scanner' in session_data and session_data['scanner']:
|
||||
session_data['scanner'].session_id = session_id
|
||||
return session_data
|
||||
@@ -179,37 +316,32 @@ class SessionManager:
|
||||
return None
|
||||
|
||||
def _save_session_data(self, session_id: str, session_data: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Serializes and saves session data back to Redis with updated TTL.
|
||||
|
||||
Returns:
|
||||
bool: True if save was successful
|
||||
"""
|
||||
"""Serialize and save session data to Redis with updated TTL."""
|
||||
try:
|
||||
session_key = self._get_session_key(session_id)
|
||||
serialized_data = pickle.dumps(session_data)
|
||||
result = self.redis_client.setex(session_key, self.session_timeout, serialized_data)
|
||||
|
||||
# Also refresh user mapping TTL if available
|
||||
if 'user_fingerprint' in session_data:
|
||||
user_session_key = self._get_user_session_key(session_data['user_fingerprint'])
|
||||
self.redis_client.setex(user_session_key, self.session_timeout, session_id.encode('utf-8'))
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to save session data for {session_id}: {e}")
|
||||
return False
|
||||
|
||||
def update_session_scanner(self, session_id: str, scanner: 'Scanner') -> bool:
|
||||
"""
|
||||
Updates just the scanner object in a session with immediate persistence.
|
||||
|
||||
Returns:
|
||||
bool: True if update was successful
|
||||
"""
|
||||
"""Update scanner object in session with immediate persistence."""
|
||||
try:
|
||||
session_data = self._get_session_data(session_id)
|
||||
if session_data:
|
||||
# Ensure scanner has the session ID
|
||||
# Ensure scanner has session ID
|
||||
scanner.session_id = session_id
|
||||
session_data['scanner'] = scanner
|
||||
session_data['last_activity'] = time.time()
|
||||
|
||||
# Immediately save to Redis for GUI updates
|
||||
success = self._save_session_data(session_id, session_data)
|
||||
if success:
|
||||
print(f"Scanner state updated for session {session_id} (status: {scanner.status})")
|
||||
@@ -224,16 +356,7 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def update_scanner_status(self, session_id: str, status: str) -> bool:
|
||||
"""
|
||||
Quickly update just the scanner status for immediate GUI feedback.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
status: New scanner status
|
||||
|
||||
Returns:
|
||||
bool: True if update was successful
|
||||
"""
|
||||
"""Quickly update scanner status for immediate GUI feedback."""
|
||||
try:
|
||||
session_data = self._get_session_data(session_id)
|
||||
if session_data and 'scanner' in session_data:
|
||||
@@ -252,9 +375,7 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def get_session(self, session_id: str) -> Optional[Scanner]:
|
||||
"""
|
||||
Get scanner instance for a session from Redis with session ID management.
|
||||
"""
|
||||
"""Get scanner instance for session with session ID management."""
|
||||
if not session_id:
|
||||
return None
|
||||
|
||||
@@ -269,21 +390,13 @@ class SessionManager:
|
||||
|
||||
scanner = session_data.get('scanner')
|
||||
if scanner:
|
||||
# Ensure the scanner can check the Redis-based stop signal
|
||||
# Ensure scanner can check Redis-based stop signal
|
||||
scanner.session_id = session_id
|
||||
|
||||
return scanner
|
||||
|
||||
def get_session_status_only(self, session_id: str) -> Optional[str]:
|
||||
"""
|
||||
Get just the scanner status without full session retrieval (for performance).
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
Scanner status string or None if not found
|
||||
"""
|
||||
"""Get scanner status without full session retrieval (for performance)."""
|
||||
try:
|
||||
session_data = self._get_session_data(session_id)
|
||||
if session_data and 'scanner' in session_data:
|
||||
@@ -294,16 +407,18 @@ class SessionManager:
|
||||
return None
|
||||
|
||||
def terminate_session(self, session_id: str) -> bool:
|
||||
"""
|
||||
Terminate a specific session in Redis with reliable stop signal and immediate status update.
|
||||
"""
|
||||
"""Terminate specific session with reliable stop signal and immediate status update."""
|
||||
return self._terminate_session_internal(session_id, cleanup_user_mapping=True)
|
||||
|
||||
def _terminate_session_internal(self, session_id: str, cleanup_user_mapping: bool = True) -> bool:
|
||||
"""Internal session termination with configurable user mapping cleanup."""
|
||||
print(f"=== TERMINATING SESSION {session_id} ===")
|
||||
|
||||
try:
|
||||
# First, set the stop signal
|
||||
# Set stop signal first
|
||||
self.set_stop_signal(session_id)
|
||||
|
||||
# Update scanner status to stopped immediately for GUI feedback
|
||||
# Update scanner status immediately for GUI feedback
|
||||
self.update_scanner_status(session_id, 'stopped')
|
||||
|
||||
session_data = self._get_session_data(session_id)
|
||||
@@ -314,16 +429,19 @@ class SessionManager:
|
||||
scanner = session_data.get('scanner')
|
||||
if scanner and scanner.status == 'running':
|
||||
print(f"Stopping scan for session: {session_id}")
|
||||
# The scanner will check the Redis stop signal
|
||||
scanner.stop_scan()
|
||||
|
||||
# Update the scanner state immediately
|
||||
self.update_session_scanner(session_id, scanner)
|
||||
|
||||
# Wait a moment for graceful shutdown
|
||||
# Wait for graceful shutdown
|
||||
time.sleep(0.5)
|
||||
|
||||
# Delete session data and stop signal from Redis
|
||||
# Clean up user mapping if requested
|
||||
if cleanup_user_mapping and 'user_fingerprint' in session_data:
|
||||
user_session_key = self._get_user_session_key(session_data['user_fingerprint'])
|
||||
self.redis_client.delete(user_session_key)
|
||||
print(f"Cleaned up user mapping for {session_data['user_fingerprint']}")
|
||||
|
||||
# Delete session data and stop signal
|
||||
session_key = self._get_session_key(session_id)
|
||||
stop_key = self._get_stop_signal_key(session_id)
|
||||
self.redis_client.delete(session_key)
|
||||
@@ -337,35 +455,72 @@ class SessionManager:
|
||||
return False
|
||||
|
||||
def _cleanup_loop(self) -> None:
|
||||
"""
|
||||
Background thread to cleanup inactive sessions and orphaned stop signals.
|
||||
"""
|
||||
"""Background thread to cleanup inactive sessions and orphaned signals."""
|
||||
while True:
|
||||
try:
|
||||
# Clean up orphaned stop signals
|
||||
stop_keys = self.redis_client.keys("dnsrecon:stop:*")
|
||||
for stop_key in stop_keys:
|
||||
# Extract session ID from stop key
|
||||
session_id = stop_key.decode('utf-8').split(':')[-1]
|
||||
session_key = self._get_session_key(session_id)
|
||||
|
||||
# If session doesn't exist but stop signal does, clean it up
|
||||
if not self.redis_client.exists(session_key):
|
||||
self.redis_client.delete(stop_key)
|
||||
print(f"Cleaned up orphaned stop signal for session {session_id}")
|
||||
|
||||
# Clean up orphaned user mappings
|
||||
user_keys = self.redis_client.keys("dnsrecon:user:*")
|
||||
for user_key in user_keys:
|
||||
session_id_bytes = self.redis_client.get(user_key)
|
||||
if session_id_bytes:
|
||||
session_id = session_id_bytes.decode('utf-8')
|
||||
session_key = self._get_session_key(session_id)
|
||||
|
||||
if not self.redis_client.exists(session_key):
|
||||
self.redis_client.delete(user_key)
|
||||
print(f"Cleaned up orphaned user mapping for session {session_id}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in cleanup loop: {e}")
|
||||
|
||||
time.sleep(300) # Sleep for 5 minutes
|
||||
|
||||
def list_active_sessions(self) -> List[Dict[str, Any]]:
|
||||
"""List all active sessions for admin purposes."""
|
||||
try:
|
||||
session_keys = self.redis_client.keys("dnsrecon:session:*")
|
||||
sessions = []
|
||||
|
||||
for session_key in session_keys:
|
||||
session_id = session_key.decode('utf-8').split(':')[-1]
|
||||
session_data = self._get_session_data(session_id)
|
||||
|
||||
if session_data:
|
||||
scanner = session_data.get('scanner')
|
||||
sessions.append({
|
||||
'session_id': session_id,
|
||||
'user_fingerprint': session_data.get('user_fingerprint', 'unknown'),
|
||||
'client_ip': session_data.get('client_ip', 'unknown'),
|
||||
'created_at': session_data.get('created_at'),
|
||||
'last_activity': session_data.get('last_activity'),
|
||||
'scanner_status': scanner.status if scanner else 'unknown',
|
||||
'current_target': scanner.current_target if scanner else None
|
||||
})
|
||||
|
||||
return sessions
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to list active sessions: {e}")
|
||||
return []
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get session manager statistics."""
|
||||
try:
|
||||
session_keys = self.redis_client.keys("dnsrecon:session:*")
|
||||
user_keys = self.redis_client.keys("dnsrecon:user:*")
|
||||
stop_keys = self.redis_client.keys("dnsrecon:stop:*")
|
||||
|
||||
active_sessions = len(session_keys)
|
||||
unique_users = len(user_keys)
|
||||
running_scans = 0
|
||||
|
||||
for session_key in session_keys:
|
||||
@@ -376,16 +531,46 @@ class SessionManager:
|
||||
|
||||
return {
|
||||
'total_active_sessions': active_sessions,
|
||||
'unique_users': unique_users,
|
||||
'running_scans': running_scans,
|
||||
'total_stop_signals': len(stop_keys)
|
||||
'total_stop_signals': len(stop_keys),
|
||||
'average_sessions_per_user': round(active_sessions / unique_users, 2) if unique_users > 0 else 0
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to get statistics: {e}")
|
||||
return {
|
||||
'total_active_sessions': 0,
|
||||
'unique_users': 0,
|
||||
'running_scans': 0,
|
||||
'total_stop_signals': 0
|
||||
'total_stop_signals': 0,
|
||||
'average_sessions_per_user': 0
|
||||
}
|
||||
|
||||
def get_session_info(self, session_id: str) -> Dict[str, Any]:
|
||||
"""Get detailed information about a specific session."""
|
||||
try:
|
||||
session_data = self._get_session_data(session_id)
|
||||
if not session_data:
|
||||
return {'error': 'Session not found'}
|
||||
|
||||
scanner = session_data.get('scanner')
|
||||
|
||||
return {
|
||||
'session_id': session_id,
|
||||
'user_fingerprint': session_data.get('user_fingerprint', 'unknown'),
|
||||
'client_ip': session_data.get('client_ip', 'unknown'),
|
||||
'user_agent': session_data.get('user_agent', 'unknown'),
|
||||
'created_at': session_data.get('created_at'),
|
||||
'last_activity': session_data.get('last_activity'),
|
||||
'status': session_data.get('status'),
|
||||
'scanner_status': scanner.status if scanner else 'unknown',
|
||||
'current_target': scanner.current_target if scanner else None,
|
||||
'session_age_minutes': round((time.time() - session_data.get('created_at', time.time())) / 60, 1)
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to get session info for {session_id}: {e}")
|
||||
return {'error': f'Failed to get session info: {str(e)}'}
|
||||
|
||||
|
||||
# Global session manager instance
|
||||
session_manager = SessionManager(session_timeout_minutes=60)
|
||||
564
core/task_manager.py
Normal file
564
core/task_manager.py
Normal file
@@ -0,0 +1,564 @@
|
||||
# dnsrecon/core/task_manager.py
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Optional, Any, Set
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from collections import deque
|
||||
|
||||
from utils.helpers import _is_valid_ip, _is_valid_domain
|
||||
|
||||
|
||||
class TaskStatus(Enum):
|
||||
"""Enumeration of task execution statuses."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
SUCCEEDED = "succeeded"
|
||||
FAILED_RETRYING = "failed_retrying"
|
||||
FAILED_PERMANENT = "failed_permanent"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class TaskType(Enum):
|
||||
"""Enumeration of task types for provider queries."""
|
||||
DOMAIN_QUERY = "domain_query"
|
||||
IP_QUERY = "ip_query"
|
||||
GRAPH_UPDATE = "graph_update"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskResult:
|
||||
"""Result of a task execution."""
|
||||
success: bool
|
||||
data: Optional[Any] = None
|
||||
error: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReconTask:
|
||||
"""Represents a single reconnaissance task with retry logic."""
|
||||
task_id: str
|
||||
task_type: TaskType
|
||||
target: str
|
||||
provider_name: str
|
||||
depth: int
|
||||
status: TaskStatus = TaskStatus.PENDING
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Retry configuration
|
||||
max_retries: int = 3
|
||||
current_retry: int = 0
|
||||
base_backoff_seconds: float = 1.0
|
||||
max_backoff_seconds: float = 60.0
|
||||
|
||||
# Execution tracking
|
||||
last_attempt_at: Optional[datetime] = None
|
||||
next_retry_at: Optional[datetime] = None
|
||||
execution_history: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
# Results
|
||||
result: Optional[TaskResult] = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize additional fields after creation."""
|
||||
if not self.task_id:
|
||||
self.task_id = str(uuid.uuid4())[:8]
|
||||
|
||||
def calculate_next_retry_time(self) -> datetime:
|
||||
"""Calculate next retry time with exponential backoff and jitter."""
|
||||
if self.current_retry >= self.max_retries:
|
||||
return None
|
||||
|
||||
# Exponential backoff with jitter
|
||||
backoff_time = min(
|
||||
self.max_backoff_seconds,
|
||||
self.base_backoff_seconds * (2 ** self.current_retry)
|
||||
)
|
||||
|
||||
# Add jitter (±25%)
|
||||
jitter = backoff_time * 0.25 * (0.5 - hash(self.task_id) % 1000 / 1000.0)
|
||||
final_backoff = max(self.base_backoff_seconds, backoff_time + jitter)
|
||||
|
||||
return datetime.now(timezone.utc) + timedelta(seconds=final_backoff)
|
||||
|
||||
def should_retry(self) -> bool:
|
||||
"""Determine if task should be retried based on status and retry count."""
|
||||
if self.status != TaskStatus.FAILED_RETRYING:
|
||||
return False
|
||||
if self.current_retry >= self.max_retries:
|
||||
return False
|
||||
if self.next_retry_at and datetime.now(timezone.utc) < self.next_retry_at:
|
||||
return False
|
||||
return True
|
||||
|
||||
def mark_failed(self, error: str, metadata: Dict[str, Any] = None):
|
||||
"""Mark task as failed and prepare for retry or permanent failure."""
|
||||
self.current_retry += 1
|
||||
self.last_attempt_at = datetime.now(timezone.utc)
|
||||
|
||||
# Record execution history
|
||||
execution_record = {
|
||||
'attempt': self.current_retry,
|
||||
'timestamp': self.last_attempt_at.isoformat(),
|
||||
'error': error,
|
||||
'metadata': metadata or {}
|
||||
}
|
||||
self.execution_history.append(execution_record)
|
||||
|
||||
if self.current_retry >= self.max_retries:
|
||||
self.status = TaskStatus.FAILED_PERMANENT
|
||||
self.result = TaskResult(success=False, error=f"Permanent failure after {self.max_retries} attempts: {error}")
|
||||
else:
|
||||
self.status = TaskStatus.FAILED_RETRYING
|
||||
self.next_retry_at = self.calculate_next_retry_time()
|
||||
|
||||
def mark_succeeded(self, data: Any = None, metadata: Dict[str, Any] = None):
|
||||
"""Mark task as successfully completed."""
|
||||
self.status = TaskStatus.SUCCEEDED
|
||||
self.last_attempt_at = datetime.now(timezone.utc)
|
||||
self.result = TaskResult(success=True, data=data, metadata=metadata or {})
|
||||
|
||||
# Record successful execution
|
||||
execution_record = {
|
||||
'attempt': self.current_retry + 1,
|
||||
'timestamp': self.last_attempt_at.isoformat(),
|
||||
'success': True,
|
||||
'metadata': metadata or {}
|
||||
}
|
||||
self.execution_history.append(execution_record)
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
"""Get task summary for progress reporting."""
|
||||
return {
|
||||
'task_id': self.task_id,
|
||||
'task_type': self.task_type.value,
|
||||
'target': self.target,
|
||||
'provider': self.provider_name,
|
||||
'status': self.status.value,
|
||||
'current_retry': self.current_retry,
|
||||
'max_retries': self.max_retries,
|
||||
'created_at': self.created_at.isoformat(),
|
||||
'last_attempt_at': self.last_attempt_at.isoformat() if self.last_attempt_at else None,
|
||||
'next_retry_at': self.next_retry_at.isoformat() if self.next_retry_at else None,
|
||||
'total_attempts': len(self.execution_history),
|
||||
'has_result': self.result is not None
|
||||
}
|
||||
|
||||
|
||||
class TaskQueue:
|
||||
"""Thread-safe task queue with retry logic and priority handling."""
|
||||
|
||||
def __init__(self, max_concurrent_tasks: int = 5):
|
||||
"""Initialize task queue."""
|
||||
self.max_concurrent_tasks = max_concurrent_tasks
|
||||
self.tasks: Dict[str, ReconTask] = {}
|
||||
self.pending_queue = deque()
|
||||
self.retry_queue = deque()
|
||||
self.running_tasks: Set[str] = set()
|
||||
|
||||
self._lock = threading.Lock()
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
def __getstate__(self):
|
||||
"""Prepare TaskQueue for pickling by excluding unpicklable objects."""
|
||||
state = self.__dict__.copy()
|
||||
# Exclude the unpickleable '_lock' and '_stop_event' attributes
|
||||
if '_lock' in state:
|
||||
del state['_lock']
|
||||
if '_stop_event' in state:
|
||||
del state['_stop_event']
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""Restore TaskQueue after unpickling by reconstructing threading objects."""
|
||||
self.__dict__.update(state)
|
||||
# Re-initialize the '_lock' and '_stop_event' attributes
|
||||
self._lock = threading.Lock()
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
def add_task(self, task: ReconTask) -> str:
|
||||
"""Add task to queue."""
|
||||
with self._lock:
|
||||
self.tasks[task.task_id] = task
|
||||
self.pending_queue.append(task.task_id)
|
||||
print(f"Added task {task.task_id}: {task.provider_name} query for {task.target}")
|
||||
return task.task_id
|
||||
|
||||
def get_next_ready_task(self) -> Optional[ReconTask]:
|
||||
"""Get next task ready for execution."""
|
||||
with self._lock:
|
||||
# Check if we have room for more concurrent tasks
|
||||
if len(self.running_tasks) >= self.max_concurrent_tasks:
|
||||
return None
|
||||
|
||||
# First priority: retry queue (tasks ready for retry)
|
||||
while self.retry_queue:
|
||||
task_id = self.retry_queue.popleft()
|
||||
if task_id in self.tasks:
|
||||
task = self.tasks[task_id]
|
||||
if task.should_retry():
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.running_tasks.add(task_id)
|
||||
print(f"Retrying task {task_id} (attempt {task.current_retry + 1})")
|
||||
return task
|
||||
|
||||
# Second priority: pending queue (new tasks)
|
||||
while self.pending_queue:
|
||||
task_id = self.pending_queue.popleft()
|
||||
if task_id in self.tasks:
|
||||
task = self.tasks[task_id]
|
||||
if task.status == TaskStatus.PENDING:
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.running_tasks.add(task_id)
|
||||
print(f"Starting task {task_id}")
|
||||
return task
|
||||
|
||||
return None
|
||||
|
||||
def complete_task(self, task_id: str, success: bool, data: Any = None,
|
||||
error: str = None, metadata: Dict[str, Any] = None):
|
||||
"""Mark task as completed (success or failure)."""
|
||||
with self._lock:
|
||||
if task_id not in self.tasks:
|
||||
return
|
||||
|
||||
task = self.tasks[task_id]
|
||||
self.running_tasks.discard(task_id)
|
||||
|
||||
if success:
|
||||
task.mark_succeeded(data=data, metadata=metadata)
|
||||
print(f"Task {task_id} succeeded")
|
||||
else:
|
||||
task.mark_failed(error or "Unknown error", metadata=metadata)
|
||||
if task.status == TaskStatus.FAILED_RETRYING:
|
||||
self.retry_queue.append(task_id)
|
||||
print(f"Task {task_id} failed, scheduled for retry at {task.next_retry_at}")
|
||||
else:
|
||||
print(f"Task {task_id} permanently failed after {task.current_retry} attempts")
|
||||
|
||||
def cancel_all_tasks(self):
|
||||
"""Cancel all pending and running tasks."""
|
||||
with self._lock:
|
||||
self._stop_event.set()
|
||||
for task in self.tasks.values():
|
||||
if task.status in [TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.FAILED_RETRYING]:
|
||||
task.status = TaskStatus.CANCELLED
|
||||
self.pending_queue.clear()
|
||||
self.retry_queue.clear()
|
||||
self.running_tasks.clear()
|
||||
print("All tasks cancelled")
|
||||
|
||||
def is_complete(self) -> bool:
|
||||
"""Check if all tasks are complete (succeeded, permanently failed, or cancelled)."""
|
||||
with self._lock:
|
||||
for task in self.tasks.values():
|
||||
if task.status in [TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.FAILED_RETRYING]:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get queue statistics."""
|
||||
with self._lock:
|
||||
stats = {
|
||||
'total_tasks': len(self.tasks),
|
||||
'pending': len(self.pending_queue),
|
||||
'running': len(self.running_tasks),
|
||||
'retry_queue': len(self.retry_queue),
|
||||
'succeeded': 0,
|
||||
'failed_permanent': 0,
|
||||
'cancelled': 0,
|
||||
'failed_retrying': 0
|
||||
}
|
||||
|
||||
for task in self.tasks.values():
|
||||
if task.status == TaskStatus.SUCCEEDED:
|
||||
stats['succeeded'] += 1
|
||||
elif task.status == TaskStatus.FAILED_PERMANENT:
|
||||
stats['failed_permanent'] += 1
|
||||
elif task.status == TaskStatus.CANCELLED:
|
||||
stats['cancelled'] += 1
|
||||
elif task.status == TaskStatus.FAILED_RETRYING:
|
||||
stats['failed_retrying'] += 1
|
||||
|
||||
stats['completion_rate'] = (stats['succeeded'] / stats['total_tasks'] * 100) if stats['total_tasks'] > 0 else 0
|
||||
stats['is_complete'] = self.is_complete()
|
||||
|
||||
return stats
|
||||
|
||||
def get_task_summaries(self) -> List[Dict[str, Any]]:
|
||||
"""Get summaries of all tasks for detailed progress reporting."""
|
||||
with self._lock:
|
||||
return [task.get_summary() for task in self.tasks.values()]
|
||||
|
||||
def get_failed_tasks(self) -> List[ReconTask]:
|
||||
"""Get all permanently failed tasks for analysis."""
|
||||
with self._lock:
|
||||
return [task for task in self.tasks.values() if task.status == TaskStatus.FAILED_PERMANENT]
|
||||
|
||||
|
||||
class TaskExecutor:
|
||||
"""Executes reconnaissance tasks using providers."""
|
||||
|
||||
def __init__(self, providers: List, graph_manager, logger):
|
||||
"""Initialize task executor."""
|
||||
self.providers = {provider.get_name(): provider for provider in providers}
|
||||
self.graph = graph_manager
|
||||
self.logger = logger
|
||||
|
||||
def execute_task(self, task: ReconTask) -> TaskResult:
|
||||
"""
|
||||
Execute a single reconnaissance task.
|
||||
|
||||
Args:
|
||||
task: Task to execute
|
||||
|
||||
Returns:
|
||||
TaskResult with success/failure information
|
||||
"""
|
||||
try:
|
||||
print(f"Executing task {task.task_id}: {task.provider_name} query for {task.target}")
|
||||
|
||||
provider = self.providers.get(task.provider_name)
|
||||
if not provider:
|
||||
return TaskResult(
|
||||
success=False,
|
||||
error=f"Provider {task.provider_name} not available"
|
||||
)
|
||||
|
||||
if not provider.is_available():
|
||||
return TaskResult(
|
||||
success=False,
|
||||
error=f"Provider {task.provider_name} is not available (missing API key or configuration)"
|
||||
)
|
||||
|
||||
# Execute provider query based on task type
|
||||
if task.task_type == TaskType.DOMAIN_QUERY:
|
||||
if not _is_valid_domain(task.target):
|
||||
return TaskResult(success=False, error=f"Invalid domain: {task.target}")
|
||||
|
||||
relationships = provider.query_domain(task.target)
|
||||
|
||||
elif task.task_type == TaskType.IP_QUERY:
|
||||
if not _is_valid_ip(task.target):
|
||||
return TaskResult(success=False, error=f"Invalid IP: {task.target}")
|
||||
|
||||
relationships = provider.query_ip(task.target)
|
||||
|
||||
else:
|
||||
return TaskResult(success=False, error=f"Unsupported task type: {task.task_type}")
|
||||
|
||||
# Process results and update graph
|
||||
new_targets = set()
|
||||
relationships_added = 0
|
||||
|
||||
for source, target, rel_type, confidence, raw_data in relationships:
|
||||
# Add nodes to graph
|
||||
from core.graph_manager import NodeType
|
||||
|
||||
if _is_valid_ip(target):
|
||||
self.graph.add_node(target, NodeType.IP)
|
||||
new_targets.add(target)
|
||||
elif target.startswith('AS') and target[2:].isdigit():
|
||||
self.graph.add_node(target, NodeType.ASN)
|
||||
elif _is_valid_domain(target):
|
||||
self.graph.add_node(target, NodeType.DOMAIN)
|
||||
new_targets.add(target)
|
||||
|
||||
# Add edge to graph
|
||||
if self.graph.add_edge(source, target, rel_type, confidence, task.provider_name, raw_data):
|
||||
relationships_added += 1
|
||||
|
||||
# Log forensic information
|
||||
self.logger.logger.info(
|
||||
f"Task {task.task_id} completed: {len(relationships)} relationships found, "
|
||||
f"{relationships_added} added to graph, {len(new_targets)} new targets"
|
||||
)
|
||||
|
||||
return TaskResult(
|
||||
success=True,
|
||||
data={
|
||||
'relationships': relationships,
|
||||
'new_targets': list(new_targets),
|
||||
'relationships_added': relationships_added
|
||||
},
|
||||
metadata={
|
||||
'provider': task.provider_name,
|
||||
'target': task.target,
|
||||
'depth': task.depth,
|
||||
'execution_time': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Task execution failed: {str(e)}"
|
||||
print(f"ERROR: {error_msg} for task {task.task_id}")
|
||||
self.logger.logger.error(error_msg)
|
||||
|
||||
return TaskResult(
|
||||
success=False,
|
||||
error=error_msg,
|
||||
metadata={
|
||||
'provider': task.provider_name,
|
||||
'target': task.target,
|
||||
'exception_type': type(e).__name__
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TaskManager:
|
||||
"""High-level task management for reconnaissance scans."""
|
||||
|
||||
def __init__(self, providers: List, graph_manager, logger, max_concurrent_tasks: int = 5):
|
||||
"""Initialize task manager."""
|
||||
self.task_queue = TaskQueue(max_concurrent_tasks)
|
||||
self.task_executor = TaskExecutor(providers, graph_manager, logger)
|
||||
self.logger = logger
|
||||
|
||||
# Execution control
|
||||
self._stop_event = threading.Event()
|
||||
self._execution_threads: List[threading.Thread] = []
|
||||
self._is_running = False
|
||||
|
||||
def create_provider_tasks(self, target: str, depth: int, providers: List) -> List[str]:
|
||||
"""
|
||||
Create tasks for querying all eligible providers for a target.
|
||||
|
||||
Args:
|
||||
target: Domain or IP to query
|
||||
depth: Current recursion depth
|
||||
providers: List of available providers
|
||||
|
||||
Returns:
|
||||
List of created task IDs
|
||||
"""
|
||||
task_ids = []
|
||||
is_ip = _is_valid_ip(target)
|
||||
target_key = 'ips' if is_ip else 'domains'
|
||||
task_type = TaskType.IP_QUERY if is_ip else TaskType.DOMAIN_QUERY
|
||||
|
||||
for provider in providers:
|
||||
if provider.get_eligibility().get(target_key) and provider.is_available():
|
||||
task = ReconTask(
|
||||
task_id=str(uuid.uuid4())[:8],
|
||||
task_type=task_type,
|
||||
target=target,
|
||||
provider_name=provider.get_name(),
|
||||
depth=depth,
|
||||
max_retries=3 # Configure retries per task type/provider
|
||||
)
|
||||
|
||||
task_id = self.task_queue.add_task(task)
|
||||
task_ids.append(task_id)
|
||||
|
||||
return task_ids
|
||||
|
||||
def start_execution(self, max_workers: int = 3):
|
||||
"""Start task execution with specified number of worker threads."""
|
||||
if self._is_running:
|
||||
print("Task execution already running")
|
||||
return
|
||||
|
||||
self._is_running = True
|
||||
self._stop_event.clear()
|
||||
|
||||
print(f"Starting task execution with {max_workers} workers")
|
||||
|
||||
for i in range(max_workers):
|
||||
worker_thread = threading.Thread(
|
||||
target=self._worker_loop,
|
||||
name=f"TaskWorker-{i+1}",
|
||||
daemon=True
|
||||
)
|
||||
worker_thread.start()
|
||||
self._execution_threads.append(worker_thread)
|
||||
|
||||
def stop_execution(self):
|
||||
"""Stop task execution and cancel all tasks."""
|
||||
print("Stopping task execution")
|
||||
self._stop_event.set()
|
||||
self.task_queue.cancel_all_tasks()
|
||||
self._is_running = False
|
||||
|
||||
# Wait for worker threads to finish
|
||||
for thread in self._execution_threads:
|
||||
thread.join(timeout=5.0)
|
||||
|
||||
self._execution_threads.clear()
|
||||
print("Task execution stopped")
|
||||
|
||||
def _worker_loop(self):
|
||||
"""Worker thread loop for executing tasks."""
|
||||
thread_name = threading.current_thread().name
|
||||
print(f"{thread_name} started")
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
# Get next task to execute
|
||||
task = self.task_queue.get_next_ready_task()
|
||||
|
||||
if task is None:
|
||||
# No tasks ready, check if we should exit
|
||||
if self.task_queue.is_complete() or self._stop_event.is_set():
|
||||
break
|
||||
time.sleep(0.1) # Brief sleep before checking again
|
||||
continue
|
||||
|
||||
# Execute the task
|
||||
result = self.task_executor.execute_task(task)
|
||||
|
||||
# Complete the task in queue
|
||||
self.task_queue.complete_task(
|
||||
task.task_id,
|
||||
success=result.success,
|
||||
data=result.data,
|
||||
error=result.error,
|
||||
metadata=result.metadata
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Worker {thread_name} encountered error: {e}")
|
||||
# Continue running even if individual task fails
|
||||
continue
|
||||
|
||||
print(f"{thread_name} finished")
|
||||
|
||||
def wait_for_completion(self, timeout_seconds: int = 300) -> bool:
|
||||
"""
|
||||
Wait for all tasks to complete.
|
||||
|
||||
Args:
|
||||
timeout_seconds: Maximum time to wait
|
||||
|
||||
Returns:
|
||||
True if all tasks completed, False if timeout
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < timeout_seconds:
|
||||
if self.task_queue.is_complete():
|
||||
return True
|
||||
|
||||
if self._stop_event.is_set():
|
||||
return False
|
||||
|
||||
time.sleep(1.0) # Check every second
|
||||
|
||||
print(f"Timeout waiting for task completion after {timeout_seconds} seconds")
|
||||
return False
|
||||
|
||||
def get_progress_report(self) -> Dict[str, Any]:
|
||||
"""Get detailed progress report for UI updates."""
|
||||
stats = self.task_queue.get_statistics()
|
||||
failed_tasks = self.task_queue.get_failed_tasks()
|
||||
|
||||
return {
|
||||
'statistics': stats,
|
||||
'failed_tasks': [task.get_summary() for task in failed_tasks],
|
||||
'is_running': self._is_running,
|
||||
'worker_count': len(self._execution_threads),
|
||||
'detailed_tasks': self.task_queue.get_task_summaries() if stats['total_tasks'] < 50 else [] # Limit detail for performance
|
||||
}
|
||||
@@ -3,14 +3,18 @@
|
||||
import time
|
||||
import requests
|
||||
import threading
|
||||
import os
|
||||
import json
|
||||
import hashlib
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from core.logger import get_forensic_logger
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""Simple rate limiter for API calls."""
|
||||
"""Thread-safe rate limiter for API calls."""
|
||||
|
||||
def __init__(self, requests_per_minute: int):
|
||||
"""
|
||||
@@ -22,36 +26,152 @@ class RateLimiter:
|
||||
self.requests_per_minute = requests_per_minute
|
||||
self.min_interval = 60.0 / requests_per_minute
|
||||
self.last_request_time = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def __getstate__(self):
|
||||
"""RateLimiter is fully picklable, return full state."""
|
||||
return self.__dict__.copy()
|
||||
state = self.__dict__.copy()
|
||||
# Exclude unpickleable lock
|
||||
if '_lock' in state:
|
||||
del state['_lock']
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""Restore RateLimiter state."""
|
||||
self.__dict__.update(state)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def wait_if_needed(self) -> None:
|
||||
"""Wait if necessary to respect rate limits."""
|
||||
current_time = time.time()
|
||||
time_since_last = current_time - self.last_request_time
|
||||
with self._lock:
|
||||
current_time = time.time()
|
||||
time_since_last = current_time - self.last_request_time
|
||||
|
||||
if time_since_last < self.min_interval:
|
||||
sleep_time = self.min_interval - time_since_last
|
||||
time.sleep(sleep_time)
|
||||
if time_since_last < self.min_interval:
|
||||
sleep_time = self.min_interval - time_since_last
|
||||
time.sleep(sleep_time)
|
||||
|
||||
self.last_request_time = time.time()
|
||||
self.last_request_time = time.time()
|
||||
|
||||
|
||||
class ProviderCache:
|
||||
"""Thread-safe global cache for provider queries."""
|
||||
|
||||
def __init__(self, provider_name: str, cache_expiry_hours: int = 12):
|
||||
"""
|
||||
Initialize provider-specific cache.
|
||||
|
||||
Args:
|
||||
provider_name: Name of the provider for cache directory
|
||||
cache_expiry_hours: Cache expiry time in hours
|
||||
"""
|
||||
self.provider_name = provider_name
|
||||
self.cache_expiry = cache_expiry_hours * 3600 # Convert to seconds
|
||||
self.cache_dir = os.path.join('.cache', provider_name)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Ensure cache directory exists with thread-safe creation
|
||||
os.makedirs(self.cache_dir, exist_ok=True)
|
||||
|
||||
def _generate_cache_key(self, method: str, url: str, params: Optional[Dict[str, Any]]) -> str:
|
||||
"""Generate unique cache key for request."""
|
||||
cache_data = f"{method}:{url}:{json.dumps(params or {}, sort_keys=True)}"
|
||||
return hashlib.md5(cache_data.encode()).hexdigest() + ".json"
|
||||
|
||||
def get_cached_response(self, method: str, url: str, params: Optional[Dict[str, Any]]) -> Optional[requests.Response]:
|
||||
"""
|
||||
Retrieve cached response if available and not expired.
|
||||
|
||||
Returns:
|
||||
Cached Response object or None if cache miss/expired
|
||||
"""
|
||||
cache_key = self._generate_cache_key(method, url, params)
|
||||
cache_path = os.path.join(self.cache_dir, cache_key)
|
||||
|
||||
with self._lock:
|
||||
if not os.path.exists(cache_path):
|
||||
return None
|
||||
|
||||
# Check if cache is expired
|
||||
cache_age = time.time() - os.path.getmtime(cache_path)
|
||||
if cache_age >= self.cache_expiry:
|
||||
try:
|
||||
os.remove(cache_path)
|
||||
except OSError:
|
||||
pass # File might have been removed by another thread
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(cache_path, 'r', encoding='utf-8') as f:
|
||||
cached_data = json.load(f)
|
||||
|
||||
# Reconstruct Response object
|
||||
response = requests.Response()
|
||||
response.status_code = cached_data['status_code']
|
||||
response._content = cached_data['content'].encode('utf-8')
|
||||
response.headers.update(cached_data['headers'])
|
||||
|
||||
return response
|
||||
|
||||
except (json.JSONDecodeError, KeyError, IOError) as e:
|
||||
# Cache file corrupted, remove it
|
||||
try:
|
||||
os.remove(cache_path)
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
def cache_response(self, method: str, url: str, params: Optional[Dict[str, Any]],
|
||||
response: requests.Response) -> bool:
|
||||
"""
|
||||
Cache successful response to disk.
|
||||
|
||||
Returns:
|
||||
True if cached successfully, False otherwise
|
||||
"""
|
||||
if response.status_code != 200:
|
||||
return False
|
||||
|
||||
cache_key = self._generate_cache_key(method, url, params)
|
||||
cache_path = os.path.join(self.cache_dir, cache_key)
|
||||
|
||||
with self._lock:
|
||||
try:
|
||||
cache_data = {
|
||||
'status_code': response.status_code,
|
||||
'content': response.text,
|
||||
'headers': dict(response.headers),
|
||||
'cached_at': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
|
||||
# Write to temporary file first, then rename for atomic operation
|
||||
temp_path = cache_path + '.tmp'
|
||||
with open(temp_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(cache_data, f)
|
||||
|
||||
# Atomic rename to prevent partial cache files
|
||||
os.rename(temp_path, cache_path)
|
||||
return True
|
||||
|
||||
except (IOError, OSError) as e:
|
||||
# Clean up temp file if it exists
|
||||
try:
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
class BaseProvider(ABC):
|
||||
"""
|
||||
Abstract base class for all DNSRecon data providers.
|
||||
Now supports session-specific configuration.
|
||||
Now supports global provider-specific caching and session-specific configuration.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, rate_limit: int = 60, timeout: int = 30, session_config=None):
|
||||
"""
|
||||
Initialize base provider with session-specific configuration.
|
||||
Initialize base provider with global caching and session-specific configuration.
|
||||
|
||||
Args:
|
||||
name: Provider name for logging
|
||||
@@ -78,20 +198,25 @@ class BaseProvider(ABC):
|
||||
self.logger = get_forensic_logger()
|
||||
self._stop_event = None
|
||||
|
||||
# GLOBAL provider-specific caching (not session-based)
|
||||
self.cache = ProviderCache(name, cache_expiry_hours=12)
|
||||
|
||||
# Statistics (per provider instance)
|
||||
self.total_requests = 0
|
||||
self.successful_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.total_relationships_found = 0
|
||||
self.cache_hits = 0
|
||||
self.cache_misses = 0
|
||||
|
||||
print(f"Initialized {name} provider with global cache and session config (rate: {actual_rate_limit}/min)")
|
||||
|
||||
def __getstate__(self):
|
||||
"""Prepare BaseProvider for pickling by excluding unpicklable objects."""
|
||||
state = self.__dict__.copy()
|
||||
# Exclude the unpickleable '_local' attribute and stop event
|
||||
unpicklable_attrs = ['_local', '_stop_event']
|
||||
for attr in unpicklable_attrs:
|
||||
if attr in state:
|
||||
del state[attr]
|
||||
state['_local'] = None
|
||||
state['_stop_event'] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
@@ -106,7 +231,7 @@ class BaseProvider(ABC):
|
||||
if not hasattr(self._local, 'session'):
|
||||
self._local.session = requests.Session()
|
||||
self._local.session.headers.update({
|
||||
'User-Agent': 'DNSRecon/1.0 (Passive Reconnaissance Tool)'
|
||||
'User-Agent': 'DNSRecon/2.0 (Passive Reconnaissance Tool)'
|
||||
})
|
||||
return self._local.session
|
||||
|
||||
@@ -164,79 +289,153 @@ class BaseProvider(ABC):
|
||||
def make_request(self, url: str, method: str = "GET",
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
target_indicator: str = "") -> Optional[requests.Response]:
|
||||
target_indicator: str = "",
|
||||
max_retries: int = 3) -> Optional[requests.Response]:
|
||||
"""
|
||||
Make a rate-limited HTTP request.
|
||||
Make a rate-limited HTTP request with global caching and aggressive stop signal handling.
|
||||
"""
|
||||
# Check for cancellation before starting
|
||||
if self._is_stop_requested():
|
||||
print(f"Request cancelled before start: {url}")
|
||||
return None
|
||||
|
||||
self.rate_limiter.wait_if_needed()
|
||||
# Check global cache first
|
||||
cached_response = self.cache.get_cached_response(method, url, params)
|
||||
if cached_response is not None:
|
||||
print(f"Cache hit for {self.name}: {url}")
|
||||
self.cache_hits += 1
|
||||
return cached_response
|
||||
|
||||
self.cache_misses += 1
|
||||
|
||||
start_time = time.time()
|
||||
response = None
|
||||
error = None
|
||||
# Determine effective max_retries based on stop signal
|
||||
effective_max_retries = 0 if self._is_stop_requested() else max_retries
|
||||
last_exception = None
|
||||
|
||||
try:
|
||||
self.total_requests += 1
|
||||
for attempt in range(effective_max_retries + 1):
|
||||
# Check for cancellation before each attempt
|
||||
if self._is_stop_requested():
|
||||
print(f"Request cancelled during attempt {attempt + 1}: {url}")
|
||||
return None
|
||||
|
||||
request_headers = dict(self.session.headers).copy()
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
# Apply rate limiting with cancellation awareness
|
||||
if not self._wait_with_cancellation_check():
|
||||
print(f"Request cancelled during rate limiting: {url}")
|
||||
return None
|
||||
|
||||
print(f"Making {method} request to: {url}")
|
||||
# Final check before making HTTP request
|
||||
if self._is_stop_requested():
|
||||
print(f"Request cancelled before HTTP call: {url}")
|
||||
return None
|
||||
|
||||
if method.upper() == "GET":
|
||||
response = self.session.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=request_headers,
|
||||
timeout=self.timeout
|
||||
start_time = time.time()
|
||||
response = None
|
||||
error = None
|
||||
|
||||
try:
|
||||
self.total_requests += 1
|
||||
|
||||
# Prepare request
|
||||
request_headers = self.session.headers.copy()
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
|
||||
print(f"Making {method} request to: {url} (attempt {attempt + 1})")
|
||||
|
||||
# Use shorter timeout if termination is requested
|
||||
request_timeout = 2 if self._is_stop_requested() else self.timeout
|
||||
|
||||
# Make request
|
||||
if method.upper() == "GET":
|
||||
response = self.session.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=request_headers,
|
||||
timeout=request_timeout
|
||||
)
|
||||
elif method.upper() == "POST":
|
||||
response = self.session.post(
|
||||
url,
|
||||
json=params,
|
||||
headers=request_headers,
|
||||
timeout=request_timeout
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported HTTP method: {method}")
|
||||
|
||||
print(f"Response status: {response.status_code}")
|
||||
response.raise_for_status()
|
||||
self.successful_requests += 1
|
||||
|
||||
# Success - log, cache, and return
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
self.logger.log_api_request(
|
||||
provider=self.name,
|
||||
url=url,
|
||||
method=method.upper(),
|
||||
status_code=response.status_code,
|
||||
response_size=len(response.content),
|
||||
duration_ms=duration_ms,
|
||||
error=None,
|
||||
target_indicator=target_indicator
|
||||
)
|
||||
elif method.upper() == "POST":
|
||||
response = self.session.post(
|
||||
url,
|
||||
json=params,
|
||||
headers=request_headers,
|
||||
timeout=self.timeout
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported HTTP method: {method}")
|
||||
|
||||
# Cache the successful response globally
|
||||
self.cache.cache_response(method, url, params, response)
|
||||
return response
|
||||
|
||||
print(f"Response status: {response.status_code}")
|
||||
response.raise_for_status()
|
||||
self.successful_requests += 1
|
||||
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
self.logger.log_api_request(
|
||||
provider=self.name,
|
||||
url=url,
|
||||
method=method.upper(),
|
||||
status_code=response.status_code,
|
||||
response_size=len(response.content),
|
||||
duration_ms=duration_ms,
|
||||
error=None,
|
||||
target_indicator=target_indicator
|
||||
)
|
||||
|
||||
return response
|
||||
except requests.exceptions.RequestException as e:
|
||||
error = str(e)
|
||||
self.failed_requests += 1
|
||||
print(f"Request failed (attempt {attempt + 1}): {error}")
|
||||
last_exception = e
|
||||
|
||||
# Immediately abort retries if stop requested
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested - aborting retries for: {url}")
|
||||
break
|
||||
|
||||
# Check if we should retry
|
||||
if attempt < effective_max_retries and self._should_retry(e):
|
||||
# Exponential backoff with jitter for 429 errors
|
||||
if isinstance(e, requests.exceptions.HTTPError) and e.response and e.response.status_code == 429:
|
||||
backoff_time = min(60, 10 * (2 ** attempt))
|
||||
print(f"Rate limit hit. Retrying in {backoff_time} seconds...")
|
||||
else:
|
||||
backoff_time = min(2.0, (2 ** attempt) * 0.5)
|
||||
print(f"Retrying in {backoff_time} seconds...")
|
||||
|
||||
if not self._sleep_with_cancellation_check(backoff_time):
|
||||
print(f"Stop requested during backoff - aborting: {url}")
|
||||
return None
|
||||
continue
|
||||
else:
|
||||
break
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
error = str(e)
|
||||
self.failed_requests += 1
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
self.logger.log_api_request(
|
||||
provider=self.name,
|
||||
url=url,
|
||||
method=method.upper(),
|
||||
status_code=response.status_code if response else None,
|
||||
response_size=len(response.content) if response else None,
|
||||
duration_ms=duration_ms,
|
||||
error=error,
|
||||
target_indicator=target_indicator
|
||||
)
|
||||
raise e
|
||||
except Exception as e:
|
||||
error = f"Unexpected error: {str(e)}"
|
||||
self.failed_requests += 1
|
||||
print(f"Unexpected error: {error}")
|
||||
last_exception = e
|
||||
break
|
||||
|
||||
# All attempts failed - log and return None
|
||||
duration_ms = (time.time() - start_time) * 1000
|
||||
self.logger.log_api_request(
|
||||
provider=self.name,
|
||||
url=url,
|
||||
method=method.upper(),
|
||||
status_code=response.status_code if response else None,
|
||||
response_size=len(response.content) if response else None,
|
||||
duration_ms=duration_ms,
|
||||
error=error,
|
||||
target_indicator=target_indicator
|
||||
)
|
||||
|
||||
if error and last_exception:
|
||||
raise last_exception
|
||||
|
||||
return None
|
||||
|
||||
def _is_stop_requested(self) -> bool:
|
||||
"""
|
||||
@@ -246,6 +445,43 @@ class BaseProvider(ABC):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _wait_with_cancellation_check(self) -> bool:
|
||||
"""
|
||||
Wait for rate limiting while aggressively checking for cancellation.
|
||||
Returns False if cancelled during wait.
|
||||
"""
|
||||
current_time = time.time()
|
||||
time_since_last = current_time - self.rate_limiter.last_request_time
|
||||
|
||||
if time_since_last < self.rate_limiter.min_interval:
|
||||
sleep_time = self.rate_limiter.min_interval - time_since_last
|
||||
if not self._sleep_with_cancellation_check(sleep_time):
|
||||
return False
|
||||
|
||||
self.rate_limiter.last_request_time = time.time()
|
||||
return True
|
||||
|
||||
def _sleep_with_cancellation_check(self, sleep_time: float) -> bool:
|
||||
"""
|
||||
Sleep for the specified time while aggressively checking for cancellation.
|
||||
|
||||
Args:
|
||||
sleep_time: Time to sleep in seconds
|
||||
|
||||
Returns:
|
||||
bool: True if sleep completed, False if cancelled
|
||||
"""
|
||||
sleep_start = time.time()
|
||||
check_interval = 0.05 # Check every 50ms for aggressive responsiveness
|
||||
|
||||
while time.time() - sleep_start < sleep_time:
|
||||
if self._is_stop_requested():
|
||||
return False
|
||||
remaining_time = sleep_time - (time.time() - sleep_start)
|
||||
time.sleep(min(check_interval, remaining_time))
|
||||
|
||||
return True
|
||||
|
||||
def set_stop_event(self, stop_event: threading.Event) -> None:
|
||||
"""
|
||||
Set the stop event for this provider to enable cancellation.
|
||||
@@ -255,6 +491,28 @@ class BaseProvider(ABC):
|
||||
"""
|
||||
self._stop_event = stop_event
|
||||
|
||||
def _should_retry(self, exception: requests.exceptions.RequestException) -> bool:
|
||||
"""
|
||||
Determine if a request should be retried based on the exception.
|
||||
|
||||
Args:
|
||||
exception: The request exception that occurred
|
||||
|
||||
Returns:
|
||||
True if the request should be retried
|
||||
"""
|
||||
# Retry on connection errors and timeouts
|
||||
if isinstance(exception, (requests.exceptions.ConnectionError,
|
||||
requests.exceptions.Timeout)):
|
||||
return True
|
||||
|
||||
if isinstance(exception, requests.exceptions.HTTPError):
|
||||
if hasattr(exception, 'response') and exception.response:
|
||||
# Retry on server errors (5xx) AND on rate-limiting errors (429)
|
||||
return exception.response.status_code >= 500 or exception.response.status_code == 429
|
||||
|
||||
return False
|
||||
|
||||
def log_relationship_discovery(self, source_node: str, target_node: str,
|
||||
relationship_type: str,
|
||||
confidence_score: float,
|
||||
@@ -285,7 +543,7 @@ class BaseProvider(ABC):
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get provider statistics.
|
||||
Get provider statistics including cache performance.
|
||||
|
||||
Returns:
|
||||
Dictionary containing provider performance metrics
|
||||
@@ -297,5 +555,8 @@ class BaseProvider(ABC):
|
||||
'failed_requests': self.failed_requests,
|
||||
'success_rate': (self.successful_requests / self.total_requests * 100) if self.total_requests > 0 else 0,
|
||||
'relationships_found': self.total_relationships_found,
|
||||
'rate_limit': self.rate_limiter.requests_per_minute
|
||||
'rate_limit': self.rate_limiter.requests_per_minute,
|
||||
'cache_hits': self.cache_hits,
|
||||
'cache_misses': self.cache_misses,
|
||||
'cache_hit_rate': (self.cache_hits / (self.cache_hits + self.cache_misses) * 100) if (self.cache_hits + self.cache_misses) > 0 else 0
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
# dnsrecon/providers/crtsh_provider.py
|
||||
"""
|
||||
Certificate Transparency provider using crt.sh.
|
||||
Discovers domain relationships through certificate SAN analysis with comprehensive certificate tracking.
|
||||
Stores certificates as metadata on domain nodes rather than creating certificate nodes.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Tuple, Set
|
||||
from urllib.parse import quote
|
||||
from datetime import datetime, timezone
|
||||
@@ -16,10 +18,10 @@ from utils.helpers import _is_valid_domain
|
||||
class CrtShProvider(BaseProvider):
|
||||
"""
|
||||
Provider for querying crt.sh certificate transparency database.
|
||||
Now uses session-specific configuration and caching with accumulative behavior.
|
||||
Now uses session-specific configuration and caching.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, session_config=None):
|
||||
def __init__(self, session_config=None):
|
||||
"""Initialize CrtSh provider with session-specific configuration."""
|
||||
super().__init__(
|
||||
name="crtsh",
|
||||
@@ -29,10 +31,6 @@ class CrtShProvider(BaseProvider):
|
||||
)
|
||||
self.base_url = "https://crt.sh/"
|
||||
self._stop_event = None
|
||||
|
||||
# Initialize cache directory
|
||||
self.cache_dir = Path('cache') / 'crtsh'
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_name(self) -> str:
|
||||
"""Return the provider name."""
|
||||
@@ -58,158 +56,6 @@ class CrtShProvider(BaseProvider):
|
||||
"""
|
||||
return True
|
||||
|
||||
def _get_cache_file_path(self, domain: str) -> Path:
|
||||
"""Generate cache file path for a domain."""
|
||||
# Sanitize domain for filename safety
|
||||
safe_domain = domain.replace('.', '_').replace('/', '_').replace('\\', '_')
|
||||
return self.cache_dir / f"{safe_domain}.json"
|
||||
|
||||
def _get_cache_status(self, cache_file_path: Path) -> str:
|
||||
"""
|
||||
Check cache status for a domain.
|
||||
Returns: 'not_found', 'fresh', or 'stale'
|
||||
"""
|
||||
if not cache_file_path.exists():
|
||||
return "not_found"
|
||||
|
||||
try:
|
||||
with open(cache_file_path, 'r') as f:
|
||||
cache_data = json.load(f)
|
||||
|
||||
last_query_str = cache_data.get("last_upstream_query")
|
||||
if not last_query_str:
|
||||
return "stale" # Invalid cache format
|
||||
|
||||
last_query = datetime.fromisoformat(last_query_str.replace('Z', '+00:00'))
|
||||
hours_since_query = (datetime.now(timezone.utc) - last_query).total_seconds() / 3600
|
||||
|
||||
cache_timeout = self.config.cache_timeout_hours
|
||||
if hours_since_query < cache_timeout:
|
||||
return "fresh"
|
||||
else:
|
||||
return "stale"
|
||||
|
||||
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||
self.logger.logger.warning(f"Invalid cache file format for {cache_file_path}: {e}")
|
||||
return "stale"
|
||||
|
||||
def _load_cached_certificates(self, cache_file_path: Path) -> List[Dict[str, Any]]:
|
||||
"""Load certificates from cache file."""
|
||||
try:
|
||||
with open(cache_file_path, 'r') as f:
|
||||
cache_data = json.load(f)
|
||||
return cache_data.get('certificates', [])
|
||||
except (json.JSONDecodeError, FileNotFoundError, KeyError) as e:
|
||||
self.logger.logger.error(f"Failed to load cached certificates from {cache_file_path}: {e}")
|
||||
return []
|
||||
|
||||
def _query_crtsh_api(self, domain: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query crt.sh API for raw certificate data.
|
||||
Raises exceptions for network errors to allow core logic to retry.
|
||||
"""
|
||||
url = f"{self.base_url}?q={quote(domain)}&output=json"
|
||||
response = self.make_request(url, target_indicator=domain)
|
||||
|
||||
if not response or response.status_code != 200:
|
||||
# This could be a temporary error - raise exception so core can retry
|
||||
raise requests.exceptions.RequestException(f"crt.sh API returned status {response.status_code if response else 'None'}")
|
||||
|
||||
certificates = response.json()
|
||||
if not certificates:
|
||||
return []
|
||||
|
||||
return certificates
|
||||
|
||||
def _create_cache_file(self, cache_file_path: Path, domain: str, certificates: List[Dict[str, Any]]) -> None:
|
||||
"""Create new cache file with certificates."""
|
||||
try:
|
||||
cache_data = {
|
||||
"domain": domain,
|
||||
"first_cached": datetime.now(timezone.utc).isoformat(),
|
||||
"last_upstream_query": datetime.now(timezone.utc).isoformat(),
|
||||
"upstream_query_count": 1,
|
||||
"certificates": certificates
|
||||
}
|
||||
|
||||
cache_file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(cache_file_path, 'w') as f:
|
||||
json.dump(cache_data, f, separators=(',', ':'))
|
||||
|
||||
self.logger.logger.info(f"Created cache file for {domain} with {len(certificates)} certificates")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.logger.warning(f"Failed to create cache file for {domain}: {e}")
|
||||
|
||||
def _append_to_cache(self, cache_file_path: Path, new_certificates: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Append new certificates to existing cache and return all certificates."""
|
||||
try:
|
||||
# Load existing cache
|
||||
with open(cache_file_path, 'r') as f:
|
||||
cache_data = json.load(f)
|
||||
|
||||
# Track existing certificate IDs to avoid duplicates
|
||||
existing_ids = {cert.get('id') for cert in cache_data.get('certificates', [])}
|
||||
|
||||
# Add only new certificates
|
||||
added_count = 0
|
||||
for cert in new_certificates:
|
||||
cert_id = cert.get('id')
|
||||
if cert_id and cert_id not in existing_ids:
|
||||
cache_data['certificates'].append(cert)
|
||||
existing_ids.add(cert_id)
|
||||
added_count += 1
|
||||
|
||||
# Update metadata
|
||||
cache_data['last_upstream_query'] = datetime.now(timezone.utc).isoformat()
|
||||
cache_data['upstream_query_count'] = cache_data.get('upstream_query_count', 0) + 1
|
||||
|
||||
# Write updated cache
|
||||
with open(cache_file_path, 'w') as f:
|
||||
json.dump(cache_data, f, separators=(',', ':'))
|
||||
|
||||
total_certs = len(cache_data['certificates'])
|
||||
self.logger.logger.info(f"Appended {added_count} new certificates to cache. Total: {total_certs}")
|
||||
|
||||
return cache_data['certificates']
|
||||
|
||||
except Exception as e:
|
||||
self.logger.logger.warning(f"Failed to append to cache: {e}")
|
||||
return new_certificates # Fallback to new certificates only
|
||||
|
||||
def _parse_issuer_organization(self, issuer_dn: str) -> str:
|
||||
"""
|
||||
Parse the issuer Distinguished Name to extract just the organization name.
|
||||
|
||||
Args:
|
||||
issuer_dn: Full issuer DN string (e.g., "C=US, O=Let's Encrypt, CN=R11")
|
||||
|
||||
Returns:
|
||||
Organization name (e.g., "Let's Encrypt") or original string if parsing fails
|
||||
"""
|
||||
if not issuer_dn:
|
||||
return issuer_dn
|
||||
|
||||
try:
|
||||
# Split by comma and look for O= component
|
||||
components = [comp.strip() for comp in issuer_dn.split(',')]
|
||||
|
||||
for component in components:
|
||||
if component.startswith('O='):
|
||||
# Extract the value after O=
|
||||
org_name = component[2:].strip()
|
||||
# Remove quotes if present
|
||||
if org_name.startswith('"') and org_name.endswith('"'):
|
||||
org_name = org_name[1:-1]
|
||||
return org_name
|
||||
|
||||
# If no O= component found, return the original string
|
||||
return issuer_dn
|
||||
|
||||
except Exception as e:
|
||||
self.logger.logger.debug(f"Failed to parse issuer DN '{issuer_dn}': {e}")
|
||||
return issuer_dn
|
||||
|
||||
def _parse_certificate_date(self, date_string: str) -> datetime:
|
||||
"""
|
||||
Parse certificate date from crt.sh format.
|
||||
@@ -287,15 +133,10 @@ class CrtShProvider(BaseProvider):
|
||||
Returns:
|
||||
Comprehensive certificate metadata dictionary
|
||||
"""
|
||||
# Parse the issuer name to get just the organization
|
||||
raw_issuer_name = cert_data.get('issuer_name', '')
|
||||
parsed_issuer_name = self._parse_issuer_organization(raw_issuer_name)
|
||||
|
||||
metadata = {
|
||||
'certificate_id': cert_data.get('id'),
|
||||
'serial_number': cert_data.get('serial_number'),
|
||||
'issuer_name': parsed_issuer_name, # Use parsed organization name
|
||||
#'issuer_name_full': raw_issuer_name, # deliberately left out, because its not useful in most cases
|
||||
'issuer_name': cert_data.get('issuer_name'),
|
||||
'issuer_ca_id': cert_data.get('issuer_ca_id'),
|
||||
'common_name': cert_data.get('common_name'),
|
||||
'not_before': cert_data.get('not_before'),
|
||||
@@ -326,8 +167,7 @@ class CrtShProvider(BaseProvider):
|
||||
|
||||
def query_domain(self, domain: str) -> List[Tuple[str, str, str, float, Dict[str, Any]]]:
|
||||
"""
|
||||
Query crt.sh for certificates containing the domain with caching support.
|
||||
Properly raises exceptions for network errors to allow core logic retries.
|
||||
Query crt.sh for certificates containing the domain.
|
||||
"""
|
||||
if not _is_valid_domain(domain):
|
||||
return []
|
||||
@@ -337,167 +177,131 @@ class CrtShProvider(BaseProvider):
|
||||
print(f"CrtSh query cancelled before start for domain: {domain}")
|
||||
return []
|
||||
|
||||
# === CACHING LOGIC ===
|
||||
cache_file = self._get_cache_file_path(domain)
|
||||
cache_status = self._get_cache_status(cache_file)
|
||||
|
||||
certificates = []
|
||||
|
||||
try:
|
||||
if cache_status == "fresh":
|
||||
# Use cached data
|
||||
certificates = self._load_cached_certificates(cache_file)
|
||||
self.logger.logger.info(f"Using cached data for {domain} ({len(certificates)} certificates)")
|
||||
|
||||
elif cache_status == "not_found":
|
||||
# Fresh query, create new cache
|
||||
certificates = self._query_crtsh_api(domain)
|
||||
if certificates: # Only cache if we got results
|
||||
self._create_cache_file(cache_file, domain, certificates)
|
||||
self.logger.logger.info(f"Cached fresh data for {domain} ({len(certificates)} certificates)")
|
||||
else:
|
||||
self.logger.logger.info(f"No certificates found for {domain}, not caching")
|
||||
|
||||
elif cache_status == "stale":
|
||||
# Append query, update existing cache
|
||||
try:
|
||||
new_certificates = self._query_crtsh_api(domain)
|
||||
if new_certificates:
|
||||
certificates = self._append_to_cache(cache_file, new_certificates)
|
||||
self.logger.logger.info(f"Refreshed and appended cache for {domain}")
|
||||
else:
|
||||
# Use existing cache if API returns no results
|
||||
certificates = self._load_cached_certificates(cache_file)
|
||||
self.logger.logger.info(f"API returned no new results, using existing cache for {domain}")
|
||||
except requests.exceptions.RequestException:
|
||||
# If API call fails for stale cache, use cached data and re-raise for retry logic
|
||||
certificates = self._load_cached_certificates(cache_file)
|
||||
if certificates:
|
||||
self.logger.logger.warning(f"API call failed for {domain}, using stale cache data ({len(certificates)} certificates)")
|
||||
# Don't re-raise here, just use cached data
|
||||
else:
|
||||
# No cached data and API failed - re-raise for retry
|
||||
raise
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
# Network/API errors should be re-raised so core logic can retry
|
||||
self.logger.logger.error(f"API query failed for {domain}: {e}")
|
||||
raise e
|
||||
except json.JSONDecodeError as e:
|
||||
# JSON parsing errors should also be raised for retry
|
||||
self.logger.logger.error(f"Failed to parse JSON response from crt.sh for {domain}: {e}")
|
||||
raise e
|
||||
|
||||
# Check for cancellation after cache operations
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh query cancelled after cache operations for domain: {domain}")
|
||||
return []
|
||||
|
||||
if not certificates:
|
||||
return []
|
||||
|
||||
return self._process_certificates_to_relationships(domain, certificates)
|
||||
|
||||
def _process_certificates_to_relationships(self, domain: str, certificates: List[Dict[str, Any]]) -> List[Tuple[str, str, str, float, Dict[str, Any]]]:
|
||||
"""
|
||||
Process certificates to relationships using existing logic.
|
||||
This method contains the original processing logic from query_domain.
|
||||
"""
|
||||
relationships = []
|
||||
|
||||
# Check for cancellation before processing
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh processing cancelled before processing for domain: {domain}")
|
||||
return []
|
||||
|
||||
# Aggregate certificate data by domain
|
||||
domain_certificates = {}
|
||||
all_discovered_domains = set()
|
||||
|
||||
# Process certificates with cancellation checking
|
||||
for i, cert_data in enumerate(certificates):
|
||||
# Check for cancellation every 5 certificates for faster response
|
||||
if i % 5 == 0 and self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh processing cancelled at certificate {i} for domain: {domain}")
|
||||
break
|
||||
|
||||
cert_metadata = self._extract_certificate_metadata(cert_data)
|
||||
cert_domains = self._extract_domains_from_certificate(cert_data)
|
||||
try:
|
||||
# Query crt.sh for certificates
|
||||
url = f"{self.base_url}?q={quote(domain)}&output=json"
|
||||
response = self.make_request(url, target_indicator=domain, max_retries=3)
|
||||
|
||||
# Add all domains from this certificate to our tracking
|
||||
all_discovered_domains.update(cert_domains)
|
||||
for cert_domain in cert_domains:
|
||||
if not _is_valid_domain(cert_domain):
|
||||
if not response or response.status_code != 200:
|
||||
return []
|
||||
|
||||
# Check for cancellation after request
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh query cancelled after request for domain: {domain}")
|
||||
return []
|
||||
|
||||
certificates = response.json()
|
||||
|
||||
if not certificates:
|
||||
return []
|
||||
|
||||
# Check for cancellation before processing
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh query cancelled before processing for domain: {domain}")
|
||||
return []
|
||||
|
||||
# Aggregate certificate data by domain
|
||||
domain_certificates = {}
|
||||
all_discovered_domains = set()
|
||||
|
||||
# Process certificates with cancellation checking
|
||||
for i, cert_data in enumerate(certificates):
|
||||
# Check for cancellation every 5 certificates instead of 10 for faster response
|
||||
if i % 5 == 0 and self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh processing cancelled at certificate {i} for domain: {domain}")
|
||||
break
|
||||
|
||||
cert_metadata = self._extract_certificate_metadata(cert_data)
|
||||
cert_domains = self._extract_domains_from_certificate(cert_data)
|
||||
|
||||
# Add all domains from this certificate to our tracking
|
||||
for cert_domain in cert_domains:
|
||||
# Additional stop check during domain processing
|
||||
if i % 20 == 0 and self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh domain processing cancelled for domain: {domain}")
|
||||
break
|
||||
|
||||
if not _is_valid_domain(cert_domain):
|
||||
continue
|
||||
|
||||
all_discovered_domains.add(cert_domain)
|
||||
|
||||
# Initialize domain certificate list if needed
|
||||
if cert_domain not in domain_certificates:
|
||||
domain_certificates[cert_domain] = []
|
||||
|
||||
# Add this certificate to the domain's certificate list
|
||||
domain_certificates[cert_domain].append(cert_metadata)
|
||||
|
||||
# Final cancellation check before creating relationships
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh query cancelled before relationship creation for domain: {domain}")
|
||||
return []
|
||||
|
||||
# Create relationships from query domain to ALL discovered domains with stop checking
|
||||
for i, discovered_domain in enumerate(all_discovered_domains):
|
||||
if discovered_domain == domain:
|
||||
continue # Skip self-relationships
|
||||
|
||||
# Check for cancellation every 10 relationships
|
||||
if i % 10 == 0 and self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh relationship creation cancelled for domain: {domain}")
|
||||
break
|
||||
|
||||
if not _is_valid_domain(discovered_domain):
|
||||
continue
|
||||
|
||||
# Initialize domain certificate list if needed
|
||||
if cert_domain not in domain_certificates:
|
||||
domain_certificates[cert_domain] = []
|
||||
# Get certificates for both domains
|
||||
query_domain_certs = domain_certificates.get(domain, [])
|
||||
discovered_domain_certs = domain_certificates.get(discovered_domain, [])
|
||||
|
||||
# Add this certificate to the domain's certificate list
|
||||
domain_certificates[cert_domain].append(cert_metadata)
|
||||
|
||||
# Final cancellation check before creating relationships
|
||||
if self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh query cancelled before relationship creation for domain: {domain}")
|
||||
return []
|
||||
|
||||
# Create relationships from query domain to ALL discovered domains with stop checking
|
||||
for i, discovered_domain in enumerate(all_discovered_domains):
|
||||
if discovered_domain == domain:
|
||||
continue # Skip self-relationships
|
||||
|
||||
# Check for cancellation every 10 relationships
|
||||
if i % 10 == 0 and self._stop_event and self._stop_event.is_set():
|
||||
print(f"CrtSh relationship creation cancelled for domain: {domain}")
|
||||
break
|
||||
|
||||
if not _is_valid_domain(discovered_domain):
|
||||
continue
|
||||
|
||||
# Get certificates for both domains
|
||||
query_domain_certs = domain_certificates.get(domain, [])
|
||||
discovered_domain_certs = domain_certificates.get(discovered_domain, [])
|
||||
|
||||
# Find shared certificates (for metadata purposes)
|
||||
shared_certificates = self._find_shared_certificates(query_domain_certs, discovered_domain_certs)
|
||||
|
||||
# Calculate confidence based on relationship type and shared certificates
|
||||
confidence = self._calculate_domain_relationship_confidence(
|
||||
domain, discovered_domain, shared_certificates, all_discovered_domains
|
||||
)
|
||||
|
||||
# Create comprehensive raw data for the relationship
|
||||
relationship_raw_data = {
|
||||
'relationship_type': 'certificate_discovery',
|
||||
'shared_certificates': shared_certificates,
|
||||
'total_shared_certs': len(shared_certificates),
|
||||
'discovery_context': self._determine_relationship_context(discovered_domain, domain),
|
||||
'domain_certificates': {
|
||||
domain: self._summarize_certificates(query_domain_certs),
|
||||
discovered_domain: self._summarize_certificates(discovered_domain_certs)
|
||||
# Find shared certificates (for metadata purposes)
|
||||
shared_certificates = self._find_shared_certificates(query_domain_certs, discovered_domain_certs)
|
||||
|
||||
# Calculate confidence based on relationship type and shared certificates
|
||||
confidence = self._calculate_domain_relationship_confidence(
|
||||
domain, discovered_domain, shared_certificates, all_discovered_domains
|
||||
)
|
||||
|
||||
# Create comprehensive raw data for the relationship
|
||||
relationship_raw_data = {
|
||||
'relationship_type': 'certificate_discovery',
|
||||
'shared_certificates': shared_certificates,
|
||||
'total_shared_certs': len(shared_certificates),
|
||||
'discovery_context': self._determine_relationship_context(discovered_domain, domain),
|
||||
'domain_certificates': {
|
||||
domain: self._summarize_certificates(query_domain_certs),
|
||||
discovered_domain: self._summarize_certificates(discovered_domain_certs)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Create domain -> domain relationship
|
||||
relationships.append((
|
||||
domain,
|
||||
discovered_domain,
|
||||
'san_certificate',
|
||||
confidence,
|
||||
relationship_raw_data
|
||||
))
|
||||
|
||||
# Log the relationship discovery
|
||||
self.log_relationship_discovery(
|
||||
source_node=domain,
|
||||
target_node=discovered_domain,
|
||||
relationship_type='san_certificate',
|
||||
confidence_score=confidence,
|
||||
raw_data=relationship_raw_data,
|
||||
discovery_method="certificate_transparency_analysis"
|
||||
)
|
||||
|
||||
# Create domain -> domain relationship
|
||||
relationships.append((
|
||||
domain,
|
||||
discovered_domain,
|
||||
'san_certificate',
|
||||
confidence,
|
||||
relationship_raw_data
|
||||
))
|
||||
|
||||
# Log the relationship discovery
|
||||
self.log_relationship_discovery(
|
||||
source_node=domain,
|
||||
target_node=discovered_domain,
|
||||
relationship_type='san_certificate',
|
||||
confidence_score=confidence,
|
||||
raw_data=relationship_raw_data,
|
||||
discovery_method="certificate_transparency_analysis"
|
||||
)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
self.logger.logger.error(f"Failed to parse JSON response from crt.sh: {e}")
|
||||
except requests.exceptions.RequestException as e:
|
||||
self.logger.logger.error(f"HTTP request to crt.sh failed: {e}")
|
||||
|
||||
|
||||
return relationships
|
||||
|
||||
def _find_shared_certificates(self, certs1: List[Dict[str, Any]], certs2: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
@@ -548,7 +352,7 @@ class CrtShProvider(BaseProvider):
|
||||
expired_count = len(certificates) - valid_count
|
||||
expires_soon_count = sum(1 for cert in certificates if cert.get('expires_soon'))
|
||||
|
||||
# Get unique issuers (using parsed organization names)
|
||||
# Get unique issuers
|
||||
unique_issuers = list(set(cert.get('issuer_name') for cert in certificates if cert.get('issuer_name')))
|
||||
|
||||
# Find the most recent certificate
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# dnsrecon/providers/dns_provider.py
|
||||
|
||||
from dns import resolver, reversename
|
||||
import dns.resolver
|
||||
import dns.reversename
|
||||
from typing import List, Dict, Any, Tuple
|
||||
from .base_provider import BaseProvider
|
||||
from utils.helpers import _is_valid_ip, _is_valid_domain
|
||||
@@ -12,7 +13,7 @@ class DNSProvider(BaseProvider):
|
||||
Now uses session-specific configuration.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, session_config=None):
|
||||
def __init__(self, session_config=None):
|
||||
"""Initialize DNS provider with session-specific configuration."""
|
||||
super().__init__(
|
||||
name="dns",
|
||||
@@ -22,7 +23,7 @@ class DNSProvider(BaseProvider):
|
||||
)
|
||||
|
||||
# Configure DNS resolver
|
||||
self.resolver = resolver.Resolver()
|
||||
self.resolver = dns.resolver.Resolver()
|
||||
self.resolver.timeout = 5
|
||||
self.resolver.lifetime = 10
|
||||
#self.resolver.nameservers = ['127.0.0.1']
|
||||
@@ -50,7 +51,12 @@ class DNSProvider(BaseProvider):
|
||||
def query_domain(self, domain: str) -> List[Tuple[str, str, str, float, Dict[str, Any]]]:
|
||||
"""
|
||||
Query DNS records for the domain to discover relationships.
|
||||
...
|
||||
|
||||
Args:
|
||||
domain: Domain to investigate
|
||||
|
||||
Returns:
|
||||
List of relationships discovered from DNS analysis
|
||||
"""
|
||||
if not _is_valid_domain(domain):
|
||||
return []
|
||||
@@ -59,15 +65,7 @@ class DNSProvider(BaseProvider):
|
||||
|
||||
# Query all record types
|
||||
for record_type in ['A', 'AAAA', 'CNAME', 'MX', 'NS', 'SOA', 'TXT', 'SRV', 'CAA']:
|
||||
try:
|
||||
relationships.extend(self._query_record(domain, record_type))
|
||||
except resolver.NoAnswer:
|
||||
# This is not an error, just a confirmation that the record doesn't exist.
|
||||
self.logger.logger.debug(f"No {record_type} record found for {domain}")
|
||||
except Exception as e:
|
||||
self.failed_requests += 1
|
||||
self.logger.logger.debug(f"{record_type} record query failed for {domain}: {e}")
|
||||
# Optionally, you might want to re-raise other, more serious exceptions.
|
||||
relationships.extend(self._query_record(domain, record_type))
|
||||
|
||||
return relationships
|
||||
|
||||
@@ -89,7 +87,7 @@ class DNSProvider(BaseProvider):
|
||||
try:
|
||||
# Perform reverse DNS lookup
|
||||
self.total_requests += 1
|
||||
reverse_name = reversename.from_address(ip)
|
||||
reverse_name = dns.reversename.from_address(ip)
|
||||
response = self.resolver.resolve(reverse_name, 'PTR')
|
||||
self.successful_requests += 1
|
||||
|
||||
@@ -121,14 +119,9 @@ class DNSProvider(BaseProvider):
|
||||
discovery_method="reverse_dns_lookup"
|
||||
)
|
||||
|
||||
except resolver.NXDOMAIN:
|
||||
self.failed_requests += 1
|
||||
self.logger.logger.debug(f"Reverse DNS lookup failed for {ip}: NXDOMAIN")
|
||||
except Exception as e:
|
||||
self.failed_requests += 1
|
||||
self.logger.logger.debug(f"Reverse DNS lookup failed for {ip}: {e}")
|
||||
# Re-raise the exception so the scanner can handle the failure
|
||||
raise e
|
||||
|
||||
return relationships
|
||||
|
||||
@@ -192,7 +185,5 @@ class DNSProvider(BaseProvider):
|
||||
except Exception as e:
|
||||
self.failed_requests += 1
|
||||
self.logger.logger.debug(f"{record_type} record query failed for {domain}: {e}")
|
||||
# Re-raise the exception so the scanner can handle it
|
||||
raise e
|
||||
|
||||
return relationships
|
||||
@@ -1,4 +1,7 @@
|
||||
# dnsrecon/providers/shodan_provider.py
|
||||
"""
|
||||
Shodan provider for DNSRecon.
|
||||
Discovers IP relationships and infrastructure context through Shodan API.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import List, Dict, Any, Tuple
|
||||
@@ -12,7 +15,7 @@ class ShodanProvider(BaseProvider):
|
||||
Now uses session-specific API keys.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, session_config=None):
|
||||
def __init__(self, session_config=None):
|
||||
"""Initialize Shodan provider with session-specific configuration."""
|
||||
super().__init__(
|
||||
name="shodan",
|
||||
|
||||
@@ -6,5 +6,4 @@ Werkzeug>=2.3.7
|
||||
urllib3>=2.0.0
|
||||
dnspython>=2.4.2
|
||||
gunicorn
|
||||
redis
|
||||
python-dotenv
|
||||
redis
|
||||
@@ -272,24 +272,8 @@ input[type="text"]:focus, select:focus {
|
||||
text-shadow: 0 0 3px rgba(0, 255, 65, 0.3);
|
||||
}
|
||||
|
||||
.progress-container {
|
||||
padding: 0 1.5rem 1.5rem;
|
||||
}
|
||||
|
||||
.progress-info {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
font-size: 0.8rem;
|
||||
color: #999;
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
|
||||
#progress-compact {
|
||||
color: #00ff41;
|
||||
font-weight: 500;
|
||||
}
|
||||
|
||||
.progress-bar {
|
||||
margin: 1rem 1.5rem;
|
||||
height: 8px;
|
||||
background-color: #1a1a1a;
|
||||
border: 1px solid #444;
|
||||
@@ -533,7 +517,7 @@ input[type="text"]:focus, select:focus {
|
||||
color: #e0e0e0;
|
||||
}
|
||||
|
||||
.provider-stats, .provider-task-stats {
|
||||
.provider-stats {
|
||||
font-size: 0.8rem;
|
||||
color: #999;
|
||||
display: grid;
|
||||
@@ -542,13 +526,6 @@ input[type="text"]:focus, select:focus {
|
||||
margin-top: 0.5rem;
|
||||
}
|
||||
|
||||
.provider-task-stats {
|
||||
border-top: 1px solid #333;
|
||||
padding-top: 0.5rem;
|
||||
margin-top: 0.5rem;
|
||||
}
|
||||
|
||||
|
||||
.provider-stat {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
|
||||
@@ -216,8 +216,12 @@ class GraphManager {
|
||||
}
|
||||
});
|
||||
|
||||
// FIX: Comment out the problematic context menu handler
|
||||
this.network.on('oncontext', (params) => {
|
||||
params.event.preventDefault();
|
||||
// if (params.nodes.length > 0) {
|
||||
// this.showNodeContextMenu(params.pointer.DOM, params.nodes[0]);
|
||||
// }
|
||||
});
|
||||
|
||||
// Stabilization events with progress
|
||||
@@ -376,7 +380,7 @@ class GraphManager {
|
||||
// Single correlation value
|
||||
const value = Array.isArray(values) && values.length > 0 ? values[0] : (metadata.value || 'Unknown');
|
||||
const displayValue = typeof value === 'string' && value.length > 20 ? value.substring(0, 17) + '...' : value;
|
||||
processedNode.label = `${displayValue}`;
|
||||
processedNode.label = `Corr: ${displayValue}`;
|
||||
processedNode.title = `Correlation: ${value}`;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
/**
|
||||
* Main application logic for DNSRecon web interface
|
||||
* Handles UI interactions, API communication, and data flow
|
||||
* DEBUG VERSION WITH EXTRA LOGGING
|
||||
*/
|
||||
|
||||
class DNSReconApp {
|
||||
@@ -60,8 +61,9 @@ class DNSReconApp {
|
||||
scanStatus: document.getElementById('scan-status'),
|
||||
targetDisplay: document.getElementById('target-display'),
|
||||
depthDisplay: document.getElementById('depth-display'),
|
||||
progressDisplay: document.getElementById('progress-display'),
|
||||
indicatorsDisplay: document.getElementById('indicators-display'),
|
||||
relationshipsDisplay: document.getElementById('relationships-display'),
|
||||
progressCompact: document.getElementById('progress-compact'),
|
||||
progressFill: document.getElementById('progress-fill'),
|
||||
|
||||
// Provider elements
|
||||
@@ -445,7 +447,7 @@ class DNSReconApp {
|
||||
// Handle status changes
|
||||
if (status.status !== this.scanStatus) {
|
||||
console.log(`*** STATUS CHANGED: ${this.scanStatus} -> ${status.status} ***`);
|
||||
this.handleStatusChange(status.status, status.task_queue_size);
|
||||
this.handleStatusChange(status.status);
|
||||
}
|
||||
|
||||
this.scanStatus = status.status;
|
||||
@@ -540,19 +542,17 @@ class DNSReconApp {
|
||||
if (this.elements.depthDisplay) {
|
||||
this.elements.depthDisplay.textContent = `${status.current_depth}/${status.max_depth}`;
|
||||
}
|
||||
if (this.elements.progressDisplay) {
|
||||
this.elements.progressDisplay.textContent = `${status.progress_percentage.toFixed(1)}%`;
|
||||
}
|
||||
if (this.elements.indicatorsDisplay) {
|
||||
this.elements.indicatorsDisplay.textContent = status.indicators_processed || 0;
|
||||
}
|
||||
|
||||
// Update progress bar and compact display
|
||||
// Update progress bar with smooth animation
|
||||
if (this.elements.progressFill) {
|
||||
const completed = status.indicators_completed || 0;
|
||||
const enqueued = status.task_queue_size || 0;
|
||||
const totalTasks = completed + enqueued;
|
||||
const progressPercentage = totalTasks > 0 ? (completed / totalTasks) * 100 : 0;
|
||||
|
||||
this.elements.progressFill.style.width = `${progressPercentage}%`;
|
||||
if (this.elements.progressCompact) {
|
||||
this.elements.progressCompact.textContent = `${completed}/${totalTasks} - ${Math.round(progressPercentage)}%`;
|
||||
}
|
||||
|
||||
this.elements.progressFill.style.width = `${status.progress_percentage}%`;
|
||||
|
||||
// Add pulsing animation for active scans
|
||||
if (status.status === 'running') {
|
||||
this.elements.progressFill.parentElement.classList.add('scanning');
|
||||
@@ -574,8 +574,6 @@ class DNSReconApp {
|
||||
this.elements.sessionId.textContent = 'Session: Loading...';
|
||||
}
|
||||
}
|
||||
|
||||
this.setUIState(status.status, status.task_queue_size);
|
||||
|
||||
console.log('Status display updated successfully');
|
||||
} catch (error) {
|
||||
@@ -587,12 +585,12 @@ class DNSReconApp {
|
||||
* Handle status changes with improved state synchronization
|
||||
* @param {string} newStatus - New scan status
|
||||
*/
|
||||
handleStatusChange(newStatus, task_queue_size) {
|
||||
handleStatusChange(newStatus) {
|
||||
console.log(`=== STATUS CHANGE: ${this.scanStatus} -> ${newStatus} ===`);
|
||||
|
||||
switch (newStatus) {
|
||||
case 'running':
|
||||
this.setUIState('scanning', task_queue_size);
|
||||
this.setUIState('scanning');
|
||||
this.showSuccess('Scan is running');
|
||||
// Increase polling frequency for active scans
|
||||
this.startPolling(1000); // Poll every 1 second for running scans
|
||||
@@ -600,7 +598,7 @@ class DNSReconApp {
|
||||
break;
|
||||
|
||||
case 'completed':
|
||||
this.setUIState('completed', task_queue_size);
|
||||
this.setUIState('completed');
|
||||
this.stopPolling();
|
||||
this.showSuccess('Scan completed successfully');
|
||||
this.updateConnectionStatus('completed');
|
||||
@@ -611,7 +609,7 @@ class DNSReconApp {
|
||||
break;
|
||||
|
||||
case 'failed':
|
||||
this.setUIState('failed', task_queue_size);
|
||||
this.setUIState('failed');
|
||||
this.stopPolling();
|
||||
this.showError('Scan failed');
|
||||
this.updateConnectionStatus('error');
|
||||
@@ -619,7 +617,7 @@ class DNSReconApp {
|
||||
break;
|
||||
|
||||
case 'stopped':
|
||||
this.setUIState('stopped', task_queue_size);
|
||||
this.setUIState('stopped');
|
||||
this.stopPolling();
|
||||
this.showSuccess('Scan stopped');
|
||||
this.updateConnectionStatus('stopped');
|
||||
@@ -627,7 +625,7 @@ class DNSReconApp {
|
||||
break;
|
||||
|
||||
case 'idle':
|
||||
this.setUIState('idle', task_queue_size);
|
||||
this.setUIState('idle');
|
||||
this.stopPolling();
|
||||
this.updateConnectionStatus('idle');
|
||||
break;
|
||||
@@ -672,11 +670,9 @@ class DNSReconApp {
|
||||
/**
|
||||
* UI state management with immediate button updates
|
||||
*/
|
||||
setUIState(state, task_queue_size) {
|
||||
setUIState(state) {
|
||||
console.log(`Setting UI state to: ${state}`);
|
||||
|
||||
const isQueueEmpty = task_queue_size === 0;
|
||||
|
||||
switch (state) {
|
||||
case 'scanning':
|
||||
this.isScanning = true;
|
||||
@@ -705,12 +701,12 @@ class DNSReconApp {
|
||||
case 'stopped':
|
||||
this.isScanning = false;
|
||||
if (this.elements.startScan) {
|
||||
this.elements.startScan.disabled = !isQueueEmpty;
|
||||
this.elements.startScan.disabled = false;
|
||||
this.elements.startScan.classList.remove('loading');
|
||||
this.elements.startScan.innerHTML = '<span class="btn-icon">[RUN]</span><span>Start Reconnaissance</span>';
|
||||
}
|
||||
if (this.elements.addToGraph) {
|
||||
this.elements.addToGraph.disabled = !isQueueEmpty;
|
||||
this.elements.addToGraph.disabled = false;
|
||||
this.elements.addToGraph.classList.remove('loading');
|
||||
}
|
||||
if (this.elements.stopScan) {
|
||||
|
||||
@@ -90,20 +90,22 @@
|
||||
<span class="status-label">Depth:</span>
|
||||
<span id="depth-display" class="status-value">0/0</span>
|
||||
</div>
|
||||
<div class="status-row">
|
||||
<span class="status-label">Progress:</span>
|
||||
<span id="progress-display" class="status-value">0%</span>
|
||||
</div>
|
||||
<div class="status-row">
|
||||
<span class="status-label">Indicators:</span>
|
||||
<span id="indicators-display" class="status-value">0</span>
|
||||
</div>
|
||||
<div class="status-row">
|
||||
<span class="status-label">Relationships:</span>
|
||||
<span id="relationships-display" class="status-value">0</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="progress-container">
|
||||
<div class="progress-info">
|
||||
<span id="progress-label">Progress:</span>
|
||||
<span id="progress-compact">0/0</span>
|
||||
</div>
|
||||
<div class="progress-bar">
|
||||
<div id="progress-fill" class="progress-fill"></div>
|
||||
</div>
|
||||
<div class="progress-bar">
|
||||
<div id="progress-fill" class="progress-fill"></div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
|
||||
Reference in New Issue
Block a user