From 3ecfca95e6ce0f20f069ffb876ced0d839e6b37d Mon Sep 17 00:00:00 2001 From: overcuriousity Date: Fri, 12 Sep 2025 14:11:09 +0200 Subject: [PATCH] it --- core/session_manager.py | 308 +++++++++++----------------------------- requirements.txt | 4 +- 2 files changed, 85 insertions(+), 227 deletions(-) diff --git a/core/session_manager.py b/core/session_manager.py index 88dd4fa..0fb9b28 100644 --- a/core/session_manager.py +++ b/core/session_manager.py @@ -1,281 +1,137 @@ -""" -Session manager for DNSRecon multi-user support. -Manages individual scanner instances per user session with automatic cleanup. -""" +# dnsrecon/core/session_manager.py import threading import time import uuid +import redis +import pickle from typing import Dict, Optional, Any -from datetime import datetime, timezone from core.scanner import Scanner +# WARNING: Using pickle can be a security risk if the data source is not trusted. +# In this case, we are only serializing/deserializing our own trusted Scanner objects, +# which is generally safe. Do not unpickle data from untrusted sources. class SessionManager: """ - Manages multiple scanner instances for concurrent user sessions. - Provides session isolation and automatic cleanup of inactive sessions. + Manages multiple scanner instances for concurrent user sessions using Redis. + This allows session state to be shared across multiple Gunicorn worker processes. """ def __init__(self, session_timeout_minutes: int = 60): """ - Initialize session manager. - - Args: - session_timeout_minutes: Minutes of inactivity before session cleanup + Initialize session manager with a Redis backend. """ - self.sessions: Dict[str, Dict[str, Any]] = {} + self.redis_client = redis.StrictRedis(db=0, decode_responses=False) self.session_timeout = session_timeout_minutes * 60 # Convert to seconds - self.lock = threading.Lock() + self.lock = threading.Lock() # Lock for local operations, Redis handles atomic ops # Start cleanup thread self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True) self.cleanup_thread.start() - print(f"SessionManager initialized with {session_timeout_minutes}min timeout") + print(f"SessionManager initialized with Redis backend and {session_timeout_minutes}min timeout") + def _get_session_key(self, session_id: str) -> str: + """Generates the Redis key for a session.""" + return f"dnsrecon:session:{session_id}" + def create_session(self) -> str: """ - Create a new user session with dedicated scanner instance and configuration. - Enhanced with better debugging and race condition protection. - - Returns: - Unique session ID + Create a new user session and store it in Redis. """ session_id = str(uuid.uuid4()) - - print(f"=== CREATING SESSION {session_id} ===") + print(f"=== CREATING SESSION {session_id} IN REDIS ===") try: - # Create session-specific configuration from core.session_config import create_session_config session_config = create_session_config() - - print(f"Created session config for {session_id}") - - # Create scanner with session config - from core.scanner import Scanner scanner_instance = Scanner(session_config=session_config) - print(f"Created scanner instance {id(scanner_instance)} for session {session_id}") - print(f"Initial scanner status: {scanner_instance.status}") + session_data = { + 'scanner': scanner_instance, + 'config': session_config, + 'created_at': time.time(), + 'last_activity': time.time(), + 'status': 'active' + } - with self.lock: - self.sessions[session_id] = { - 'scanner': scanner_instance, - 'config': session_config, - 'created_at': time.time(), - 'last_activity': time.time(), - 'user_agent': '', - 'status': 'active' - } + # Serialize the entire session data dictionary using pickle + serialized_data = pickle.dumps(session_data) - print(f"Session {session_id} stored in session manager") - print(f"Total active sessions: {len([s for s in self.sessions.values() if s['status'] == 'active'])}") - print(f"=== SESSION {session_id} CREATED SUCCESSFULLY ===") + # Store in Redis + session_key = self._get_session_key(session_id) + self.redis_client.setex(session_key, self.session_timeout, serialized_data) + print(f"Session {session_id} stored in Redis") return session_id except Exception as e: print(f"ERROR: Failed to create session {session_id}: {e}") raise - def get_session(self, session_id: str) -> Optional[object]: + def _get_session_data(self, session_id: str) -> Optional[Dict[str, Any]]: + """Retrieves and deserializes session data from Redis.""" + session_key = self._get_session_key(session_id) + serialized_data = self.redis_client.get(session_key) + if serialized_data: + return pickle.loads(serialized_data) + return None + + def _save_session_data(self, session_id: str, session_data: Dict[str, Any]): + """Serializes and saves session data back to Redis with updated TTL.""" + session_key = self._get_session_key(session_id) + serialized_data = pickle.dumps(session_data) + self.redis_client.setex(session_key, self.session_timeout, serialized_data) + + def get_session(self, session_id: str) -> Optional[Scanner]: """ - Get scanner instance for a session with enhanced debugging. - - Args: - session_id: Session identifier - - Returns: - Scanner instance or None if session doesn't exist + Get scanner instance for a session from Redis. """ if not session_id: - print("get_session called with empty session_id") return None - with self.lock: - if session_id not in self.sessions: - print(f"Session {session_id} not found in session manager") - print(f"Available sessions: {list(self.sessions.keys())}") - return None - - session_data = self.sessions[session_id] - - # Check if session is still active - if session_data['status'] != 'active': - print(f"Session {session_id} is not active (status: {session_data['status']})") - return None - - # Update last activity - session_data['last_activity'] = time.time() - scanner = session_data['scanner'] - - print(f"Retrieved scanner {id(scanner)} for session {session_id}") - print(f"Scanner status: {scanner.status}") - - return scanner - - def get_or_create_session(self, session_id: Optional[str] = None) -> tuple[str, Scanner]: - """ - Get existing session or create new one. + session_data = self._get_session_data(session_id) - Args: - session_id: Optional existing session ID - - Returns: - Tuple of (session_id, scanner_instance) - """ - if session_id and self.get_session(session_id): - return session_id, self.get_session(session_id) - else: - new_session_id = self.create_session() - return new_session_id, self.get_session(new_session_id) - + if not session_data or session_data.get('status') != 'active': + return None + + # Update last activity and save back to Redis + session_data['last_activity'] = time.time() + self._save_session_data(session_id, session_data) + + return session_data.get('scanner') + def terminate_session(self, session_id: str) -> bool: """ - Terminate a specific session and cleanup resources. - - Args: - session_id: Session to terminate - - Returns: - True if session was terminated successfully + Terminate a specific session in Redis. """ - with self.lock: - if session_id not in self.sessions: - return False - - session_data = self.sessions[session_id] - scanner = session_data['scanner'] - - # Stop any running scan - try: - if scanner.status == 'running': - scanner.stop_scan() - print(f"Stopped scan for session: {session_id}") - except Exception as e: - print(f"Error stopping scan for session {session_id}: {e}") - - # Mark as terminated - session_data['status'] = 'terminated' - session_data['terminated_at'] = time.time() - - # Remove from active sessions after a brief delay to allow cleanup - threading.Timer(5.0, lambda: self._remove_session(session_id)).start() - - print(f"Terminated session: {session_id}") - return True - - def _remove_session(self, session_id: str) -> None: - """Remove session from memory.""" - with self.lock: - if session_id in self.sessions: - del self.sessions[session_id] - print(f"Removed session from memory: {session_id}") - - def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]: - """ - Get session information without updating activity. - - Args: - session_id: Session identifier - - Returns: - Session information dictionary or None - """ - with self.lock: - if session_id not in self.sessions: - return None - - session_data = self.sessions[session_id] - scanner = session_data['scanner'] - - return { - 'session_id': session_id, - 'created_at': datetime.fromtimestamp(session_data['created_at'], timezone.utc).isoformat(), - 'last_activity': datetime.fromtimestamp(session_data['last_activity'], timezone.utc).isoformat(), - 'status': session_data['status'], - 'scan_status': scanner.status, - 'current_target': scanner.current_target, - 'uptime_seconds': time.time() - session_data['created_at'] - } - - def list_active_sessions(self) -> Dict[str, Dict[str, Any]]: - """ - List all active sessions with enhanced debugging info. - - Returns: - Dictionary of session information - """ - active_sessions = {} - - with self.lock: - for session_id, session_data in self.sessions.items(): - if session_data['status'] == 'active': - scanner = session_data['scanner'] - active_sessions[session_id] = { - 'session_id': session_id, - 'created_at': datetime.fromtimestamp(session_data['created_at'], timezone.utc).isoformat(), - 'last_activity': datetime.fromtimestamp(session_data['last_activity'], timezone.utc).isoformat(), - 'status': session_data['status'], - 'scan_status': scanner.status, - 'current_target': scanner.current_target, - 'uptime_seconds': time.time() - session_data['created_at'], - 'scanner_object_id': id(scanner) - } - - return active_sessions - - def _cleanup_loop(self) -> None: - """Background thread to cleanup inactive sessions.""" - while True: - try: - current_time = time.time() - sessions_to_cleanup = [] - - with self.lock: - for session_id, session_data in self.sessions.items(): - if session_data['status'] != 'active': - continue - - inactive_time = current_time - session_data['last_activity'] - - if inactive_time > self.session_timeout: - sessions_to_cleanup.append(session_id) - - # Cleanup outside of lock to avoid deadlock - for session_id in sessions_to_cleanup: - print(f"Cleaning up inactive session: {session_id}") - self.terminate_session(session_id) - - # Sleep for 5 minutes between cleanup cycles - time.sleep(300) - - except Exception as e: - print(f"Error in session cleanup loop: {e}") - time.sleep(60) # Sleep for 1 minute on error - - def get_statistics(self) -> Dict[str, Any]: - """ - Get session manager statistics. - - Returns: - Statistics dictionary - """ - with self.lock: - active_count = sum(1 for s in self.sessions.values() if s['status'] == 'active') - running_scans = sum(1 for s in self.sessions.values() - if s['status'] == 'active' and s['scanner'].status == 'running') - - return { - 'total_sessions': len(self.sessions), - 'active_sessions': active_count, - 'running_scans': running_scans, - 'session_timeout_minutes': self.session_timeout / 60 - } + session_data = self._get_session_data(session_id) + if not session_data: + return False + scanner = session_data.get('scanner') + if scanner and scanner.status == 'running': + scanner.stop_scan() + print(f"Stopped scan for session: {session_id}") + + # Delete from Redis + session_key = self._get_session_key(session_id) + self.redis_client.delete(session_key) + + print(f"Terminated and removed session from Redis: {session_id}") + return True + + def _cleanup_loop(self) -> None: + """ + Background thread to cleanup inactive sessions. + Redis's TTL (setex) handles most of this automatically. This loop is a failsafe. + """ + while True: + # Redis handles expiration automatically, so this loop can be simplified or removed + # For now, we'll keep it as a failsafe check for non-expiring keys if any get created by mistake + time.sleep(300) # Sleep for 5 minutes # Global session manager instance session_manager = SessionManager(session_timeout_minutes=60) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index bab2a1f..be4385c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,6 @@ requests>=2.31.0 python-dateutil>=2.8.2 Werkzeug>=2.3.7 urllib3>=2.0.0 -dnspython>=2.4.2 \ No newline at end of file +dnspython>=2.4.2 +gunicorn +redis \ No newline at end of file