headplane/remote-access/session_manager.py

"""
Session management for Heady remote access
Tracks active connections, recordings, and audit logs
"""
import asyncio
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from pydantic import BaseModel
import logging

from auth import User

logger = logging.getLogger(__name__)


class RemoteSession(BaseModel):
    """Remote access session model"""
    id: str
    user_id: str
    user_email: str
    node_name: str
    protocol: str
    started_at: datetime
    ended_at: Optional[datetime] = None
    duration_seconds: int = 0
    recording_enabled: bool = False
    recording_path: Optional[str] = None
    client_ip: Optional[str] = None
    status: str = "active"  # active, ended, error


class SessionManager:
    """
    Manages remote access sessions
    Adapted from your BookingResourceConsole model
    """
    
    def __init__(self):
        self.active_sessions: Dict[str, RemoteSession] = {}
        self.session_history: List[RemoteSession] = []
        
    async def create_session(
        self,
        user: User,
        node_name: str,
        protocol: str,
        client_ip: Optional[str] = None
    ) -> RemoteSession:
        """Create new remote access session"""
        
        session_id = str(uuid.uuid4())
        
        # Check if recording is required for this user role
        recording_enabled = await self._should_record_session(user)
        recording_path = None
        
        if recording_enabled:
            recording_path = f"/recordings/{session_id}-{datetime.now().isoformat()}.guac"
            
        session = RemoteSession(
            id=session_id,
            user_id=user.id,
            user_email=user.email,
            node_name=node_name,
            protocol=protocol,
            started_at=datetime.now(),
            recording_enabled=recording_enabled,
            recording_path=recording_path,
            client_ip=client_ip
        )
        
        self.active_sessions[session_id] = session
        
        logger.info(
            f"Created session {session_id} for {user.email} -> {node_name} "
            f"(protocol: {protocol}, recording: {recording_enabled})"
        )
        
        return session
    
    async def close_session(self, session_id: str) -> bool:
        """Close and archive session"""
        
        if session_id not in self.active_sessions:
            return False
            
        session = self.active_sessions[session_id]
        session.ended_at = datetime.now()
        session.duration_seconds = int(
            (session.ended_at - session.started_at).total_seconds()
        )
        session.status = "ended"
        
        # Move to history
        self.session_history.append(session)
        del self.active_sessions[session_id]
        
        logger.info(
            f"Closed session {session_id} for {session.user_email} "
            f"(duration: {session.duration_seconds}s)"
        )
        
        return True
    
    async def terminate_session(self, session_id: str, requesting_user_id: str) -> bool:
        """Terminate session (admin action or self-termination)"""
        
        if session_id not in self.active_sessions:
            return False
            
        session = self.active_sessions[session_id]
        
        # Check if user can terminate this session
        if session.user_id != requesting_user_id:
            # TODO: Check if requesting user has admin privileges
            pass
            
        await self.close_session(session_id)
        return True
    
    async def get_user_sessions(self, user_id: str) -> List[RemoteSession]:
        """Get all sessions for a user"""
        
        user_sessions = []
        
        # Active sessions
        for session in self.active_sessions.values():
            if session.user_id == user_id:
                user_sessions.append(session)
                
        # Recent history (last 24 hours)
        cutoff = datetime.now() - timedelta(hours=24)
        for session in self.session_history:
            if session.user_id == user_id and session.started_at > cutoff:
                user_sessions.append(session)
                
        return sorted(user_sessions, key=lambda s: s.started_at, reverse=True)
    
    async def get_recording_config(self, session_id: str) -> Optional[Dict[str, str]]:
        """
        Get guacamole recording configuration
        Returns dict for guacd handshake parameters
        """
        
        if session_id not in self.active_sessions:
            return None
            
        session = self.active_sessions[session_id]
        
        if not session.recording_enabled or not session.recording_path:
            return None
            
        return {
            "recording-path": session.recording_path,
            "recording-name": f"session-{session_id}",
            "create-recording-path": "true"
        }
    
    async def _should_record_session(self, user: User) -> bool:
        """
        Determine if session should be recorded based on user role
        Matches your recording logic from Django
        """
        
        # Owners typically don't need recording
        if user.role == "owner":
            return False
            
        # All other privileged roles should be recorded
        if user.role in ["admin", "network_admin", "it_admin"]:
            return True
            
        # Members and auditors don't get terminal access anyway
        return False
    
    async def get_session_stats(self) -> Dict:
        """Get session statistics for monitoring"""
        
        active_count = len(self.active_sessions)
        total_sessions_today = len([
            s for s in self.session_history 
            if s.started_at.date() == datetime.now().date()
        ])
        
        return {
            "active_sessions": active_count,
            "sessions_today": total_sessions_today,
            "total_recorded": len([
                s for s in self.session_history 
                if s.recording_enabled
            ])
        }
    
    async def cleanup_old_sessions(self, max_age_days: int = 30):
        """Clean up old session history"""
        
        cutoff = datetime.now() - timedelta(days=max_age_days)
        
        original_count = len(self.session_history)
        self.session_history = [
            s for s in self.session_history 
            if s.started_at > cutoff
        ]
        
        cleaned_count = original_count - len(self.session_history)
        
        if cleaned_count > 0:
            logger.info(f"Cleaned up {cleaned_count} old sessions")
            
        return cleaned_count