dnsrecon/app.py

# 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.
"""

import json
import traceback
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 config import config
from core.graph_manager import NodeType
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 and scanner if one doesn't exist.
    """
    # Get current Flask session info for debugging
    current_flask_session_id = session.get('dnsrecon_session_id')
    
    # 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

@app.route('/')
def index():
    """Serve the main web interface."""
    return render_template('index.html')


@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.
    """
    print("=== API: /api/scan/start called ===")
    
    try:
        data = request.get_json()
        if not data or 'target' not in data:
            return jsonify({'success': False, 'error': 'Missing target in request'}), 400
        
        target = data['target'].strip()
        max_depth = data.get('max_depth', config.default_recursion_depth)
        clear_graph = data.get('clear_graph', True)
        force_rescan_target = data.get('force_rescan_target', None) # **FIX**: Get the new parameter
        
        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
        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
        
        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
        
        print(f"Using scanner {id(scanner)} in session {user_session_id}")
        
        success = scanner.start_scan(target, max_depth, clear_graph=clear_graph, force_rescan_target=force_rescan_target) # **FIX**: Pass the new parameter
        
        if success:
            return jsonify({
                'success': True,
                'message': 'Scan started successfully',
                'scan_id': scanner.logger.session_id,
                'user_session_id': user_session_id,
            })
        else:
            return jsonify({
                'success': False,
                'error': f'Failed to start scan (scanner status: {scanner.status})',
            }), 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 ===")
    
    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
        
        # 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
        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'
        })
        
    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


@app.route('/api/scan/status', methods=['GET'])
def get_scan_status():
    """Get current scan status with error handling."""
    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': {},
                    '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())
        }
        
        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
            }
        }), 500



@app.route('/api/graph', methods=['GET'])
def get_graph_data():
    """Get current graph data with error handling."""
    try:
        # Get user-specific scanner
        user_session_id, scanner = get_user_scanner()
        
        if not scanner:
            # Return empty graph if no scanner
            return jsonify({
                'success': True,
                'graph': {
                    'nodes': [],
                    'edges': [],
                    'statistics': {
                        'node_count': 0,
                        'edge_count': 0,
                        'creation_time': datetime.now(timezone.utc).isoformat(),
                        'last_modified': datetime.now(timezone.utc).isoformat()
                    }
                },
                'user_session_id': user_session_id
            })
        
        graph_data = scanner.get_graph_data()
        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()
        return jsonify({
            'success': False,
            'error': f'Internal server error: {str(e)}',
            'fallback_graph': {
                'nodes': [],
                'edges': [],
                'statistics': {'node_count': 0, 'edge_count': 0}
            }
        }), 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."""
    try:
        user_session_id, scanner = get_user_scanner()
        if not scanner:
            return jsonify({'success': False, 'error': 'No active session found'}), 404

        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

    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


@app.route('/api/graph/revert', methods=['POST'])
def revert_graph_action():
    """Reverts a graph action, such as re-adding a deleted node."""
    try:
        data = request.get_json()
        if not data or 'type' not in data or 'data' not in data:
            return jsonify({'success': False, 'error': 'Invalid revert request format'}), 400

        user_session_id, scanner = get_user_scanner()
        if not scanner:
            return jsonify({'success': False, 'error': 'No active session found'}), 404

        action_type = data['type']
        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(
                    node_id=node_to_add['id'],
                    node_type=NodeType(node_to_add['type']),
                    attributes=node_to_add.get('attributes'),
                    description=node_to_add.get('description'),
                    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'],
                        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."""
    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"
        
        # Create in-memory file
        json_data = json.dumps(results, indent=2, ensure_ascii=False)
        file_obj = io.BytesIO(json_data.encode('utf-8'))
        
        return send_file(
            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


@app.route('/api/providers', methods=['GET'])
def get_providers():
    """Get information about available providers for the user session."""
    
    try:
        # Get user-specific scanner
        user_session_id, scanner = get_user_scanner()
        
        if scanner:
            # Updated debug print to be consistent with the new progress bar logic
            completed_tasks = scanner.indicators_completed
            total_tasks = scanner.total_tasks_ever_enqueued
            print(f"DEBUG: Task Progress - Completed: {completed_tasks}, Total Enqueued: {total_tasks}")
        else:
            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
        })
    
    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


@app.route('/api/config/api-keys', methods=['POST'])
def set_api_keys():
    """
    Set API keys for providers for the user session only.
    """
    try:
        data = request.get_json()
        
        if data is None:
            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
        
        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)
        
        if updated_providers:
            # Reinitialize scanner providers to apply the new keys
            scanner._initialize_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
            })
        else:
            return jsonify({
                'success': False,
                'error': 'No valid API keys were provided or provider names were incorrect.'
            }), 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

@app.errorhandler(404)
def not_found(error):
    """Handle 404 errors."""
    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


if __name__ == '__main__':
    print("Starting DNSRecon Flask application with user session support...")
    
    # 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,
        debug=config.flask_debug,
        threaded=True
    )