dnsrecon/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 better session debugging and isolation.
"""

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


app = Flask(__name__)
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():
    """
    Get or create scanner instance for current user session with enhanced debugging.
    
    Returns:
        Tuple of (session_id, scanner_instance)
    """
    # Get current Flask session info for debugging
    current_flask_session_id = session.get('dnsrecon_session_id')
    client_ip = request.remote_addr
    user_agent = request.headers.get('User-Agent', '')[:100]  # Truncate for logging
    
    print("=== SESSION DEBUG ===")
    print(f"Client IP: {client_ip}")
    print(f"User Agent: {user_agent}")
    print(f"Flask Session ID: {current_flask_session_id}")
    print(f"Flask Session Keys: {list(session.keys())}")
    
    # Try to get existing session
    if current_flask_session_id:
        existing_scanner = session_manager.get_session(current_flask_session_id)
        if existing_scanner:
            print(f"Using existing session: {current_flask_session_id}")
            print(f"Scanner status: {existing_scanner.status}")
            return current_flask_session_id, existing_scanner
        else:
            print(f"Session {current_flask_session_id} not found in session manager")
    
    # Create new session
    print("Creating new session...")
    new_session_id = session_manager.create_session()
    new_scanner = session_manager.get_session(new_session_id)
    
    # Store in Flask session
    session['dnsrecon_session_id'] = new_session_id
    session.permanent = True
    
    print(f"Created new session: {new_session_id}")
    print(f"New scanner status: {new_scanner.status}")
    print("=== END SESSION DEBUG ===")
    
    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 for the current user session.
    Enhanced with better error handling and debugging.
    """
    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:
            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)
        clear_graph = data.get('clear_graph', True)
        
        print(f"Parsed - target_domain: '{target_domain}', max_depth: {max_depth}")
        
        # Validation
        if not target_domain:
            print("ERROR: Target domain cannot be empty")
            return jsonify({
                'success': False,
                'error': 'Target domain cannot be empty'
            }), 400
        
        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("Validation passed, getting user scanner...")
        
        # Get user-specific scanner with enhanced debugging
        user_session_id, scanner = get_user_scanner()
        print(f"Using session: {user_session_id}")
        print(f"Scanner object ID: {id(scanner)}")
        print(f"Scanner status before start: {scanner.status}")
        
        # Additional safety check - if scanner is somehow in running state, force reset
        if scanner.status == 'running':
            print(f"WARNING: Scanner in session {user_session_id} was already running - forcing reset")
            scanner.stop_scan()
            # Give it a moment to stop
            import time
            time.sleep(1)
            
            # If still running, force status reset
            if scanner.status == 'running':
                print("WARNING: Force resetting scanner status from 'running' to 'idle'")
                scanner.status = 'idle'
        
        # Start scan
        print(f"Calling start_scan on scanner {id(scanner)}...")
        success = scanner.start_scan(target_domain, max_depth, clear_graph=clear_graph)
        
        print(f"scanner.start_scan returned: {success}")
        print(f"Scanner status after start attempt: {scanner.status}")
        
        if success:
            scan_session_id = scanner.logger.session_id
            print(f"Scan started successfully with scan session ID: {scan_session_id}")
            return jsonify({
                'success': True,
                'message': 'Scan started successfully',
                'scan_id': scan_session_id,
                'user_session_id': user_session_id,
                '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

@app.route('/api/scan/stop', methods=['POST'])
def stop_scan():
    """Stop the current scan for the user session."""
    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}")
        
        success = scanner.stop_scan()
        
        if success:
            return jsonify({
                'success': True,
                'message': 'Scan stop requested',
                'user_session_id': user_session_id
            })
        else:
            return jsonify({
                'success': False,
                'error': 'No active scan to stop for this session'
            }), 400
        
    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 and progress for the user session."""
    try:
        # Get user-specific scanner
        user_session_id, scanner = get_user_scanner()
        
        status = scanner.get_scan_status()
        status['user_session_id'] = user_session_id
        
        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)}'
        }), 500


@app.route('/api/graph', methods=['GET'])
def get_graph_data():
    """Get current graph data for visualization for the user session."""
    try:
        # Get user-specific scanner
        user_session_id, scanner = get_user_scanner()
        
        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)}'
        }), 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."""
    print("=== API: /api/providers called ===")
    
    try:
        # Get user-specific scanner
        user_session_id, scanner = get_user_scanner()
        
        provider_stats = scanner.get_provider_statistics()
        
        # Add configuration information
        provider_info = {}
        for provider_name, stats in provider_stats.items():
            provider_info[provider_name] = {
                'statistics': stats,
                'enabled': config.is_provider_enabled(provider_name),
                'rate_limit': config.get_rate_limit(provider_name),
                'requires_api_key': provider_name in ['shodan', 'virustotal']
            }
        
        print(f"Returning provider info for session {user_session_id}: {list(provider_info.keys())}")
        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 not data:
            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 = []
        
        for provider, api_key in data.items():
            if provider in ['shodan', 'virustotal'] and api_key.strip():
                success = session_config.set_api_key(provider, api_key.strip())
                if success:
                    updated_providers.append(provider)
        
        if updated_providers:
            # Reinitialize scanner providers for this session only
            scanner._initialize_providers()
            
            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'
            }), 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
    
    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.route('/api/session/info', methods=['GET'])
def get_session_info():
    """Get 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."""
    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 Phase 2 information."""
    try:
        # Get session stats
        session_stats = session_manager.get_statistics()
        
        return jsonify({
            'success': True,
            'status': 'healthy',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'version': '1.0.0-phase2',
            'phase': 2,
            'features': {
                'multi_provider': True,
                'concurrent_processing': True,
                'real_time_updates': True,
                'api_key_management': True,
                'enhanced_visualization': True,
                'retry_logic': True,
                'user_sessions': True,
                'session_isolation': True
            },
            'session_statistics': session_stats
        })
    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."""
    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
    )