dnsrecon/app.py

"""
Flask application entry point for DNSRecon web interface.
Enhanced with user session management and task-based completion model.
"""

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, UserIdentifier
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():
    """
    Enhanced user scanner retrieval with user identification and session consolidation.
    Implements single session per user with seamless consolidation.
    """
    print("=== ENHANCED GET_USER_SCANNER ===")
    
    try:
        # Extract user identification from request
        client_ip, user_agent = UserIdentifier.extract_request_info(request)
        user_fingerprint = UserIdentifier.generate_user_fingerprint(client_ip, user_agent)
        
        print(f"User fingerprint: {user_fingerprint}")
        print(f"Client IP: {client_ip}")
        print(f"User Agent: {user_agent[:50]}...")
        
        # Get current Flask session info for debugging
        current_flask_session_id = session.get('dnsrecon_session_id')
        print(f"Flask session ID: {current_flask_session_id}")
        
        # Try to get existing session first
        if current_flask_session_id:
            existing_scanner = session_manager.get_session(current_flask_session_id)
            if existing_scanner:
                # Verify session belongs to current user
                session_info = session_manager.get_session_info(current_flask_session_id)
                if session_info.get('user_fingerprint') == user_fingerprint:
                    print(f"Found valid existing session {current_flask_session_id} for user {user_fingerprint}")
                    existing_scanner.session_id = current_flask_session_id
                    return current_flask_session_id, existing_scanner
                else:
                    print(f"Session {current_flask_session_id} belongs to different user, will create new session")
            else:
                print(f"Session {current_flask_session_id} not found in Redis, will create new session")
        
        # Create or replace user session (this handles consolidation automatically)
        new_session_id = session_manager.create_or_replace_user_session(client_ip, user_agent)
        new_scanner = session_manager.get_session(new_session_id)
        
        if not new_scanner:
            print(f"ERROR: Failed to retrieve newly created session {new_session_id}")
            raise Exception("Failed to create new scanner session")
        
        # Store in Flask session for browser persistence
        session['dnsrecon_session_id'] = new_session_id
        session.permanent = True
        
        # Ensure session ID is set on scanner
        new_scanner.session_id = new_session_id
        
        # Get session info for user feedback
        session_info = session_manager.get_session_info(new_session_id)
        
        print(f"Session created/consolidated successfully")
        print(f"  - Session ID: {new_session_id}")
        print(f"  - User: {user_fingerprint}")
        print(f"  - Scanner status: {new_scanner.status}")
        print(f"  - Session age: {session_info.get('session_age_minutes', 0)} minutes")
        
        return new_session_id, new_scanner
        
    except Exception as e:
        print(f"ERROR: Exception in get_user_scanner: {e}")
        traceback.print_exc()
        raise


@app.route('/')
def index():
    """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 with enhanced user session management.
    """
    print("=== API: /api/scan/start called ===")
    
    try:
        print("Getting JSON data from request...")
        data = request.get_json()
        print(f"Request data: {data}")
        
        if not data or 'target_domain' not in data:
            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}, clear_graph: {clear_graph}")
        
        # 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 session management
        user_session_id, scanner = get_user_scanner()
        
        # Ensure session ID is properly set
        if not scanner.session_id:
            scanner.session_id = user_session_id
        
        print(f"Using session: {user_session_id}")
        print(f"Scanner object ID: {id(scanner)}")
        
        # Start scan
        print(f"Calling start_scan on scanner {id(scanner)}...")
        success = scanner.start_scan(target_domain, max_depth, clear_graph=clear_graph)
        
        # Immediately update session state regardless of success
        session_manager.update_session_scanner(user_session_id, scanner)
        
        if success:
            scan_session_id = scanner.logger.session_id
            print(f"Scan started successfully with scan session ID: {scan_session_id}")
            
            # Get session info for user feedback
            session_info = session_manager.get_session_info(user_session_id)
            
            return jsonify({
                'success': True,
                'message': 'Scan started successfully',
                'scan_id': scan_session_id,
                'user_session_id': user_session_id,
                'scanner_status': scanner.status,
                'session_info': {
                    'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
                    'session_age_minutes': session_info.get('session_age_minutes', 0),
                    'consolidated': session_info.get('session_age_minutes', 0) > 0
                },
                'debug_info': {
                    'scanner_object_id': id(scanner),
                    'scanner_status': scanner.status
                }
            })
        else:
            print("ERROR: Scanner returned False")
            
            # Provide more detailed error information
            error_details = {
                'scanner_status': scanner.status,
                'scanner_object_id': id(scanner),
                'session_id': user_session_id,
                'providers_count': len(scanner.providers) if hasattr(scanner, 'providers') else 0
            }
            
            return jsonify({
                'success': False,
                'error': f'Failed to start scan (scanner status: {scanner.status})',
                'debug_info': error_details
            }), 409
        
    except Exception as e:
        print(f"ERROR: Exception in start_scan endpoint: {e}")
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': f'Internal server error: {str(e)}'
        }), 500


@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 enhanced session information."""
    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
        
        # Add enhanced session information
        session_info = session_manager.get_session_info(user_session_id)
        status['session_info'] = {
            'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
            'session_age_minutes': session_info.get('session_age_minutes', 0),
            'client_ip': session_info.get('client_ip', 'unknown'),
            'last_activity': session_info.get('last_activity')
        }
        
        # Additional debug info
        status['debug_info'] = {
            'scanner_object_id': id(scanner),
            '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/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 enhanced session information to export
        session_info = session_manager.get_session_info(user_session_id)
        results['export_metadata'] = {
            'user_session_id': user_session_id,
            'user_fingerprint': session_info.get('user_fingerprint', 'unknown'),
            'client_ip': session_info.get('client_ip', 'unknown'),
            'session_age_minutes': session_info.get('session_age_minutes', 0),
            'export_timestamp': datetime.now(timezone.utc).isoformat(),
            'export_type': 'user_session_results'
        }
        
        # Create filename with user fingerprint
        timestamp = datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')
        target = scanner.current_target or 'unknown'
        user_fp = session_info.get('user_fingerprint', 'unknown')[:8]
        filename = f"dnsrecon_{target}_{timestamp}_{user_fp}.json"
        
        # Create in-memory file
        json_data = json.dumps(results, indent=2, ensure_ascii=False)
        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_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.route('/api/session/info', methods=['GET'])
def get_session_info():
    """Get enhanced information about the current user session."""
    try:
        user_session_id, scanner = get_user_scanner()
        session_info = session_manager.get_session_info(user_session_id)
        
        return jsonify({
            'success': True,
            'session_info': session_info
        })
    
    except Exception as e:
        print(f"ERROR: Exception in get_session_info endpoint: {e}")
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': f'Internal server error: {str(e)}'
        }), 500


@app.route('/api/session/terminate', methods=['POST'])
def terminate_session():
    """Terminate the current user session."""
    try:
        user_session_id = session.get('dnsrecon_session_id')
        
        if user_session_id:
            success = session_manager.terminate_session(user_session_id)
            # Clear Flask session
            session.pop('dnsrecon_session_id', None)
            
            return jsonify({
                'success': success,
                'message': 'Session terminated' if success else 'Session not found'
            })
        else:
            return jsonify({
                'success': False,
                'error': 'No active session to terminate'
            }), 400
    
    except Exception as e:
        print(f"ERROR: Exception in terminate_session endpoint: {e}")
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': f'Internal server error: {str(e)}'
        }), 500


@app.route('/api/admin/sessions', methods=['GET'])
def list_sessions():
    """Admin endpoint to list all active sessions with enhanced information."""
    try:
        sessions = session_manager.list_active_sessions()
        stats = session_manager.get_statistics()
        
        return jsonify({
            'success': True,
            'sessions': sessions,
            'statistics': stats
        })
    
    except Exception as e:
        print(f"ERROR: Exception in list_sessions endpoint: {e}")
        traceback.print_exc()
        return jsonify({
            'success': False,
            'error': f'Internal server error: {str(e)}'
        }), 500


@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint with enhanced session statistics."""
    try:
        # Get session stats
        session_stats = session_manager.get_statistics()
        
        return jsonify({
            'success': True,
            'status': 'healthy',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'version': '2.0.0-enhanced',
            'phase': 'enhanced_architecture',
            'features': {
                'multi_provider': True,
                'concurrent_processing': True,
                'real_time_updates': True,
                'api_key_management': True,
                'visualization': True,
                'retry_logic': True,
                'user_sessions': True,
                'session_isolation': True,
                'global_provider_caching': True,
                'single_session_per_user': True,
                'session_consolidation': True,
                'task_completion_model': True
            },
            'session_statistics': session_stats,
            'cache_info': {
                'global_provider_cache': True,
                'cache_location': '.cache/<provider_name>/',
                'cache_expiry_hours': 12
            }
        })
    except Exception as e:
        print(f"ERROR: Exception in health_check endpoint: {e}")
        return jsonify({
            'success': False,
            'error': f'Health check failed: {str(e)}'
        }), 500


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