"""Session manager for per-device access control. Supports multiple concurrent sessions, one per PM3 device. Each device can have at most one active session at a time. """ import asyncio import time from typing import Optional, Dict from dataclasses import dataclass import uuid from .. import config @dataclass class Session: """Active session information.""" session_id: str device_id: Optional[str] # None for legacy single-device mode client_ip: str user_agent: Optional[str] created_at: float last_activity: float class SessionManager: """Manages per-device sessions for PM3 access. Each PM3 device can have at most one active session. Multiple devices can be used concurrently by different sessions. """ # Key used for legacy single-device mode DEFAULT_DEVICE_KEY = "_default" def __init__(self): """Initialize session manager.""" self._active_sessions: Dict[str, Session] = {} # keyed by device_id self._lock = asyncio.Lock() def _get_device_key(self, device_id: Optional[str]) -> str: """Get the key to use for session lookup. Args: device_id: Device ID or None for legacy mode Returns: Key to use in _active_sessions dict """ return device_id or self.DEFAULT_DEVICE_KEY def _cleanup_expired_sessions(self) -> None: """Remove any expired sessions from the active sessions dict.""" current_time = time.time() expired_keys = [ key for key, session in self._active_sessions.items() if current_time - session.last_activity > config.SESSION_TIMEOUT ] for key in expired_keys: del self._active_sessions[key] def has_active_session(self, device_id: Optional[str] = None) -> bool: """Check if there's an active session for a device. Args: device_id: Device ID to check. If None, checks for any active session. Returns: True if active session exists """ self._cleanup_expired_sessions() if device_id is None: # Check if ANY session is active (legacy behavior) return len(self._active_sessions) > 0 key = self._get_device_key(device_id) return key in self._active_sessions async def create_session( self, client_ip: str, user_agent: Optional[str] = None, force_takeover: bool = False, device_id: Optional[str] = None ) -> tuple[bool, Optional[str], Optional[str]]: """Create a new session for a device. Args: client_ip: Client IP address user_agent: Client user agent string force_takeover: Force takeover of existing session device_id: Device ID to create session for (None for legacy mode) Returns: Tuple of (success, session_id, error_message) """ async with self._lock: self._cleanup_expired_sessions() key = self._get_device_key(device_id) # Check if another session is active for this device if key in self._active_sessions and not force_takeover: return False, None, f"Another session is active for device {device_id or 'default'}" # Create new session session_id = str(uuid.uuid4()) current_time = time.time() self._active_sessions[key] = Session( session_id=session_id, device_id=device_id, client_ip=client_ip, user_agent=user_agent, created_at=current_time, last_activity=current_time ) return True, session_id, None async def release_session(self, session_id: str, device_id: Optional[str] = None) -> bool: """Release a session. Args: session_id: Session ID to release device_id: Device ID (if known). If None, searches all sessions. Returns: True if session was released, False if not found """ async with self._lock: if device_id is not None: # Direct lookup by device_id key = self._get_device_key(device_id) if key in self._active_sessions and self._active_sessions[key].session_id == session_id: del self._active_sessions[key] return True else: # Search all sessions for the session_id for key, session in list(self._active_sessions.items()): if session.session_id == session_id: del self._active_sessions[key] return True return False def update_activity(self, session_id: str, device_id: Optional[str] = None) -> bool: """Update session activity timestamp. Args: session_id: Session ID to update device_id: Device ID (if known). If None, searches all sessions. Returns: True if updated, False if session not found """ if device_id is not None: key = self._get_device_key(device_id) if key in self._active_sessions and self._active_sessions[key].session_id == session_id: self._active_sessions[key].last_activity = time.time() return True else: # Search all sessions for session in self._active_sessions.values(): if session.session_id == session_id: session.last_activity = time.time() return True return False def can_execute(self, session_id: Optional[str], device_id: Optional[str] = None) -> bool: """Check if a session can execute commands on a device. Args: session_id: Session ID to check (None for no session) device_id: Device ID to check (None for legacy single-device mode) Returns: True if session can execute, False otherwise """ self._cleanup_expired_sessions() key = self._get_device_key(device_id) # No active session for this device - allow execution if key not in self._active_sessions: return True # Check if the provided session ID matches the device's active session if session_id and self._active_sessions[key].session_id == session_id: return True return False def get_active_session(self, device_id: Optional[str] = None) -> Optional[Session]: """Get the currently active session for a device. Args: device_id: Device ID to get session for. If None, returns any active session. Returns: Active session or None """ self._cleanup_expired_sessions() if device_id is None: # Return first active session (legacy behavior) return next(iter(self._active_sessions.values()), None) key = self._get_device_key(device_id) return self._active_sessions.get(key) def get_session_for_device(self, device_id: str) -> Optional[Session]: """Get the active session for a specific device. Args: device_id: Device ID to get session for Returns: Active session for the device or None """ return self.get_active_session(device_id) def get_all_active_sessions(self) -> Dict[str, Session]: """Get all active sessions. Returns: Dict of device_id -> Session for all active sessions """ self._cleanup_expired_sessions() return dict(self._active_sessions)