""" Session Manager - Multi-session rotation and load balancing Handles multiple Pyrogram sessions to maximize bandwidth and avoid flood limits """ import asyncio import logging import os from typing import List, Optional, AsyncGenerator from io import BytesIO from pyrogram import Client from pyrogram.errors import FloodWait, BadRequest from pyrogram.types import Message from utils import CHUNK_SIZE logger = logging.getLogger(__name__) class TelegramSession: """Wrapper for a single Pyrogram session""" def __init__( self, session_name: str, api_id: int, api_hash: str, session_string: Optional[str] = None ): self.session_name = session_name self.api_id = api_id self.api_hash = api_hash self.session_string = session_string self.client: Optional[Client] = None self.is_active = False self.upload_count = 0 self.download_count = 0 async def initialize(self): """Initialize and start the Pyrogram client""" try: if self.session_string: self.client = Client( name=self.session_name, api_id=self.api_id, api_hash=self.api_hash, session_string=self.session_string, in_memory=True, no_updates=True ) else: self.client = Client( name=self.session_name, api_id=self.api_id, api_hash=self.api_hash, in_memory=True, no_updates=True ) await self.client.start() self.is_active = True me = await self.client.get_me() logger.info( f"Session {self.session_name} initialized: " f"@{me.username or me.first_name}" ) except Exception as e: logger.error(f"Failed to initialize session {self.session_name}: {str(e)}") self.is_active = False raise async def cleanup(self): """Stop and cleanup the session""" if self.client and self.is_active: try: await self.client.stop() self.is_active = False logger.info(f"Session {self.session_name} stopped") except Exception as e: logger.error(f"Error stopping session {self.session_name}: {str(e)}") class SessionManager: """Manages multiple Telegram sessions for load balancing""" def __init__(self): self.sessions: List[TelegramSession] = [] self.bot_token: Optional[str] = None self.bot_session: Optional[TelegramSession] = None self.current_upload_index = 0 self.current_download_index = 0 self.lock = asyncio.Lock() async def initialize(self): """Initialize all sessions from environment variables""" logger.info("Initializing Session Manager...") # Get bot token self.bot_token = os.getenv("BOT_TOKEN") if not self.bot_token: raise ValueError("BOT_TOKEN environment variable is required") # Initialize bot session api_id = int(os.getenv("API_ID", "0")) api_hash = os.getenv("API_HASH", "") if not api_id or not api_hash: raise ValueError("API_ID and API_HASH environment variables are required") self.bot_session = TelegramSession( session_name="bot_session", api_id=api_id, api_hash=api_hash ) # Override with bot token self.bot_session.client = Client( name="bot_session", api_id=api_id, api_hash=api_hash, bot_token=self.bot_token, in_memory=True, no_updates=True ) await self.bot_session.client.start() self.bot_session.is_active = True logger.info("Bot session initialized") # Initialize user sessions from SESSION_STRINGS session_strings = os.getenv("SESSION_STRINGS", "").split(",") session_strings = [s.strip() for s in session_strings if s.strip()] if not session_strings: logger.warning("No SESSION_STRINGS found, using bot session only") self.sessions = [self.bot_session] return # Create user sessions for i, session_string in enumerate(session_strings): session = TelegramSession( session_name=f"user_session_{i}", api_id=api_id, api_hash=api_hash, session_string=session_string ) try: await session.initialize() self.sessions.append(session) except Exception as e: logger.error(f"Failed to initialize user session {i}: {str(e)}") if not self.sessions: # Fallback to bot session self.sessions = [self.bot_session] logger.info(f"Session Manager initialized with {len(self.sessions)} session(s)") def get_next_upload_session(self) -> TelegramSession: """Get next session for upload (round-robin)""" if not self.sessions: raise RuntimeError("No active sessions available") session = self.sessions[self.current_upload_index] self.current_upload_index = (self.current_upload_index + 1) % len(self.sessions) return session def get_next_download_session(self) -> TelegramSession: """Get next session for download (round-robin)""" if not self.sessions: raise RuntimeError("No active sessions available") session = self.sessions[self.current_download_index] self.current_download_index = ( self.current_download_index + 1 ) % len(self.sessions) return session async def upload_part( self, data: bytes, filename: str, max_retries: int = 3 ) -> str: """ Upload a file part to Telegram Returns: file_id for later retrieval """ retry_count = 0 while retry_count < max_retries: session = self.get_next_upload_session() if not session.is_active or not session.client: retry_count += 1 continue try: # Upload to "Saved Messages" (self chat) message: Message = await session.client.send_document( chat_id="me", document=BytesIO(data), file_name=filename, force_document=True ) session.upload_count += 1 file_id = message.document.file_id logger.debug( f"Part uploaded via {session.session_name}: " f"file_id={file_id}, size={len(data)}" ) return file_id except FloodWait as e: logger.warning( f"FloodWait on {session.session_name}: waiting {e.value}s" ) await asyncio.sleep(e.value) retry_count += 1 except Exception as e: logger.error( f"Upload failed on {session.session_name}: {str(e)}" ) retry_count += 1 if retry_count < max_retries: await asyncio.sleep(2 ** retry_count) raise RuntimeError(f"Failed to upload part after {max_retries} retries") async def stream_part( self, file_id: str, offset: int = 0, limit: Optional[int] = None ) -> AsyncGenerator[bytes, None]: """ Stream a file part from Telegram Yields chunks of data """ session = self.get_next_download_session() if not session.is_active or not session.client: raise RuntimeError("No active session available for streaming") try: bytes_read = 0 async for chunk in session.client.stream_media( file_id, offset=offset, limit=limit or 0 ): if limit and bytes_read + len(chunk) > limit: # Trim final chunk yield chunk[:limit - bytes_read] break yield chunk bytes_read += len(chunk) if limit and bytes_read >= limit: break session.download_count += 1 logger.debug( f"Part streamed via {session.session_name}: " f"file_id={file_id}, bytes={bytes_read}" ) except FloodWait as e: logger.warning(f"FloodWait on download: waiting {e.value}s") await asyncio.sleep(e.value) # Retry with next session async for chunk in self.stream_part(file_id, offset, limit): yield chunk except Exception as e: logger.error(f"Stream failed: {str(e)}") raise async def delete_part(self, file_id: str) -> bool: """Delete a file part from Telegram""" session = self.get_next_upload_session() if not session.is_active or not session.client: return False try: # Get message and delete it # Note: This requires the message_id, which we don't store # For production, consider storing message_ids in metadata logger.warning("Delete operation requires message_id, not implemented") return False except Exception as e: logger.error(f"Delete failed: {str(e)}") return False async def cleanup(self): """Cleanup all sessions""" logger.info("Cleaning up Session Manager...") for session in self.sessions: await session.cleanup() if self.bot_session and self.bot_session != self.sessions[0]: await self.bot_session.cleanup() logger.info("Session Manager cleanup complete") def get_stats(self) -> dict: """Get session statistics""" return { "total_sessions": len(self.sessions), "active_sessions": sum(1 for s in self.sessions if s.is_active), "sessions": [ { "name": s.session_name, "active": s.is_active, "uploads": s.upload_count, "downloads": s.download_count } for s in self.sessions ] }