pi-pm3/app/backend/managers/session_manager.py

"""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)