#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Smart Session Management Provides session pooling, health monitoring, and automatic refresh """ import os import time import logging import datetime from typing import Optional, Dict, Any, Callable class SessionInfo: """Information about a browser session.""" def __init__(self, driver: Any, session_id: str, created_at: datetime.datetime, last_used: datetime.datetime, user_data_dir: Optional[str] = None): """Initialize session info. Args: driver: Browser/session instance session_id: Unique session identifier created_at: When session was created last_used: When session was last used user_data_dir: User data directory for session persistence """ self.driver = driver self.session_id = session_id self.created_at = created_at self.last_used = last_used self.user_data_dir = user_data_dir self.is_valid = True self.use_count = 0 def mark_used(self): """Mark session as used.""" self.last_used = datetime.datetime.now() self.use_count += 1 def is_expired(self, max_age: int) -> bool: """Check if session is expired. Args: max_age: Maximum age in seconds Returns: True if expired, False otherwise """ age = (datetime.datetime.now() - self.created_at).total_seconds() return age > max_age def needs_refresh(self, refresh_threshold: int) -> bool: """Check if session needs refresh. Args: refresh_threshold: Threshold in seconds before expiry to refresh Returns: True if needs refresh, False otherwise """ age = (datetime.datetime.now() - self.created_at).total_seconds() max_age = refresh_threshold * 2 # Assume max_age is 2x refresh_threshold return age > (max_age - refresh_threshold) class SessionManager: """Manages browser sessions with pooling and health monitoring.""" def __init__(self, config: Optional[Dict[str, Any]] = None): """Initialize session manager. Args: config: Configuration dictionary """ self.config = config or {} self.check_interval = self.config.get('check_interval', 1) self.login_timeout = self.config.get('login_timeout', 600) self.session_refresh_threshold = self.config.get('session_refresh_threshold', 300) self.max_session_age = self.config.get('max_session_age', 3600) self.sessions: Dict[str, SessionInfo] = {} self.logger = logging.getLogger(__name__) def create_session( self, session_id: str, driver: Any, user_data_dir: Optional[str] = None ) -> SessionInfo: """Create a new session. Args: session_id: Unique session identifier driver: Browser/session instance user_data_dir: User data directory for session persistence Returns: SessionInfo instance """ now = datetime.datetime.now() session_info = SessionInfo( driver=driver, session_id=session_id, created_at=now, last_used=now, user_data_dir=user_data_dir ) self.sessions[session_id] = session_info self.logger.info(f"Created new session: {session_id}") return session_info def get_session(self, session_id: str) -> Optional[SessionInfo]: """Get an existing session. Args: session_id: Session identifier Returns: SessionInfo if found, None otherwise """ session = self.sessions.get(session_id) if session: session.mark_used() return session def validate_session(self, session: SessionInfo, validation_func: Callable[[Any], bool]) -> bool: """Validate a session using a custom validation function. Args: session: SessionInfo to validate validation_func: Function that takes a browser instance and returns True if valid Returns: True if session is valid, False otherwise """ try: is_valid = validation_func(session.driver) session.is_valid = is_valid if not is_valid: self.logger.warning(f"Session {session.session_id} validation failed") return is_valid except Exception as e: self.logger.error(f"Error validating session {session.session_id}: {e}") session.is_valid = False return False def refresh_session(self, session: SessionInfo, refresh_func: Callable[[Any], bool]) -> bool: """Refresh a session using a custom refresh function. Args: session: SessionInfo to refresh refresh_func: Function that takes a browser instance and returns True if refresh successful Returns: True if refresh successful, False otherwise """ try: success = refresh_func(session.driver) if success: session.created_at = datetime.datetime.now() session.last_used = datetime.datetime.now() session.is_valid = True self.logger.info(f"Session {session.session_id} refreshed successfully") return success except Exception as e: self.logger.error(f"Error refreshing session {session.session_id}: {e}") session.is_valid = False return False def cleanup_expired_sessions(self): """Remove expired sessions.""" expired = [] for session_id, session in self.sessions.items(): if session.is_expired(self.max_session_age): expired.append(session_id) for session_id in expired: self.remove_session(session_id) self.logger.info(f"Removed expired session: {session_id}") def remove_session(self, session_id: str): """Remove a session. Args: session_id: Session identifier """ if session_id in self.sessions: session = self.sessions[session_id] try: session.driver.quit() except Exception as e: self.logger.warning(f"Error closing driver for session {session_id}: {e}") del self.sessions[session_id] self.logger.info(f"Removed session: {session_id}") def get_or_create_session( self, session_id: str, create_func: Callable[[], Any], validation_func: Optional[Callable[[Any], bool]] = None, refresh_func: Optional[Callable[[Any], bool]] = None, user_data_dir: Optional[str] = None ) -> SessionInfo: """Get existing session or create a new one. Args: session_id: Session identifier create_func: Function that creates a new browser instance validation_func: Optional function to validate session refresh_func: Optional function to refresh session user_data_dir: User data directory for session persistence Returns: SessionInfo instance """ # Clean up expired sessions first self.cleanup_expired_sessions() # Try to get existing session session = self.get_session(session_id) if session: # Validate existing session if validation_func: if not self.validate_session(session, validation_func): self.logger.warning(f"Session {session_id} is invalid, creating new one") self.remove_session(session_id) session = None elif session.needs_refresh(self.session_refresh_threshold) and refresh_func: # Try to refresh session if not self.refresh_session(session, refresh_func): self.logger.warning(f"Failed to refresh session {session_id}, creating new one") self.remove_session(session_id) session = None # Create new session if needed if not session: driver = create_func() session = self.create_session(session_id, driver, user_data_dir) return session def wait_for_manual_login( self, job_id: str, output_dir: str, driver: Any, login_url: str, validation_func: Optional[Callable[[Any], bool]] = None ) -> Optional[str]: """Wait for user to manually log in through browser. Args: job_id: Job ID for signal file naming output_dir: Output directory where signal files will be created driver: Browser instance (should be in visible mode) login_url: URL to navigate to for login validation_func: Optional function to validate login success Returns: Session ID extracted from URL after login, or None if failed """ signal_file = os.path.join(output_dir, f"waiting_for_login_{job_id}.txt") complete_file = os.path.join(output_dir, f"login_complete_{job_id}.txt") # Create signal file to notify Electron self.logger.info("Creating login signal file for Electron...") try: with open(signal_file, 'w', encoding='utf-8') as f: f.write(f"job_id={job_id}\n") f.write(f"login_url={login_url}\n") f.write(f"timestamp={datetime.datetime.now().isoformat()}\n") self.logger.info(f"✓ Signal file created: {signal_file}") except Exception as e: self.logger.error(f"Failed to create signal file: {e}") raise # Navigate to login URL self.logger.info("Navigating to login page...") self.logger.info("⚠️ Please log in through the browser window") driver.get(login_url) time.sleep(3) # Wait for login completion signal self.logger.info("Waiting for login completion signal from Electron...") start_time = time.time() while time.time() - start_time < self.login_timeout: # Check for stop signal stop_file = os.path.join(output_dir, f"stop_{job_id}.txt") if os.path.exists(stop_file): self.logger.info("Stop signal received during login wait") raise Exception("Automation stopped by user") # Check for completion signal if os.path.exists(complete_file): self.logger.info("✓ Login completion signal received") time.sleep(2) break time.sleep(self.check_interval) # Check if timeout occurred if not os.path.exists(complete_file): try: if os.path.exists(signal_file): os.remove(signal_file) except: pass raise Exception(f"Login timeout - user did not complete login within {self.login_timeout} seconds") # Verify login was successful time.sleep(2) if validation_func: if not validation_func(driver): try: if os.path.exists(signal_file): os.remove(signal_file) if os.path.exists(complete_file): os.remove(complete_file) except: pass raise Exception("Login verification failed - session may not be valid") # Extract session ID from URL (if applicable) current_url = driver.current_url self.logger.info(f"Current URL after login: {current_url}") # Clean up signal files try: if os.path.exists(signal_file): os.remove(signal_file) if os.path.exists(complete_file): os.remove(complete_file) except: pass # Try to extract session ID from URL (implementation depends on specific site) # This is a placeholder - actual implementation should match the site's URL pattern session_match = None if 'session' in current_url.lower() or 'sid' in current_url.lower(): import re # Generic pattern matching - should be customized per site session_match = re.search(r'[?&](?:session|sid|sessionid)=([^&]+)', current_url, re.IGNORECASE) if session_match: session_id = session_match.group(1) self.logger.info(f"✓ Extracted session ID: {session_id}") return session_id # Return a default session ID if extraction fails return "default_session" def cleanup_all(self): """Clean up all sessions.""" session_ids = list(self.sessions.keys()) for session_id in session_ids: self.remove_session(session_id) self.logger.info("All sessions cleaned up")