fix: sqlite

Author: An0nX

scripts/manage_account.py

@@ -34,7 +34,7 @@ async def add_account_logic(args: argparse.Namespace) -> None:
     """
     logger.info(f"--- Adding new account: {args.name} ---")
 
-    session_file: str = f"temp_cli_{args.name}.session"
+    session_file_path: str = f"temp_cli_{args.name}.session"
     temp_client: Optional[TelegramClient] = None
     try:
         async with get_db() as db:
@@ -50,7 +50,7 @@ async def add_account_logic(args: argparse.Namespace) -> None:
                 return
 
             logger.info("Initializing temporary session to verify credentials...")
-            temp_client = TelegramClient(SQLiteSession(session_file), int(api_id), api_hash)
+            temp_client = TelegramClient(SQLiteSession(session_file_path), int(api_id), api_hash)
             await temp_client.connect()
 
             if not await temp_client.is_user_authorized():
@@ -137,24 +137,14 @@ async def add_account_logic(args: argparse.Namespace) -> None:
                 return
 
             # Now extract session from file and save to DB
-            logger.info("Extracting session data and saving to the database...")
-            reader_session = SQLiteSession(session_file)
-            reader_session.load()
-            
-            update_state = reader_session.get_update_state(0)
-            pts, qts, date_ts, seq, _ = (None, None, None, None, None)
-            if update_state:
-                pts, qts, date_ts, seq, _ = update_state
+            logger.info("Reading session file and saving to the database...")
+            with open(session_file_path, 'rb') as f:
+                session_bytes: bytes = f.read()
 
             await db_manager.add_or_update_session(
                 db,
                 account_id=new_account.account_id,
-                dc_id=reader_session.dc_id,
-                server_address=reader_session.server_address,
-                port=reader_session.port,
-                auth_key_data=reader_session.auth_key.key,
-                takeout_id=reader_session.takeout_id,
-                pts=pts, qts=qts, date=date_ts, seq=seq
+                session_file=session_bytes
             )
 
             logger.info(f"\nSuccess! Account '{args.name}' was added. Restart the bot to activate.")
@@ -164,9 +154,9 @@ async def add_account_logic(args: argparse.Namespace) -> None:
     finally:
         if temp_client and temp_client.is_connected():
             await temp_client.disconnect()
-        if os.path.exists(session_file):
-            os.remove(session_file)
-            logger.info(f"Cleaned up temporary session file: {session_file}")
+        if os.path.exists(session_file_path):
+            os.remove(session_file_path)
+            logger.info(f"Cleaned up temporary session file: {session_file_path}")
 
 
 async def edit_account_logic(args: argparse.Namespace) -> None:

userbot/__init__.py

@@ -5,13 +5,14 @@
 
 from faker import Faker
 from telethon import TelegramClient as TelethonTelegramClient, events
+from telethon.tl.functions.channels import JoinChannelRequest
 from telethon.errors.rpcerrorlist import UserAlreadyParticipantError
 from python_socks import ProxyType
 
 from userbot.src.config import API_ID, API_HASH, LOG_LEVEL
 from userbot.src.db.session import initialize_database, get_db
-from userbot.src.db.models import Account
-from userbot.src.db_session import DbSession
+from userbot.src.db.models import Account, Session
+from userbot.src.memory_session import MemorySession
 from userbot.src.encrypt import encryption_manager
 import userbot.src.db_manager as db_manager
 from userbot.src.log_handler import DBLogHandler
@@ -35,12 +36,11 @@ class TelegramClient(TelethonTelegramClient):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.lang_code: str = 'ru'
+        self.account_id_override: Optional[int] = kwargs.get('account_id_override')
 
     @property
     def current_account_id(self) -> Optional[int]:
-        if hasattr(self, 'session') and isinstance(self.session, DbSession):
-            return self.session.account_id
-        return None
+        return self.account_id_override
 
     async def get_string(self, key: str, module_name: Optional[str] = None, **kwargs) -> str:
         return translator.get_string(self.lang_code, key, module_name, **kwargs)
@@ -139,21 +139,29 @@ async def manage_clients() -> None:
                      encryption_manager.decrypt(account.proxy_username).decode() if account.proxy_username else None,
                      encryption_manager.decrypt(account.proxy_password).decode() if account.proxy_password else None
                  )
+        
+        session_instance: MemorySession
+        if account.session and account.session.session_file:
+            try:
+                decrypted_session_bytes: bytes = encryption_manager.decrypt(account.session.session_file)
+                session_instance = MemorySession(decrypted_session_bytes)
+            except Exception as e:
+                logger.error(f"Could not decrypt session for '{account.account_name}'. Skipping. Error: {e}")
+                continue
+        else:
+            logger.warning(f"No session file found in DB for account '{account.account_name}'. A new login will be required if client is used directly.")
+            session_instance = MemorySession(None)
 
-        session = DbSession(
-            account_id=account.account_id,
-            user_id=account.user_telegram_id,
-            access_hash=account.access_hash
-        )
 
         new_client: TelegramClient = TelegramClient(
-            session=session, 
+            session=session_instance, 
             api_id=int(acc_api_id), 
             api_hash=acc_api_hash, 
             device_model=account.device_model, 
             system_version=account.system_version, 
             app_version=account.app_version,
-            proxy=proxy_details
+            proxy=proxy_details,
+            account_id_override=account.account_id
         )
         ACTIVE_CLIENTS[account.account_id] = new_client
         tasks.append(start_individual_client(new_client, account))

userbot/src/core_handlers.py

@@ -21,7 +21,8 @@
 
 # --- Helper ---
 async def get_account_id_from_client(client) -> int | None:
-    return next((acc_id for acc_id, c in ACTIVE_CLIENTS.items() if c == client), None)
+    # Access the override property directly
+    return client.current_account_id if hasattr(client, 'current_account_id') else None
 
 # --- Module Management ---
 async def load_account_modules(account_id: int, client_instance: TelegramClient, current_help_info: Dict[str, str]):
@@ -75,7 +76,7 @@ async def list_accounts_handler(event: events.NewMessage.Event):
 
 async def add_account_handler(event: events.NewMessage.Event):
     account_name: str = event.pattern_match.group(1)
-    session_file: str = f"temp_add_{account_name}.session"
+    session_file_path: str = f"temp_add_{account_name}.session"
     temp_client: Optional[TelegramClient] = None
 
     try:
@@ -86,7 +87,7 @@ async def add_account_handler(event: events.NewMessage.Event):
             api_hash_resp = await conv.get_response(); api_hash = api_hash_resp.text.strip()
 
             await conv.send_message(await event.client.get_string("verifying_creds"))
-            temp_client = TelegramClient(SQLiteSession(session_file), int(api_id), api_hash)
+            temp_client = TelegramClient(SQLiteSession(session_file_path), int(api_id), api_hash)
 
             await temp_client.connect()
             if not await temp_client.is_user_authorized():
@@ -141,23 +142,13 @@ async def add_account_handler(event: events.NewMessage.Event):
                     return
 
                 # Now extract session data and save it
-                reader_session = SQLiteSession(session_file)
-                reader_session.load()
-                
-                update_state = reader_session.get_update_state(0)
-                pts, qts, date_ts, seq, _ = (None, None, None, None, None)
-                if update_state:
-                    pts, qts, date_ts, seq, _ = update_state
-                
+                with open(session_file_path, 'rb') as f:
+                    session_bytes: bytes = f.read()
+
                 await db_manager.add_or_update_session(
                     db,
                     account_id=new_acc.account_id,
-                    dc_id=reader_session.dc_id,
-                    server_address=reader_session.server_address,
-                    port=reader_session.port,
-                    auth_key_data=reader_session.auth_key.key,
-                    takeout_id=reader_session.takeout_id,
-                    pts=pts, qts=qts, date=date_ts, seq=seq
+                    session_file=session_bytes
                 )
 
             await conv.send_message(await event.client.get_string("add_acc_success", account_name=account_name))
@@ -170,8 +161,8 @@ async def add_account_handler(event: events.NewMessage.Event):
     finally:
         if temp_client and temp_client.is_connected():
             await temp_client.disconnect()
-        if os.path.exists(session_file):
-            os.remove(session_file)
+        if os.path.exists(session_file_path):
+            os.remove(session_file_path)
 
 
 async def delete_account_handler(event: events.NewMessage.Event):

userbot/src/db/models.py

@@ -2,7 +2,7 @@
     Column, Integer, String, Boolean, ForeignKey, BIGINT, TEXT,
     TIMESTAMP, UniqueConstraint
 )
-from sqlalchemy.dialects.postgresql import BYTEA, JSONB
+from sqlalchemy.dialects.postgresql import BYTEA
 from sqlalchemy.orm import declarative_base, relationship
 from sqlalchemy.sql import func
 
@@ -43,18 +43,7 @@ class Session(Base):
     __tablename__ = 'sessions'
     session_id = Column(Integer, primary_key=True)
     account_id = Column(Integer, ForeignKey('accounts.account_id', ondelete='CASCADE'), nullable=False, unique=True)
-    dc_id = Column(Integer, nullable=False)
-    server_address = Column(TEXT)
-    port = Column(Integer)
-    auth_key_data = Column(BYTEA)
-    
-    # Update state fields
-    pts = Column(Integer)
-    qts = Column(Integer)
-    date = Column(BIGINT)
-    seq = Column(Integer)
-    takeout_id = Column(BIGINT)
-    
+    session_file = Column(BYTEA, nullable=False)
     last_used_at = Column(TIMESTAMP(timezone=True), server_default=func.now(), nullable=False)
     created_at = Column(TIMESTAMP(timezone=True), server_default=func.now(), nullable=False)
 
@@ -80,7 +69,7 @@ class AccountModule(Base):
     module_id = Column(Integer, ForeignKey('modules.module_id', ondelete='CASCADE'), nullable=False)
     is_active = Column(Boolean, default=True)
     is_trusted = Column(Boolean, default=False)
-    configuration = Column(JSONB)
+    configuration = Column(BYTEA) # Changed to BYTEA for encrypted JSON
     activated_at = Column(TIMESTAMP(timezone=True), server_default=func.now(), nullable=False)
     updated_at = Column(TIMESTAMP(timezone=True), onupdate=func.now())
 

userbot/src/db_manager.py

@@ -116,7 +116,7 @@ async def get_all_accounts(db: AsyncSession) -> List[Account]:
 
 async def get_all_active_accounts(db: AsyncSession) -> List[Account]:
     """Retrieves all enabled accounts from the database."""
-    result = await db.execute(select(Account).where(Account.is_enabled == True))
+    result = await db.execute(select(Account).where(Account.is_enabled == True).options(selectinload(Account.session)))
     return result.scalars().all()
 
 async def delete_account(db: AsyncSession, account_name: str) -> bool:
@@ -161,34 +161,35 @@ async def update_account_self_info(account_id: int, user_id: int, access_hash: i
 
 # --- Session CRUD ---
 async def get_session(db: AsyncSession, account_id: int) -> Optional[Session]:
-    """Retrieves a session for a given account and decrypts its auth key."""
+    """Retrieves a session for a given account ID."""
     result = await db.execute(select(Session).where(Session.account_id == account_id))
-    session = result.scalars().first()
-    if session and session.auth_key_data:
-        try:
-            session.auth_key_data = encryption_manager.decrypt(session.auth_key_data)
-        except Exception as e:
-            logger.error(f"Failed to decrypt session auth_key for account {account_id}: {e}")
-            return None
-    return session
+    return result.scalars().first()
 
-async def add_or_update_session(db: AsyncSession, **kwargs) -> Optional[Session]:
-    """Adds or updates a session in the database, encrypting the auth key."""
-    account_id = kwargs.get("account_id")
+async def add_or_update_session(db: AsyncSession, account_id: int, session_file: bytes) -> Optional[Session]:
+    """
+    Adds or updates a session in the database.
+
+    Args:
+        db (AsyncSession): The database session.
+        account_id (int): The ID of the account this session belongs to.
+        session_file (bytes): The raw, unencrypted bytes of the session file.
+
+    Returns:
+        Optional[Session]: The created or updated Session object.
+    """
     if not account_id: return None
 
     result = await db.execute(select(Session).where(Session.account_id == account_id))
     session = result.scalars().first()
 
+    encrypted_session_file = encryption_manager.encrypt(session_file)
+    
     if not session:
-        session = Session(account_id=account_id)
+        session = Session(account_id=account_id, session_file=encrypted_session_file)
         db.add(session)
-        
-    for key, value in kwargs.items():
-        if key == "auth_key_data" and value is not None:
-            value = encryption_manager.encrypt(value)
-        setattr(session, key, value)
-    
+    else:
+        session.session_file = encrypted_session_file
+
     session.last_used_at = datetime.now(timezone.utc)
     await db.flush()
     return session