dnsrecon/core/session_manager.py

# dnsrecon/core/session_manager.py

import threading
import time
import uuid
import redis
import pickle
from typing import Dict, Optional, Any

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 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 with a Redis backend.
        """
        self.redis_client = redis.StrictRedis(db=0, decode_responses=False)
        self.session_timeout = session_timeout_minutes * 60  # Convert to seconds
        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 Redis backend and {session_timeout_minutes}min timeout")
    
    def __getstate__(self):
        """Prepare SessionManager for pickling."""
        state = self.__dict__.copy()
        # Exclude unpickleable attributes - Redis client and threading objects
        unpicklable_attrs = ['lock', 'cleanup_thread', 'redis_client']
        for attr in unpicklable_attrs:
            if attr in state:
                del state[attr]
        return state

    def __setstate__(self, state):
        """Restore SessionManager after unpickling."""
        self.__dict__.update(state)
        # Re-initialize unpickleable attributes
        import redis
        self.redis_client = redis.StrictRedis(db=0, decode_responses=False)
        self.lock = threading.Lock()
        self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self.cleanup_thread.start()

    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 and store it in Redis.
        """
        session_id = str(uuid.uuid4())
        print(f"=== CREATING SESSION {session_id} IN REDIS ===")
        
        try:
            from core.session_config import create_session_config
            session_config = create_session_config()
            scanner_instance = Scanner(session_config=session_config)
            
            session_data = {
                'scanner': scanner_instance,
                'config': session_config,
                'created_at': time.time(),
                'last_activity': time.time(),
                'status': 'active'
            }
            
            # Serialize the entire session data dictionary using pickle
            serialized_data = pickle.dumps(session_data)
            
            # 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_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 update_session_scanner(self, session_id: str, scanner: 'Scanner'):
        """Updates just the scanner object in a session."""
        session_data = self._get_session_data(session_id)
        if session_data:
            session_data['scanner'] = scanner
            # We don't need to update last_activity here, as that's for user interaction
            self._save_session_data(session_id, session_data)

    def get_session(self, session_id: str) -> Optional[Scanner]:
        """
        Get scanner instance for a session from Redis.
        """
        if not session_id:
            return None
            
        session_data = self._get_session_data(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 in Redis.
        """
        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)