data-model #2
309
app.py
309
app.py
@ -18,38 +18,29 @@ from utils.helpers import is_valid_target
|
||||
|
||||
|
||||
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)
|
||||
|
||||
def get_user_scanner():
|
||||
"""
|
||||
Retrieves the scanner for the current session, or creates a new session only if none exists.
|
||||
Retrieves the scanner for the current session, or creates a new one if none exists.
|
||||
"""
|
||||
current_flask_session_id = session.get('dnsrecon_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:
|
||||
#print(f"Reusing existing session: {current_flask_session_id}")
|
||||
return current_flask_session_id, existing_scanner
|
||||
else:
|
||||
print(f"Session {current_flask_session_id} expired, will create new one")
|
||||
|
||||
# Only create new session if we absolutely don't have one
|
||||
print("Creating new session (no valid session 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
|
||||
|
||||
print(f"Created new session: {new_session_id}")
|
||||
return new_session_id, new_scanner
|
||||
|
||||
@app.route('/')
|
||||
@ -61,11 +52,8 @@ def index():
|
||||
@app.route('/api/scan/start', methods=['POST'])
|
||||
def start_scan():
|
||||
"""
|
||||
FIXED: Start a new reconnaissance scan while preserving session configuration.
|
||||
Only clears graph data, not the entire session with API keys.
|
||||
Starts a new reconnaissance scan.
|
||||
"""
|
||||
print("=== API: /api/scan/start called ===")
|
||||
|
||||
try:
|
||||
data = request.get_json()
|
||||
if not data or 'target' not in data:
|
||||
@ -76,25 +64,18 @@ def start_scan():
|
||||
clear_graph = data.get('clear_graph', True)
|
||||
force_rescan_target = data.get('force_rescan_target', None)
|
||||
|
||||
print(f"Parsed - target: '{target}', max_depth: {max_depth}, clear_graph: {clear_graph}, force_rescan: {force_rescan_target}")
|
||||
|
||||
# Validation
|
||||
if not target:
|
||||
return jsonify({'success': False, 'error': 'Target cannot be empty'}), 400
|
||||
if not is_valid_target(target):
|
||||
return jsonify({'success': False, 'error': 'Invalid target format. Please enter a valid domain or IP address.'}), 400
|
||||
return jsonify({'success': False, 'error': 'Invalid target format.'}), 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
|
||||
|
||||
# FIXED: Always reuse existing session, preserve API keys
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if not scanner:
|
||||
return jsonify({'success': False, 'error': 'Failed to get scanner instance.'}), 500
|
||||
|
||||
print(f"Using scanner {id(scanner)} in session {user_session_id}")
|
||||
|
||||
# FIXED: Pass clear_graph flag to scanner, let it handle graph clearing internally
|
||||
success = scanner.start_scan(target, max_depth, clear_graph=clear_graph, force_rescan_target=force_rescan_target)
|
||||
|
||||
if success:
|
||||
@ -102,8 +83,7 @@ def start_scan():
|
||||
'success': True,
|
||||
'message': 'Scan started successfully',
|
||||
'scan_id': scanner.logger.session_id,
|
||||
'user_session_id': user_session_id,
|
||||
'available_providers': [p.get_name() for p in scanner.providers] # Show which providers are active
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
@ -112,206 +92,98 @@ def start_scan():
|
||||
}), 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
|
||||
|
||||
@app.route('/api/scan/stop', methods=['POST'])
|
||||
def stop_scan():
|
||||
"""Stop the current scan with immediate GUI feedback."""
|
||||
print("=== API: /api/scan/stop called ===")
|
||||
|
||||
"""Stop the current scan."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
print(f"Stopping scan for session: {user_session_id}")
|
||||
|
||||
if not scanner:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'No scanner found for session'
|
||||
}), 404
|
||||
return jsonify({'success': False, 'error': 'No scanner found for session'}), 404
|
||||
|
||||
# Ensure session ID is set
|
||||
if not scanner.session_id:
|
||||
scanner.session_id = user_session_id
|
||||
|
||||
# Use the stop mechanism
|
||||
success = scanner.stop_scan()
|
||||
|
||||
# Also set the Redis stop signal directly for extra reliability
|
||||
scanner.stop_scan()
|
||||
session_manager.set_stop_signal(user_session_id)
|
||||
|
||||
# Force immediate status update
|
||||
session_manager.update_scanner_status(user_session_id, 'stopped')
|
||||
|
||||
# Update the full scanner state
|
||||
session_manager.update_session_scanner(user_session_id, scanner)
|
||||
|
||||
print(f"Stop scan completed. Success: {success}, Scanner status: {scanner.status}")
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': 'Scan stop requested - termination initiated',
|
||||
'user_session_id': user_session_id,
|
||||
'scanner_status': scanner.status,
|
||||
'stop_method': 'cross_process'
|
||||
'message': 'Scan stop requested',
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in stop_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/status', methods=['GET'])
|
||||
def get_scan_status():
|
||||
"""Get current scan status with error handling."""
|
||||
"""Get current scan status."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if not scanner:
|
||||
# Return default idle status if no scanner
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'status': {
|
||||
'status': 'idle',
|
||||
'target_domain': None,
|
||||
'current_depth': 0,
|
||||
'max_depth': 0,
|
||||
'current_indicator': '',
|
||||
'total_indicators_found': 0,
|
||||
'indicators_processed': 0,
|
||||
'progress_percentage': 0.0,
|
||||
'enabled_providers': [],
|
||||
'graph_statistics': {},
|
||||
'status': 'idle', 'target_domain': None, 'current_depth': 0,
|
||||
'max_depth': 0, 'progress_percentage': 0.0,
|
||||
'user_session_id': user_session_id
|
||||
}
|
||||
})
|
||||
|
||||
# Ensure session ID is set
|
||||
if not scanner.session_id:
|
||||
scanner.session_id = user_session_id
|
||||
|
||||
status = scanner.get_scan_status()
|
||||
status['user_session_id'] = user_session_id
|
||||
|
||||
# Additional debug info
|
||||
status['debug_info'] = {
|
||||
'scanner_object_id': id(scanner),
|
||||
'session_id_set': bool(scanner.session_id),
|
||||
'has_scan_thread': bool(scanner.scan_thread and scanner.scan_thread.is_alive()),
|
||||
'provider_count': len(scanner.providers),
|
||||
'provider_names': [p.get_name() for p in scanner.providers]
|
||||
}
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'status': status
|
||||
})
|
||||
return jsonify({'success': True, 'status': status})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_scan_status endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}',
|
||||
'fallback_status': {
|
||||
'status': 'error',
|
||||
'target_domain': None,
|
||||
'current_depth': 0,
|
||||
'max_depth': 0,
|
||||
'progress_percentage': 0.0
|
||||
}
|
||||
'success': False, 'error': f'Internal server error: {str(e)}',
|
||||
'fallback_status': {'status': 'error', 'progress_percentage': 0.0}
|
||||
}), 500
|
||||
|
||||
|
||||
|
||||
@app.route('/api/graph', methods=['GET'])
|
||||
def get_graph_data():
|
||||
"""Get current graph data with error handling and proper empty graph structure."""
|
||||
"""Get current graph data."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if not scanner:
|
||||
# FIXED: Return proper empty graph structure instead of None
|
||||
empty_graph = {
|
||||
'nodes': [],
|
||||
'edges': [],
|
||||
'statistics': {
|
||||
'node_count': 0,
|
||||
'edge_count': 0,
|
||||
'creation_time': datetime.now(timezone.utc).isoformat(),
|
||||
'last_modified': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
}
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'graph': empty_graph,
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
|
||||
graph_data = scanner.get_graph_data()
|
||||
|
||||
if not graph_data:
|
||||
graph_data = {
|
||||
'nodes': [],
|
||||
'edges': [],
|
||||
'statistics': {
|
||||
'node_count': 0,
|
||||
'edge_count': 0,
|
||||
'creation_time': datetime.now(timezone.utc).isoformat(),
|
||||
'last_modified': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
'nodes': [], 'edges': [],
|
||||
'statistics': {'node_count': 0, 'edge_count': 0}
|
||||
}
|
||||
|
||||
# FIXED: Ensure required fields exist
|
||||
if 'nodes' not in graph_data:
|
||||
graph_data['nodes'] = []
|
||||
if 'edges' not in graph_data:
|
||||
graph_data['edges'] = []
|
||||
if 'statistics' not in graph_data:
|
||||
graph_data['statistics'] = {
|
||||
'node_count': len(graph_data['nodes']),
|
||||
'edge_count': len(graph_data['edges']),
|
||||
'creation_time': datetime.now(timezone.utc).isoformat(),
|
||||
'last_modified': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
if not scanner:
|
||||
return jsonify({'success': True, 'graph': empty_graph, 'user_session_id': user_session_id})
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'graph': graph_data,
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
graph_data = scanner.get_graph_data() or empty_graph
|
||||
|
||||
return jsonify({'success': True, 'graph': graph_data, 'user_session_id': user_session_id})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_graph_data endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
# FIXED: Return proper error structure with empty graph fallback
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Internal server error: {str(e)}',
|
||||
'fallback_graph': {
|
||||
'nodes': [],
|
||||
'edges': [],
|
||||
'statistics': {
|
||||
'node_count': 0,
|
||||
'edge_count': 0,
|
||||
'creation_time': datetime.now(timezone.utc).isoformat(),
|
||||
'last_modified': datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
}
|
||||
'success': False, 'error': f'Internal server error: {str(e)}',
|
||||
'fallback_graph': {'nodes': [], 'edges': [], 'statistics': {}}
|
||||
}), 500
|
||||
|
||||
@app.route('/api/graph/large-entity/extract', methods=['POST'])
|
||||
def extract_from_large_entity():
|
||||
"""Extract a node from a large entity, making it a standalone node."""
|
||||
"""Extract a node from a large entity."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
large_entity_id = data.get('large_entity_id')
|
||||
@ -333,13 +205,12 @@ def extract_from_large_entity():
|
||||
return jsonify({'success': False, 'error': f'Failed to extract node {node_id}.'}), 500
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in extract_from_large_entity endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({'success': False, 'error': f'Internal server error: {str(e)}'}), 500
|
||||
|
||||
@app.route('/api/graph/node/<node_id>', methods=['DELETE'])
|
||||
def delete_graph_node(node_id):
|
||||
"""Delete a node from the graph for the current user session."""
|
||||
"""Delete a node from the graph."""
|
||||
try:
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
if not scanner:
|
||||
@ -348,14 +219,12 @@ def delete_graph_node(node_id):
|
||||
success = scanner.graph.remove_node(node_id)
|
||||
|
||||
if success:
|
||||
# Persist the change
|
||||
session_manager.update_session_scanner(user_session_id, scanner)
|
||||
return jsonify({'success': True, 'message': f'Node {node_id} deleted successfully.'})
|
||||
else:
|
||||
return jsonify({'success': False, 'error': f'Node {node_id} not found in graph.'}), 404
|
||||
return jsonify({'success': False, 'error': f'Node {node_id} not found.'}), 404
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in delete_graph_node endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({'success': False, 'error': f'Internal server error: {str(e)}'}), 500
|
||||
|
||||
@ -376,7 +245,6 @@ def revert_graph_action():
|
||||
action_data = data['data']
|
||||
|
||||
if action_type == 'delete':
|
||||
# Re-add the node
|
||||
node_to_add = action_data.get('node')
|
||||
if node_to_add:
|
||||
scanner.graph.add_node(
|
||||
@ -387,204 +255,121 @@ def revert_graph_action():
|
||||
metadata=node_to_add.get('metadata')
|
||||
)
|
||||
|
||||
# Re-add the edges
|
||||
edges_to_add = action_data.get('edges', [])
|
||||
for edge in edges_to_add:
|
||||
# Add edge only if both nodes exist to prevent errors
|
||||
if scanner.graph.graph.has_node(edge['from']) and scanner.graph.graph.has_node(edge['to']):
|
||||
scanner.graph.add_edge(
|
||||
source_id=edge['from'],
|
||||
target_id=edge['to'],
|
||||
source_id=edge['from'], target_id=edge['to'],
|
||||
relationship_type=edge['metadata']['relationship_type'],
|
||||
confidence_score=edge['metadata']['confidence_score'],
|
||||
source_provider=edge['metadata']['source_provider'],
|
||||
raw_data=edge.get('raw_data', {})
|
||||
)
|
||||
|
||||
# Persist the change
|
||||
session_manager.update_session_scanner(user_session_id, scanner)
|
||||
return jsonify({'success': True, 'message': 'Delete action reverted successfully.'})
|
||||
|
||||
return jsonify({'success': False, 'error': f'Unknown revert action type: {action_type}'}), 400
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in revert_graph_action endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({'success': False, 'error': f'Internal server error: {str(e)}'}), 500
|
||||
|
||||
|
||||
@app.route('/api/export', methods=['GET'])
|
||||
def export_results():
|
||||
"""Export complete scan results as downloadable JSON for the user session."""
|
||||
"""Export scan results as a JSON file."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
# Get complete results
|
||||
results = scanner.export_results()
|
||||
|
||||
# Add session information to export
|
||||
results['export_metadata'] = {
|
||||
'user_session_id': user_session_id,
|
||||
'export_timestamp': datetime.now(timezone.utc).isoformat(),
|
||||
'export_type': 'user_session_results'
|
||||
}
|
||||
|
||||
# Create filename with timestamp
|
||||
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"
|
||||
filename = f"dnsrecon_{target}_{timestamp}.json"
|
||||
|
||||
# Create in-memory file
|
||||
json_data = json.dumps(results, indent=2, ensure_ascii=False)
|
||||
json_data = json.dumps(results, indent=2)
|
||||
file_obj = io.BytesIO(json_data.encode('utf-8'))
|
||||
|
||||
return send_file(
|
||||
file_obj,
|
||||
as_attachment=True,
|
||||
download_name=filename,
|
||||
mimetype='application/json'
|
||||
file_obj, as_attachment=True,
|
||||
download_name=filename, mimetype='application/json'
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in export_results endpoint: {e}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'Export failed: {str(e)}'
|
||||
}), 500
|
||||
return jsonify({'success': False, 'error': f'Export failed: {str(e)}'}), 500
|
||||
|
||||
|
||||
@app.route('/api/providers', methods=['GET'])
|
||||
def get_providers():
|
||||
"""Get information about available providers for the user session."""
|
||||
|
||||
"""Get information about available providers."""
|
||||
try:
|
||||
# Get user-specific scanner
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
|
||||
if scanner and scanner.status == 'running':
|
||||
status = scanner.get_scan_status()
|
||||
currently_processing = status.get('currently_processing')
|
||||
if currently_processing:
|
||||
provider_name, target_item = currently_processing[0]
|
||||
print(f"DEBUG: RUNNING Task - Provider: {provider_name}, Target: {target_item}")
|
||||
|
||||
print(f"DEBUG: Task Queue Status - In Queue: {status.get('tasks_in_queue', 0)}, Completed: {status.get('tasks_completed', 0)}, Skipped: {status.get('tasks_skipped', 0)}, Rescheduled: {status.get('tasks_rescheduled', 0)}")
|
||||
elif not scanner:
|
||||
print("DEBUG: No active scanner session found.")
|
||||
|
||||
provider_info = scanner.get_provider_info()
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'providers': provider_info,
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
return jsonify({'success': True, 'providers': provider_info, 'user_session_id': user_session_id})
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_providers 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/config/api-keys', methods=['POST'])
|
||||
def set_api_keys():
|
||||
"""
|
||||
"""
|
||||
"""Set API keys for the current session."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if data is None:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'No API keys provided'
|
||||
}), 400
|
||||
return jsonify({'success': False, 'error': 'No API keys provided'}), 400
|
||||
|
||||
# Get user-specific scanner and config
|
||||
user_session_id, scanner = get_user_scanner()
|
||||
session_config = scanner.config
|
||||
|
||||
print(f"Setting API keys for session {user_session_id}: {list(data.keys())}")
|
||||
|
||||
updated_providers = []
|
||||
|
||||
# Iterate over the API keys provided in the request data
|
||||
for provider_name, api_key in data.items():
|
||||
# This allows us to both set and clear keys. The config
|
||||
# handles enabling/disabling based on if the key is empty.
|
||||
api_key_value = str(api_key or '').strip()
|
||||
success = session_config.set_api_key(provider_name.lower(), api_key_value)
|
||||
|
||||
if success:
|
||||
updated_providers.append(provider_name)
|
||||
print(f"API key {'set' if api_key_value else 'cleared'} for {provider_name}")
|
||||
|
||||
if updated_providers:
|
||||
# FIXED: Reinitialize scanner providers to apply the new keys
|
||||
print("Reinitializing providers with new API keys...")
|
||||
old_provider_count = len(scanner.providers)
|
||||
scanner._initialize_providers()
|
||||
new_provider_count = len(scanner.providers)
|
||||
|
||||
print(f"Providers reinitialized: {old_provider_count} -> {new_provider_count}")
|
||||
print(f"Available providers: {[p.get_name() for p in scanner.providers]}")
|
||||
|
||||
# Persist the updated scanner object back to the user's session
|
||||
session_manager.update_session_scanner(user_session_id, scanner)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': f'API keys updated for session {user_session_id}: {", ".join(updated_providers)}',
|
||||
'updated_providers': updated_providers,
|
||||
'user_session_id': user_session_id,
|
||||
'available_providers': [p.get_name() for p in scanner.providers]
|
||||
'message': f'API keys updated for: {", ".join(updated_providers)}',
|
||||
'user_session_id': user_session_id
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'No valid API keys were provided or provider names were incorrect.'
|
||||
}), 400
|
||||
return jsonify({'success': False, 'error': 'No valid API keys were provided.'}), 400
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in set_api_keys 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.errorhandler(404)
|
||||
def not_found(error):
|
||||
"""Handle 404 errors."""
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Endpoint not found'
|
||||
}), 404
|
||||
return jsonify({'success': False, 'error': 'Endpoint not found'}), 404
|
||||
|
||||
|
||||
@app.errorhandler(500)
|
||||
def internal_error(error):
|
||||
"""Handle 500 errors."""
|
||||
print(f"ERROR: 500 Internal Server Error: {error}")
|
||||
traceback.print_exc()
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Internal server error'
|
||||
}), 500
|
||||
return jsonify({'success': False, 'error': 'Internal server error'}), 500
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("Starting DNSRecon Flask application with streamlined session management...")
|
||||
|
||||
# Load configuration from environment
|
||||
config.load_from_env()
|
||||
|
||||
# Start Flask application
|
||||
print(f"Starting server on {config.flask_host}:{config.flask_port}")
|
||||
app.run(
|
||||
host=config.flask_host,
|
||||
port=config.flask_port,
|
||||
|
||||
@ -149,7 +149,7 @@ class GraphManager:
|
||||
|
||||
if self.graph.has_node(node_id) and not self.graph.has_edge(node_id, correlation_node_id):
|
||||
# Format relationship label as "corr_provider_attribute"
|
||||
relationship_label = f"corr_{provider}_{attribute}"
|
||||
relationship_label = f"{provider}_{attribute}"
|
||||
|
||||
self.add_edge(
|
||||
source_id=node_id,
|
||||
|
||||
329
core/scanner.py
329
core/scanner.py
@ -5,6 +5,7 @@ import traceback
|
||||
import os
|
||||
import importlib
|
||||
import redis
|
||||
import time
|
||||
from typing import List, Set, Dict, Any, Tuple, Optional
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from collections import defaultdict
|
||||
@ -30,13 +31,11 @@ class ScanStatus:
|
||||
class Scanner:
|
||||
"""
|
||||
Main scanning orchestrator for DNSRecon passive reconnaissance.
|
||||
FIXED: Now preserves session configuration including API keys when clearing graphs.
|
||||
UNIFIED: Combines comprehensive features with improved display formatting.
|
||||
"""
|
||||
|
||||
def __init__(self, session_config=None):
|
||||
"""Initialize scanner with session-specific configuration."""
|
||||
print("Initializing Scanner instance...")
|
||||
|
||||
try:
|
||||
# Use provided session config or create default
|
||||
if session_config is None:
|
||||
@ -57,15 +56,18 @@ class Scanner:
|
||||
self.target_retries = defaultdict(int)
|
||||
self.scan_failed_due_to_retries = False
|
||||
|
||||
# **NEW**: Track currently processing tasks to prevent processing after stop
|
||||
# Thread-safe processing tracking (from Document 1)
|
||||
self.currently_processing = set()
|
||||
self.processing_lock = threading.Lock()
|
||||
# Display-friendly processing list (from Document 2)
|
||||
self.currently_processing_display = []
|
||||
|
||||
# Scanning progress tracking
|
||||
self.total_indicators_found = 0
|
||||
self.indicators_processed = 0
|
||||
self.indicators_completed = 0
|
||||
self.tasks_re_enqueued = 0
|
||||
self.tasks_skipped = 0 # BUGFIX: Initialize tasks_skipped
|
||||
self.total_tasks_ever_enqueued = 0
|
||||
self.current_indicator = ""
|
||||
|
||||
@ -73,19 +75,19 @@ class Scanner:
|
||||
self.max_workers = self.config.max_concurrent_requests
|
||||
self.executor = None
|
||||
|
||||
# Status logger thread with improved formatting
|
||||
self.status_logger_thread = None
|
||||
self.status_logger_stop_event = threading.Event()
|
||||
|
||||
# Initialize providers with session config
|
||||
print("Calling _initialize_providers with session config...")
|
||||
self._initialize_providers()
|
||||
|
||||
# Initialize logger
|
||||
print("Initializing forensic logger...")
|
||||
self.logger = get_forensic_logger()
|
||||
|
||||
# Initialize global rate limiter
|
||||
self.rate_limiter = GlobalRateLimiter(redis.StrictRedis(db=0))
|
||||
|
||||
print("Scanner initialization complete")
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Scanner initialization failed: {e}")
|
||||
traceback.print_exc()
|
||||
@ -96,17 +98,14 @@ class Scanner:
|
||||
Check if stop is requested using both local and Redis-based signals.
|
||||
This ensures reliable termination across process boundaries.
|
||||
"""
|
||||
# Check local threading event first (fastest)
|
||||
if self.stop_event.is_set():
|
||||
return True
|
||||
|
||||
# Check Redis-based stop signal if session ID is available
|
||||
if self.session_id:
|
||||
try:
|
||||
from core.session_manager import session_manager
|
||||
return session_manager.is_stop_requested(self.session_id)
|
||||
except Exception as e:
|
||||
print(f"Error checking Redis stop signal: {e}")
|
||||
# Fall back to local event
|
||||
return self.stop_event.is_set()
|
||||
|
||||
@ -116,22 +115,19 @@ class Scanner:
|
||||
"""
|
||||
Set stop signal both locally and in Redis.
|
||||
"""
|
||||
# Set local event
|
||||
self.stop_event.set()
|
||||
|
||||
# Set Redis signal if session ID is available
|
||||
if self.session_id:
|
||||
try:
|
||||
from core.session_manager import session_manager
|
||||
session_manager.set_stop_signal(self.session_id)
|
||||
except Exception as e:
|
||||
print(f"Error setting Redis stop signal: {e}")
|
||||
pass
|
||||
|
||||
def __getstate__(self):
|
||||
"""Prepare object for pickling by excluding unpicklable attributes."""
|
||||
state = self.__dict__.copy()
|
||||
|
||||
# Remove unpicklable threading objects
|
||||
unpicklable_attrs = [
|
||||
'stop_event',
|
||||
'scan_thread',
|
||||
@ -139,14 +135,15 @@ class Scanner:
|
||||
'processing_lock',
|
||||
'task_queue',
|
||||
'rate_limiter',
|
||||
'logger'
|
||||
'logger',
|
||||
'status_logger_thread',
|
||||
'status_logger_stop_event'
|
||||
]
|
||||
|
||||
for attr in unpicklable_attrs:
|
||||
if attr in state:
|
||||
del state[attr]
|
||||
|
||||
# Handle providers separately to ensure they're picklable
|
||||
if 'providers' in state:
|
||||
for provider in state['providers']:
|
||||
if hasattr(provider, '_stop_event'):
|
||||
@ -158,7 +155,6 @@ class Scanner:
|
||||
"""Restore object after unpickling by reconstructing threading objects."""
|
||||
self.__dict__.update(state)
|
||||
|
||||
# Reconstruct threading objects
|
||||
self.stop_event = threading.Event()
|
||||
self.scan_thread = None
|
||||
self.executor = None
|
||||
@ -166,15 +162,18 @@ class Scanner:
|
||||
self.task_queue = PriorityQueue()
|
||||
self.rate_limiter = GlobalRateLimiter(redis.StrictRedis(db=0))
|
||||
self.logger = get_forensic_logger()
|
||||
self.status_logger_thread = None
|
||||
self.status_logger_stop_event = threading.Event()
|
||||
|
||||
if not hasattr(self, 'providers') or not self.providers:
|
||||
print("Providers not found after loading session, re-initializing...")
|
||||
self._initialize_providers()
|
||||
|
||||
if not hasattr(self, 'currently_processing'):
|
||||
self.currently_processing = set()
|
||||
|
||||
# Re-set stop events for providers
|
||||
if not hasattr(self, 'currently_processing_display'):
|
||||
self.currently_processing_display = []
|
||||
|
||||
if hasattr(self, 'providers'):
|
||||
for provider in self.providers:
|
||||
if hasattr(provider, 'set_stop_event'):
|
||||
@ -183,8 +182,6 @@ class Scanner:
|
||||
def _initialize_providers(self) -> None:
|
||||
"""Initialize all available providers based on session configuration."""
|
||||
self.providers = []
|
||||
print("Initializing providers with session config...")
|
||||
|
||||
provider_dir = os.path.join(os.path.dirname(__file__), '..', 'providers')
|
||||
for filename in os.listdir(provider_dir):
|
||||
if filename.endswith('_provider.py') and not filename.startswith('base'):
|
||||
@ -202,99 +199,98 @@ class Scanner:
|
||||
if provider.is_available():
|
||||
provider.set_stop_event(self.stop_event)
|
||||
self.providers.append(provider)
|
||||
print(f"✓ {provider.get_display_name()} provider initialized successfully for session")
|
||||
else:
|
||||
print(f"✗ {provider.get_display_name()} provider is not available")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to initialize provider from {filename}: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
print(f"Initialized {len(self.providers)} providers for session")
|
||||
def _status_logger_thread(self):
|
||||
"""Periodically prints a clean, formatted scan status to the terminal."""
|
||||
# Color codes for improved display (from Document 2)
|
||||
HEADER = "\033[95m"
|
||||
CYAN = "\033[96m"
|
||||
GREEN = "\033[92m"
|
||||
YELLOW = "\033[93m"
|
||||
ENDC = "\033[0m"
|
||||
BOLD = "\033[1m"
|
||||
|
||||
def update_session_config(self, new_config) -> None:
|
||||
"""Update session configuration and reinitialize providers."""
|
||||
print("Updating session configuration...")
|
||||
self.config = new_config
|
||||
self.max_workers = self.config.max_concurrent_requests
|
||||
self._initialize_providers()
|
||||
print("Session configuration updated")
|
||||
last_status_str = ""
|
||||
while not self.status_logger_stop_event.is_set():
|
||||
try:
|
||||
# Use thread-safe copy of currently processing
|
||||
with self.processing_lock:
|
||||
in_flight_tasks = list(self.currently_processing)
|
||||
# Update display list for consistent formatting
|
||||
self.currently_processing_display = in_flight_tasks.copy()
|
||||
|
||||
status_str = (
|
||||
f"{BOLD}{HEADER}Scan Status: {self.status.upper()}{ENDC} | "
|
||||
f"{CYAN}Queued: {self.task_queue.qsize()}{ENDC} | "
|
||||
f"{YELLOW}In-Flight: {len(in_flight_tasks)}{ENDC} | "
|
||||
f"{GREEN}Completed: {self.indicators_completed}{ENDC} | "
|
||||
f"Skipped: {self.tasks_skipped} | "
|
||||
f"Rescheduled: {self.tasks_re_enqueued}"
|
||||
)
|
||||
|
||||
if status_str != last_status_str:
|
||||
print(f"\n{'-'*80}")
|
||||
print(status_str)
|
||||
if in_flight_tasks:
|
||||
print(f"{BOLD}{YELLOW}Currently Processing:{ENDC}")
|
||||
# Display up to 3 currently processing tasks
|
||||
display_tasks = [f" - {p}: {t}" for p, t in in_flight_tasks[:3]]
|
||||
print("\n".join(display_tasks))
|
||||
if len(in_flight_tasks) > 3:
|
||||
print(f" ... and {len(in_flight_tasks) - 3} more")
|
||||
print(f"{'-'*80}")
|
||||
last_status_str = status_str
|
||||
|
||||
except Exception:
|
||||
# Silently fail to avoid crashing the logger
|
||||
pass
|
||||
|
||||
time.sleep(2) # Update interval
|
||||
|
||||
def start_scan(self, target: str, max_depth: int = 2, clear_graph: bool = True, force_rescan_target: Optional[str] = None) -> bool:
|
||||
"""
|
||||
FIXED: Start a new reconnaissance scan preserving session configuration.
|
||||
Only clears graph data when requested, never destroys session/API keys.
|
||||
Starts a new reconnaissance scan.
|
||||
"""
|
||||
print(f"=== STARTING SCAN IN SCANNER {id(self)} ===")
|
||||
print(f"Session ID: {self.session_id}")
|
||||
print(f"Initial scanner status: {self.status}")
|
||||
print(f"Clear graph requested: {clear_graph}")
|
||||
print(f"Current providers: {[p.get_name() for p in self.providers]}")
|
||||
self.total_tasks_ever_enqueued = 0
|
||||
|
||||
# FIXED: Improved cleanup of previous scan without destroying session config
|
||||
if self.scan_thread and self.scan_thread.is_alive():
|
||||
print("A previous scan thread is still alive. Forcing termination...")
|
||||
|
||||
# Set stop signals immediately
|
||||
self._set_stop_signal()
|
||||
self.status = ScanStatus.STOPPED
|
||||
|
||||
# Clear all processing state
|
||||
with self.processing_lock:
|
||||
self.currently_processing.clear()
|
||||
self.currently_processing_display = []
|
||||
self.task_queue = PriorityQueue()
|
||||
|
||||
# Shutdown executor aggressively
|
||||
if self.executor:
|
||||
print("Shutting down executor forcefully...")
|
||||
self.executor.shutdown(wait=False, cancel_futures=True)
|
||||
self.executor = None
|
||||
self.scan_thread.join(5.0)
|
||||
|
||||
# Wait for thread termination with shorter timeout
|
||||
print("Waiting for previous scan thread to terminate...")
|
||||
self.scan_thread.join(5.0) # Reduced from 10 seconds
|
||||
|
||||
if self.scan_thread.is_alive():
|
||||
print("WARNING: Previous scan thread is still alive after 5 seconds")
|
||||
self.logger.logger.warning("Previous scan thread failed to terminate cleanly")
|
||||
|
||||
# FIXED: Reset scan state but preserve session configuration (API keys, etc.)
|
||||
print("Resetting scanner state for new scan (preserving session config)...")
|
||||
self.status = ScanStatus.IDLE
|
||||
self.stop_event.clear()
|
||||
|
||||
# Clear Redis stop signal explicitly
|
||||
if self.session_id:
|
||||
from core.session_manager import session_manager
|
||||
session_manager.clear_stop_signal(self.session_id)
|
||||
|
||||
with self.processing_lock:
|
||||
self.currently_processing.clear()
|
||||
self.currently_processing_display = []
|
||||
|
||||
# Reset scan-specific state but keep providers and config intact
|
||||
self.task_queue = PriorityQueue()
|
||||
self.target_retries.clear()
|
||||
self.scan_failed_due_to_retries = False
|
||||
self.tasks_skipped = 0 # BUGFIX: Reset tasks_skipped for new scan
|
||||
|
||||
# Update session state immediately for GUI feedback
|
||||
self._update_session_state()
|
||||
print("Scanner state reset complete (providers preserved).")
|
||||
|
||||
try:
|
||||
if not hasattr(self, 'providers') or not self.providers:
|
||||
print(f"ERROR: No providers available in scanner {id(self)}, cannot start scan")
|
||||
return False
|
||||
|
||||
print(f"Scanner {id(self)} validation passed, providers available: {[p.get_name() for p in self.providers]}")
|
||||
|
||||
# FIXED: Only clear graph if explicitly requested, don't destroy session
|
||||
if clear_graph:
|
||||
print("Clearing graph data (preserving session configuration)")
|
||||
self.graph.clear()
|
||||
|
||||
# Handle force rescan by clearing provider states for that specific node
|
||||
if force_rescan_target and self.graph.graph.has_node(force_rescan_target):
|
||||
print(f"Forcing rescan of {force_rescan_target}, clearing provider states.")
|
||||
node_data = self.graph.graph.nodes[force_rescan_target]
|
||||
if 'metadata' in node_data and 'provider_states' in node_data['metadata']:
|
||||
node_data['metadata']['provider_states'] = {}
|
||||
@ -307,17 +303,12 @@ class Scanner:
|
||||
self.indicators_processed = 0
|
||||
self.indicators_completed = 0
|
||||
self.tasks_re_enqueued = 0
|
||||
self.total_tasks_ever_enqueued = 0
|
||||
self.current_indicator = self.current_target
|
||||
|
||||
# Update GUI with scan preparation state
|
||||
self._update_session_state()
|
||||
|
||||
# Start new forensic session (but don't reinitialize providers)
|
||||
print(f"Starting new forensic session for scanner {id(self)}...")
|
||||
self.logger = new_session()
|
||||
|
||||
# Start scan in a separate thread
|
||||
print(f"Starting scan thread for scanner {id(self)}...")
|
||||
self.scan_thread = threading.Thread(
|
||||
target=self._execute_scan,
|
||||
args=(self.current_target, max_depth),
|
||||
@ -325,12 +316,14 @@ class Scanner:
|
||||
)
|
||||
self.scan_thread.start()
|
||||
|
||||
print(f"=== SCAN STARTED SUCCESSFULLY IN SCANNER {id(self)} ===")
|
||||
print(f"Active providers for this scan: {[p.get_name() for p in self.providers]}")
|
||||
# Start the status logger thread
|
||||
self.status_logger_stop_event.clear()
|
||||
self.status_logger_thread = threading.Thread(target=self._status_logger_thread, daemon=True)
|
||||
self.status_logger_thread.start()
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in start_scan for scanner {id(self)}: {e}")
|
||||
traceback.print_exc()
|
||||
self.status = ScanStatus.FAILED
|
||||
self._update_session_state()
|
||||
@ -347,11 +340,9 @@ class Scanner:
|
||||
|
||||
def _execute_scan(self, target: str, max_depth: int) -> None:
|
||||
"""Execute the reconnaissance scan with proper termination handling."""
|
||||
print(f"_execute_scan started for {target} with depth {max_depth}")
|
||||
self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
|
||||
processed_tasks = set()
|
||||
|
||||
# Initial task population for the main target
|
||||
is_ip = _is_valid_ip(target)
|
||||
initial_providers = self._get_eligible_providers(target, is_ip, False)
|
||||
for provider in initial_providers:
|
||||
@ -366,23 +357,19 @@ class Scanner:
|
||||
enabled_providers = [provider.get_name() for provider in self.providers]
|
||||
self.logger.log_scan_start(target, max_depth, enabled_providers)
|
||||
|
||||
# Determine initial node type
|
||||
node_type = NodeType.IP if is_ip else NodeType.DOMAIN
|
||||
self.graph.add_node(target, node_type)
|
||||
|
||||
self._initialize_provider_states(target)
|
||||
|
||||
# Better termination checking in main loop
|
||||
while not self.task_queue.empty() and not self._is_stop_requested():
|
||||
try:
|
||||
priority, (provider_name, target_item, depth) = self.task_queue.get()
|
||||
except IndexError:
|
||||
# Queue became empty during processing
|
||||
break
|
||||
|
||||
task_tuple = (provider_name, target_item)
|
||||
if task_tuple in processed_tasks:
|
||||
self.indicators_completed += 1
|
||||
self.tasks_skipped += 1
|
||||
continue
|
||||
|
||||
if depth > max_depth:
|
||||
@ -394,7 +381,6 @@ class Scanner:
|
||||
|
||||
with self.processing_lock:
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested before processing {target_item}")
|
||||
break
|
||||
self.currently_processing.add(task_tuple)
|
||||
|
||||
@ -404,7 +390,6 @@ class Scanner:
|
||||
self._update_session_state()
|
||||
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested during processing setup for {target_item}")
|
||||
break
|
||||
|
||||
provider = next((p for p in self.providers if p.get_name() == provider_name), None)
|
||||
@ -413,18 +398,15 @@ class Scanner:
|
||||
new_targets, large_entity_members, success = self._query_single_provider_for_target(provider, target_item, depth)
|
||||
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested after querying providers for {target_item}")
|
||||
break
|
||||
|
||||
if not success:
|
||||
self.target_retries[task_tuple] += 1
|
||||
if self.target_retries[task_tuple] <= self.config.max_retries_per_target:
|
||||
print(f"Re-queueing task {task_tuple} (attempt {self.target_retries[task_tuple]})")
|
||||
self.task_queue.put((priority, (provider_name, target_item, depth)))
|
||||
self.tasks_re_enqueued += 1
|
||||
self.total_tasks_ever_enqueued += 1
|
||||
else:
|
||||
print(f"ERROR: Max retries exceeded for task {task_tuple}")
|
||||
self.scan_failed_due_to_retries = True
|
||||
self._log_target_processing_error(str(task_tuple), "Max retries exceeded")
|
||||
else:
|
||||
@ -446,21 +428,14 @@ class Scanner:
|
||||
with self.processing_lock:
|
||||
self.currently_processing.discard(task_tuple)
|
||||
|
||||
if self._is_stop_requested():
|
||||
print("Scan terminated due to stop request")
|
||||
self.logger.logger.info("Scan terminated by user request")
|
||||
elif self.task_queue.empty():
|
||||
print("Scan completed - no more targets to process")
|
||||
self.logger.logger.info("Scan completed - all targets processed")
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Scan execution failed with error: {e}")
|
||||
traceback.print_exc()
|
||||
self.status = ScanStatus.FAILED
|
||||
self.logger.logger.error(f"Scan failed: {e}")
|
||||
finally:
|
||||
with self.processing_lock:
|
||||
self.currently_processing.clear()
|
||||
self.currently_processing_display = []
|
||||
|
||||
if self._is_stop_requested():
|
||||
self.status = ScanStatus.STOPPED
|
||||
@ -469,31 +444,27 @@ class Scanner:
|
||||
else:
|
||||
self.status = ScanStatus.COMPLETED
|
||||
|
||||
# Stop the status logger
|
||||
self.status_logger_stop_event.set()
|
||||
if self.status_logger_thread:
|
||||
self.status_logger_thread.join()
|
||||
|
||||
self._update_session_state()
|
||||
self.logger.log_scan_complete()
|
||||
if self.executor:
|
||||
self.executor.shutdown(wait=False, cancel_futures=True)
|
||||
self.executor = None
|
||||
stats = self.graph.get_statistics()
|
||||
print("Final scan statistics:")
|
||||
print(f" - Total nodes: {stats['basic_metrics']['total_nodes']}")
|
||||
print(f" - Total edges: {stats['basic_metrics']['total_edges']}")
|
||||
print(f" - Tasks processed: {len(processed_tasks)}")
|
||||
|
||||
def _query_single_provider_for_target(self, provider: BaseProvider, target: str, depth: int) -> Tuple[Set[str], Set[str], bool]:
|
||||
"""
|
||||
Query a single provider and process the unified ProviderResult.
|
||||
Now provider-agnostic - handles any provider that returns ProviderResult.
|
||||
"""
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested before querying {provider.get_name()} for {target}")
|
||||
return set(), set(), False
|
||||
|
||||
is_ip = _is_valid_ip(target)
|
||||
target_type = NodeType.IP if is_ip else NodeType.DOMAIN
|
||||
print(f"Querying {provider.get_name()} for {target_type.value}: {target} at depth {depth}")
|
||||
|
||||
# Ensure target node exists in graph
|
||||
self.graph.add_node(target, target_type)
|
||||
self._initialize_provider_states(target)
|
||||
|
||||
@ -502,13 +473,11 @@ class Scanner:
|
||||
provider_successful = True
|
||||
|
||||
try:
|
||||
# Query provider - now returns unified ProviderResult
|
||||
provider_result = self._query_single_provider_unified(provider, target, is_ip, depth)
|
||||
|
||||
if provider_result is None:
|
||||
provider_successful = False
|
||||
elif not self._is_stop_requested():
|
||||
# Process the unified result
|
||||
discovered, is_large_entity = self._process_provider_result_unified(
|
||||
target, provider, provider_result, depth
|
||||
)
|
||||
@ -517,8 +486,6 @@ class Scanner:
|
||||
else:
|
||||
new_targets.update(discovered)
|
||||
self.graph.process_correlations_for_node(target)
|
||||
else:
|
||||
print(f"Stop requested after processing results from {provider.get_name()}")
|
||||
except Exception as e:
|
||||
provider_successful = False
|
||||
self._log_provider_error(target, provider.get_name(), str(e))
|
||||
@ -527,58 +494,45 @@ class Scanner:
|
||||
|
||||
def _query_single_provider_unified(self, provider: BaseProvider, target: str, is_ip: bool, current_depth: int) -> Optional[ProviderResult]:
|
||||
"""
|
||||
Query a single provider with stop signal checking, now returns ProviderResult.
|
||||
Query a single provider with stop signal checking.
|
||||
"""
|
||||
provider_name = provider.get_name()
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested before querying {provider_name} for {target}")
|
||||
return None
|
||||
|
||||
print(f"Querying {provider_name} for {target}")
|
||||
|
||||
self.logger.logger.info(f"Attempting {provider_name} query for {target} at depth {current_depth}")
|
||||
|
||||
try:
|
||||
# Query the provider - returns unified ProviderResult
|
||||
if is_ip:
|
||||
result = provider.query_ip(target)
|
||||
else:
|
||||
result = provider.query_domain(target)
|
||||
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested after querying {provider_name} for {target}")
|
||||
return None
|
||||
|
||||
# Update provider state with relationship count (more meaningful than raw result count)
|
||||
relationship_count = result.get_relationship_count() if result else 0
|
||||
self._update_provider_state(target, provider_name, 'success', relationship_count, None, start_time)
|
||||
|
||||
print(f"✓ {provider_name} returned {relationship_count} relationships for {target}")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._update_provider_state(target, provider_name, 'failed', 0, str(e), start_time)
|
||||
print(f"✗ {provider_name} failed for {target}: {e}")
|
||||
return None
|
||||
|
||||
def _process_provider_result_unified(self, target: str, provider: BaseProvider,
|
||||
provider_result: ProviderResult, current_depth: int) -> Tuple[Set[str], bool]:
|
||||
"""
|
||||
Process a unified ProviderResult object to update the graph.
|
||||
Returns (discovered_targets, is_large_entity).
|
||||
"""
|
||||
provider_name = provider.get_name()
|
||||
discovered_targets = set()
|
||||
|
||||
if self._is_stop_requested():
|
||||
print(f"Stop requested before processing results from {provider_name} for {target}")
|
||||
return discovered_targets, False
|
||||
|
||||
attributes_by_node = defaultdict(list)
|
||||
for attribute in provider_result.attributes:
|
||||
# Convert the StandardAttribute object to a dictionary that the frontend can use
|
||||
attr_dict = {
|
||||
"name": attribute.name,
|
||||
"value": attribute.value,
|
||||
@ -589,10 +543,8 @@ class Scanner:
|
||||
}
|
||||
attributes_by_node[attribute.target_node].append(attr_dict)
|
||||
|
||||
# Add attributes to nodes
|
||||
for node_id, node_attributes_list in attributes_by_node.items():
|
||||
if self.graph.graph.has_node(node_id):
|
||||
# Determine node type
|
||||
if _is_valid_ip(node_id):
|
||||
node_type = NodeType.IP
|
||||
elif node_id.startswith('AS') and node_id[2:].isdigit():
|
||||
@ -600,26 +552,19 @@ class Scanner:
|
||||
else:
|
||||
node_type = NodeType.DOMAIN
|
||||
|
||||
# Add node with the list of attributes
|
||||
self.graph.add_node(node_id, node_type, attributes=node_attributes_list)
|
||||
|
||||
# Check for large entity based on relationship count
|
||||
if provider_result.get_relationship_count() > self.config.large_entity_threshold:
|
||||
print(f"Large entity detected: {provider_name} returned {provider_result.get_relationship_count()} relationships for {target}")
|
||||
members = self._create_large_entity_from_provider_result(target, provider_name, provider_result, current_depth)
|
||||
return members, True
|
||||
|
||||
# Process relationships
|
||||
for i, relationship in enumerate(provider_result.relationships):
|
||||
if i % 5 == 0 and self._is_stop_requested(): # Check periodically for stop
|
||||
print(f"Stop requested while processing relationships from {provider_name} for {target}")
|
||||
if i % 5 == 0 and self._is_stop_requested():
|
||||
break
|
||||
|
||||
# Add nodes for relationship endpoints
|
||||
source_node = relationship.source_node
|
||||
target_node = relationship.target_node
|
||||
|
||||
# Determine node types
|
||||
source_type = NodeType.IP if _is_valid_ip(source_node) else NodeType.DOMAIN
|
||||
if target_node.startswith('AS') and target_node[2:].isdigit():
|
||||
target_type = NodeType.ASN
|
||||
@ -628,11 +573,9 @@ class Scanner:
|
||||
else:
|
||||
target_type = NodeType.DOMAIN
|
||||
|
||||
# Add nodes to graph
|
||||
self.graph.add_node(source_node, source_type)
|
||||
self.graph.add_node(target_node, target_type)
|
||||
|
||||
# Add edge to graph
|
||||
if self.graph.add_edge(
|
||||
source_node, target_node,
|
||||
relationship.relationship_type,
|
||||
@ -640,9 +583,8 @@ class Scanner:
|
||||
provider_name,
|
||||
relationship.raw_data
|
||||
):
|
||||
print(f"Added relationship: {source_node} -> {target_node} ({relationship.relationship_type})")
|
||||
pass
|
||||
|
||||
# Track discovered targets for further processing
|
||||
if _is_valid_domain(target_node) or _is_valid_ip(target_node):
|
||||
discovered_targets.add(target_node)
|
||||
|
||||
@ -651,11 +593,10 @@ class Scanner:
|
||||
def _create_large_entity_from_provider_result(self, source: str, provider_name: str,
|
||||
provider_result: ProviderResult, current_depth: int) -> Set[str]:
|
||||
"""
|
||||
Create a large entity node from a ProviderResult and return the members for DNS processing.
|
||||
Create a large entity node from a ProviderResult.
|
||||
"""
|
||||
entity_id = f"large_entity_{provider_name}_{hash(source) & 0x7FFFFFFF}"
|
||||
|
||||
# Extract target nodes from relationships
|
||||
targets = [rel.target_node for rel in provider_result.relationships]
|
||||
node_type = 'unknown'
|
||||
|
||||
@ -665,7 +606,6 @@ class Scanner:
|
||||
elif _is_valid_ip(targets[0]):
|
||||
node_type = 'ip'
|
||||
|
||||
# Create nodes in graph (they exist but are grouped)
|
||||
for target in targets:
|
||||
target_node_type = NodeType.DOMAIN if node_type == 'domain' else NodeType.IP
|
||||
self.graph.add_node(target, target_node_type)
|
||||
@ -694,106 +634,76 @@ class Scanner:
|
||||
|
||||
self.graph.add_node(entity_id, NodeType.LARGE_ENTITY, attributes=attributes_list, description=description)
|
||||
|
||||
# Create edge from source to large entity
|
||||
if provider_result.relationships:
|
||||
rel_type = provider_result.relationships[0].relationship_type
|
||||
self.graph.add_edge(source, entity_id, rel_type, 0.9, provider_name,
|
||||
{'large_entity_info': f'Contains {len(targets)} {node_type}s'})
|
||||
|
||||
self.logger.logger.warning(f"Large entity created: {entity_id} contains {len(targets)} targets from {provider_name}")
|
||||
print(f"Created large entity {entity_id} for {len(targets)} {node_type}s from {provider_name}")
|
||||
|
||||
return set(targets)
|
||||
|
||||
def stop_scan(self) -> bool:
|
||||
"""Request immediate scan termination with proper cleanup."""
|
||||
try:
|
||||
print("=== INITIATING IMMEDIATE SCAN TERMINATION ===")
|
||||
self.logger.logger.info("Scan termination requested by user")
|
||||
|
||||
# **IMPROVED**: More aggressive stop signal setting
|
||||
self._set_stop_signal()
|
||||
self.status = ScanStatus.STOPPED
|
||||
|
||||
# **NEW**: Clear processing state immediately
|
||||
with self.processing_lock:
|
||||
currently_processing_copy = self.currently_processing.copy()
|
||||
self.currently_processing.clear()
|
||||
print(f"Cleared {len(currently_processing_copy)} currently processing targets: {currently_processing_copy}")
|
||||
self.currently_processing_display = []
|
||||
|
||||
# **IMPROVED**: Clear task queue and log what was discarded
|
||||
discarded_tasks = []
|
||||
while not self.task_queue.empty():
|
||||
discarded_tasks.append(self.task_queue.get())
|
||||
self.task_queue = PriorityQueue()
|
||||
print(f"Discarded {len(discarded_tasks)} pending tasks")
|
||||
|
||||
# **IMPROVED**: Aggressively shut down executor
|
||||
if self.executor:
|
||||
print("Shutting down executor with immediate cancellation...")
|
||||
try:
|
||||
# Cancel all pending futures
|
||||
self.executor.shutdown(wait=False, cancel_futures=True)
|
||||
print("Executor shutdown completed")
|
||||
except Exception as e:
|
||||
print(f"Error during executor shutdown: {e}")
|
||||
pass
|
||||
|
||||
# Immediately update GUI with stopped status
|
||||
self._update_session_state()
|
||||
|
||||
print("Termination signals sent. The scan will stop as soon as possible.")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in stop_scan: {e}")
|
||||
self.logger.logger.error(f"Error during scan termination: {e}")
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def extract_node_from_large_entity(self, large_entity_id: str, node_id_to_extract: str) -> bool:
|
||||
"""
|
||||
Extracts a node from a large entity, re-creates its original edge, and
|
||||
re-queues it for full scanning.
|
||||
Extracts a node from a large entity and re-queues it for scanning.
|
||||
"""
|
||||
if not self.graph.graph.has_node(large_entity_id):
|
||||
print(f"ERROR: Large entity {large_entity_id} not found.")
|
||||
return False
|
||||
|
||||
# 1. Get the original source node that discovered the large entity
|
||||
predecessors = list(self.graph.graph.predecessors(large_entity_id))
|
||||
if not predecessors:
|
||||
print(f"ERROR: No source node found for large entity {large_entity_id}.")
|
||||
return False
|
||||
source_node_id = predecessors[0]
|
||||
|
||||
# Get the original edge data to replicate it for the extracted node
|
||||
original_edge_data = self.graph.graph.get_edge_data(source_node_id, large_entity_id)
|
||||
if not original_edge_data:
|
||||
print(f"ERROR: Could not find original edge data from {source_node_id} to {large_entity_id}.")
|
||||
return False
|
||||
|
||||
# 2. Modify the graph data structure first
|
||||
success = self.graph.extract_node_from_large_entity(large_entity_id, node_id_to_extract)
|
||||
if not success:
|
||||
print(f"ERROR: Node {node_id_to_extract} could not be removed from {large_entity_id}'s attributes.")
|
||||
return False
|
||||
|
||||
# 3. Create the direct edge from the original source to the newly extracted node
|
||||
print(f"Re-creating direct edge from {source_node_id} to extracted node {node_id_to_extract}")
|
||||
self.graph.add_edge(
|
||||
source_id=source_node_id,
|
||||
target_id=node_id_to_extract,
|
||||
relationship_type=original_edge_data.get('relationship_type', 'extracted_from_large_entity'),
|
||||
confidence_score=original_edge_data.get('confidence_score', 0.85), # Slightly lower confidence
|
||||
confidence_score=original_edge_data.get('confidence_score', 0.85),
|
||||
source_provider=original_edge_data.get('source_provider', 'unknown'),
|
||||
raw_data={'context': f'Extracted from large entity {large_entity_id}'}
|
||||
)
|
||||
|
||||
# 4. Re-queue the extracted node for full processing by all eligible providers
|
||||
print(f"Re-queueing extracted node {node_id_to_extract} for full reconnaissance...")
|
||||
is_ip = _is_valid_ip(node_id_to_extract)
|
||||
|
||||
# FIX: Correctly retrieve discovery_depth from the list of attributes
|
||||
large_entity_attributes = self.graph.graph.nodes[large_entity_id].get('attributes', [])
|
||||
discovery_depth_attr = next((attr for attr in large_entity_attributes if attr.get('name') == 'discovery_depth'), None)
|
||||
current_depth = discovery_depth_attr['value'] if discovery_depth_attr else 0
|
||||
@ -804,9 +714,7 @@ class Scanner:
|
||||
self.task_queue.put((self._get_priority(provider_name), (provider_name, node_id_to_extract, current_depth)))
|
||||
self.total_tasks_ever_enqueued += 1
|
||||
|
||||
# 5. If the scanner is not running, we need to kickstart it to process this one item.
|
||||
if self.status != ScanStatus.RUNNING:
|
||||
print("Scanner is idle. Starting a mini-scan to process the extracted node.")
|
||||
self.status = ScanStatus.RUNNING
|
||||
self._update_session_state()
|
||||
|
||||
@ -818,25 +726,21 @@ class Scanner:
|
||||
)
|
||||
self.scan_thread.start()
|
||||
|
||||
print(f"Successfully extracted and re-queued {node_id_to_extract} from {large_entity_id}.")
|
||||
return True
|
||||
|
||||
def _update_session_state(self) -> None:
|
||||
"""
|
||||
Update the scanner state in Redis for GUI updates.
|
||||
This ensures the web interface sees real-time updates.
|
||||
"""
|
||||
if self.session_id:
|
||||
try:
|
||||
from core.session_manager import session_manager
|
||||
success = session_manager.update_session_scanner(self.session_id, self)
|
||||
if not success:
|
||||
print(f"WARNING: Failed to update session state for {self.session_id}")
|
||||
session_manager.update_session_scanner(self.session_id, self)
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to update session state: {e}")
|
||||
pass
|
||||
|
||||
def get_scan_status(self) -> Dict[str, Any]:
|
||||
"""Get current scan status with processing information."""
|
||||
"""Get current scan status with comprehensive processing information."""
|
||||
try:
|
||||
with self.processing_lock:
|
||||
currently_processing_count = len(self.currently_processing)
|
||||
@ -860,31 +764,18 @@ class Scanner:
|
||||
'currently_processing': currently_processing_list[:5],
|
||||
'tasks_in_queue': self.task_queue.qsize(),
|
||||
'tasks_completed': self.indicators_completed,
|
||||
'tasks_skipped': self.total_tasks_ever_enqueued - self.task_queue.qsize() - self.indicators_completed - self.tasks_re_enqueued,
|
||||
'tasks_skipped': self.tasks_skipped,
|
||||
'tasks_rescheduled': self.tasks_re_enqueued,
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"ERROR: Exception in get_scan_status: {e}")
|
||||
traceback.print_exc()
|
||||
return {
|
||||
'status': 'error',
|
||||
'target_domain': None,
|
||||
'current_depth': 0,
|
||||
'max_depth': 0,
|
||||
'current_indicator': '',
|
||||
'indicators_processed': 0,
|
||||
'indicators_completed': 0,
|
||||
'tasks_re_enqueued': 0,
|
||||
'progress_percentage': 0.0,
|
||||
'enabled_providers': [],
|
||||
'graph_statistics': {},
|
||||
'task_queue_size': 0,
|
||||
'currently_processing_count': 0,
|
||||
'currently_processing': [],
|
||||
'tasks_in_queue': 0,
|
||||
'tasks_completed': 0,
|
||||
'tasks_skipped': 0,
|
||||
'tasks_rescheduled': 0,
|
||||
'status': 'error', 'target_domain': None, 'current_depth': 0, 'max_depth': 0,
|
||||
'current_indicator': '', 'indicators_processed': 0, 'indicators_completed': 0,
|
||||
'tasks_re_enqueued': 0, 'progress_percentage': 0.0, 'enabled_providers': [],
|
||||
'graph_statistics': {}, 'task_queue_size': 0, 'currently_processing_count': 0,
|
||||
'currently_processing': [], 'tasks_in_queue': 0, 'tasks_completed': 0,
|
||||
'tasks_skipped': 0, 'tasks_rescheduled': 0,
|
||||
}
|
||||
|
||||
def _initialize_provider_states(self, target: str) -> None:
|
||||
@ -910,8 +801,6 @@ class Scanner:
|
||||
if provider.get_eligibility().get(target_key):
|
||||
if not self._already_queried_provider(target, provider.get_name()):
|
||||
eligible.append(provider)
|
||||
else:
|
||||
print(f"Skipping {provider.get_name()} for {target} - already queried")
|
||||
|
||||
return eligible
|
||||
|
||||
@ -923,7 +812,6 @@ class Scanner:
|
||||
node_data = self.graph.graph.nodes[target]
|
||||
provider_states = node_data.get('metadata', {}).get('provider_states', {})
|
||||
|
||||
# A provider has been successfully queried if a state exists and its status is 'success'
|
||||
provider_state = provider_states.get(provider_name)
|
||||
return provider_state is not None and provider_state.get('status') == 'success'
|
||||
|
||||
@ -947,8 +835,6 @@ class Scanner:
|
||||
'duration_ms': (datetime.now(timezone.utc) - start_time).total_seconds() * 1000
|
||||
}
|
||||
|
||||
self.logger.logger.info(f"Provider state updated: {target} -> {provider_name} -> {status} ({results_count} results)")
|
||||
|
||||
def _log_target_processing_error(self, target: str, error: str) -> None:
|
||||
"""Log target processing errors for forensic trail."""
|
||||
self.logger.logger.error(f"Target processing failed for {target}: {error}")
|
||||
@ -957,11 +843,6 @@ class Scanner:
|
||||
"""Log provider query errors for forensic trail."""
|
||||
self.logger.logger.error(f"Provider {provider_name} failed for {target}: {error}")
|
||||
|
||||
def _log_no_eligible_providers(self, target: str, is_ip: bool) -> None:
|
||||
"""Log when no providers are eligible for a target."""
|
||||
target_type = 'IP' if is_ip else 'domain'
|
||||
self.logger.logger.warning(f"No eligible providers for {target_type}: {target}")
|
||||
|
||||
def _calculate_progress(self) -> float:
|
||||
"""Calculate scan progress percentage based on task completion."""
|
||||
if self.total_tasks_ever_enqueued == 0:
|
||||
@ -996,13 +877,6 @@ class Scanner:
|
||||
}
|
||||
return export_data
|
||||
|
||||
def get_provider_statistics(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get statistics for all providers with forensic information."""
|
||||
stats = {}
|
||||
for provider in self.providers:
|
||||
stats[provider.get_name()] = provider.get_statistics()
|
||||
return stats
|
||||
|
||||
def get_provider_info(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get information about all available providers."""
|
||||
info = {}
|
||||
@ -1016,11 +890,9 @@ class Scanner:
|
||||
attribute = getattr(module, attribute_name)
|
||||
if isinstance(attribute, type) and issubclass(attribute, BaseProvider) and attribute is not BaseProvider:
|
||||
provider_class = attribute
|
||||
# Instantiate to get metadata, even if not fully configured
|
||||
temp_provider = provider_class(name=attribute_name, session_config=self.config)
|
||||
provider_name = temp_provider.get_name()
|
||||
|
||||
# Find the actual provider instance if it exists, to get live stats
|
||||
live_provider = next((p for p in self.providers if p.get_name() == provider_name), None)
|
||||
|
||||
info[provider_name] = {
|
||||
@ -1031,6 +903,5 @@ class Scanner:
|
||||
'rate_limit': self.config.get_rate_limit(provider_name),
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to get info for provider from {filename}: {e}")
|
||||
traceback.print_exc()
|
||||
return info
|
||||
Loading…
x
Reference in New Issue
Block a user