feat: access_hash gathering

Author: An0nX

userbot/src/db/models.py

@@ -11,19 +11,17 @@
 class Account(Base):
     __tablename__ = 'accounts'
     account_id = Column(Integer, primary_key=True)
-    user_telegram_id = Column(BIGINT, unique=True, nullable=True) # Can be null until first login
+    user_telegram_id = Column(BIGINT, unique=True, nullable=True)
     api_id = Column(BYTEA, nullable=False)
     api_hash = Column(BYTEA, nullable=False)
     account_name = Column(String(255), unique=True, nullable=False)
     is_enabled = Column(Boolean, default=True, nullable=False)
     lang_code = Column(String(10), default='ru', nullable=False)
 
-    # Device info
     device_model = Column(TEXT)
     system_version = Column(TEXT)
     app_version = Column(TEXT)
 
-    # Proxy info
     proxy_type = Column(TEXT)
     proxy_ip = Column(TEXT)
     proxy_port = Column(Integer)
@@ -46,11 +44,18 @@ class Session(Base):
     server_address = Column(TEXT)
     port = Column(Integer)
     auth_key_data = Column(BYTEA)
+    
+    # New fields for self-entity caching
+    self_user_id = Column(BIGINT)
+    self_access_hash = Column(BIGINT)
+
+    # Update state fields
     pts = Column(Integer)
     qts = Column(Integer)
     date = Column(BIGINT)
     seq = Column(Integer)
     takeout_id = Column(BIGINT)
+    
     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)
 

userbot/src/db_manager.py

@@ -28,37 +28,26 @@ async def add_account(db: AsyncSession, account_name: str, api_id: str, api_hash
         )
         db.add(new_account)
         await db.flush()
-        logger.info(f"Added account '{account_name}' with ID: {new_account.account_id}")
         return new_account
     except IntegrityError:
-        logger.warning(f"Account with name '{account_name}' or user_id '{user_telegram_id}' already exists.")
-        await db.rollback()
-        return None
+        await db.rollback(); return None
     except Exception as e:
-        logger.error(f"Error adding account '{account_name}': {e}")
-        await db.rollback()
-        raise
+        await db.rollback(); raise e
 
 async def get_account(db: AsyncSession, account_name: str) -> Optional[Account]:
-    result = await db.execute(select(Account).where(Account.account_name == account_name))
-    return result.scalars().first()
+    result = await db.execute(select(Account).where(Account.account_name == account_name)); return result.scalars().first()
 
 async def get_account_by_id(db: AsyncSession, account_id: int) -> Optional[Account]:
-    """Retrieves an account by its primary key ID."""
-    result = await db.execute(select(Account).where(Account.account_id == account_id))
-    return result.scalars().first()
+    result = await db.execute(select(Account).where(Account.account_id == account_id)); return result.scalars().first()
 
 async def get_account_by_user_id(db: AsyncSession, user_id: int) -> Optional[Account]:
-    result = await db.execute(select(Account).where(Account.user_telegram_id == user_id))
-    return result.scalars().first()
+    result = await db.execute(select(Account).where(Account.user_telegram_id == user_id)); return result.scalars().first()
 
 async def get_all_accounts(db: AsyncSession) -> List[Account]:
-    result = await db.execute(select(Account).options(selectinload(Account.session)).order_by(Account.account_id))
-    return result.scalars().all()
+    result = await db.execute(select(Account).options(selectinload(Account.session)).order_by(Account.account_id)); return result.scalars().all()
 
 async def get_all_active_accounts(db: AsyncSession) -> List[Account]:
-    result = await db.execute(select(Account).where(Account.is_enabled == True))
-    return result.scalars().all()
+    result = await db.execute(select(Account).where(Account.is_enabled == True)); return result.scalars().all()
 
 async def delete_account(db: AsyncSession, account_name: str) -> bool:
     account = await get_account(db, account_name)

userbot/src/db_session.py

@@ -4,7 +4,7 @@
 
 from telethon.sessions.abstract import Session
 from telethon.crypto import AuthKey
-from telethon.tl.types import InputPeerSelf
+from telethon.tl.types import User, InputPeerUser
 
 import userbot.src.db_manager as db_manager
 from userbot.src.db.session import get_db
@@ -14,17 +14,12 @@
 class DbSession(Session):
     """
     A Telethon session class that stores session data in a PostgreSQL database.
-    This implementation correctly provides all the abstract methods and properties
-    required by Telethon's base Session class.
+    It implements entity processing to cache the current user's ID and access_hash,
+    which is crucial for the client's startup and authorization checks.
     """
 
     def __init__(self, account_id: int):
-        """
-        Initializes the database-backed session.
-
-        Args:
-            account_id (int): The unique identifier for the account this session belongs to.
-        """
+        """Initializes the database-backed session."""
         super().__init__()
         if not isinstance(account_id, int):
             raise ValueError("DbSession requires a valid integer account_id.")
@@ -36,16 +31,18 @@ def __init__(self, account_id: int):
         self._server_address: Optional[str] = None
         self._port: int = 443
         self._takeout_id: Optional[int] = None
+        
+        # In-memory cache for the self user entity
+        self._self_user_id: Optional[int] = None
+        self._self_access_hash: Optional[int] = None
 
         self._pts: Optional[int] = None
         self._qts: Optional[int] = None
         self._date: Optional[int] = None
         self._seq: Optional[int] = None
 
     async def load(self) -> None:
-        """
-        Loads the session data for the current account_id from the database.
-        """
+        """Loads session data, including the self-user cache, from the database."""
         logger.debug(f"Attempting to load session for account_id: {self.account_id}")
         async with get_db() as db:
             session_data = await db_manager.get_session(db, self.account_id)
@@ -61,15 +58,20 @@ async def load(self) -> None:
             self._qts = session_data.qts
             self._date = session_data.date
             self._seq = session_data.seq
+            
+            # Load self entity from cache
+            self._self_user_id = session_data.self_user_id
+            self._self_access_hash = session_data.self_access_hash
         else:
-            logger.info(f"No session data in DB for account_id: {self.account_id}. New login required.")
+            logger.info(f"No session data in DB for account_id: {self.account_id}.")
 
     async def save(self) -> None:
-        """Saves the current session data to the database."""
+        """Saves session data, including the self-user cache, to the database."""
         session_data = {
             "account_id": self.account_id, "dc_id": self._dc_id,
             "server_address": self._server_address, "port": self._port,
             "auth_key_data": self._auth_key.key if self._auth_key else None,
+            "self_user_id": self._self_user_id, "self_access_hash": self._self_access_hash,
             "pts": self._pts, "qts": self._qts, "date": self._date,
             "seq": self._seq, "takeout_id": self._takeout_id,
         }
@@ -78,7 +80,6 @@ async def save(self) -> None:
         logger.info(f"Session saved for account_id: {self.account_id}")
 
     async def delete(self) -> None:
-        """Deletes the session for the current account from the database."""
         async with get_db() as db:
             await db_manager.delete_session(db, self.account_id)
         self._auth_key = None; self._dc_id = 0
@@ -108,8 +109,7 @@ def get_update_state(self, entity_id: int) -> Optional[Tuple[int, int, int, int,
     def set_update_state(self, entity_id: int, state: Any):
         if isinstance(state.date, datetime):
             date_ts = int(state.date.replace(tzinfo=timezone.utc).timestamp())
-        else:
-            date_ts = int(state.date)
+        else: date_ts = int(state.date)
         self._pts, self._qts, self._date, self._seq = state.pts, state.qts, date_ts, state.seq
 
     async def close(self) -> None: pass
@@ -118,18 +118,27 @@ def get_update_states(self) -> List[Tuple[int, int, int, int, int, int]]:
         if self._pts is None: return []
         return [(0, self._pts, self._qts, self._date, self._seq, 0)]
 
-    def process_entities(self, tlo: object) -> None: pass
-
     def get_input_entity(self, key: Any) -> Any:
+        if key == 0:
+            if self._self_user_id and self._self_access_hash:
+                return InputPeerUser(self._self_user_id, self._self_access_hash)
+        raise KeyError("Entity not found in DbSession cache. It should be populated by process_entities.")
+
+    def process_entities(self, tlo: object) -> None:
         """
-        Returns InputPeerSelf() if key is 0, which tells Telethon to use the
-        currently authorized user. This is the correct way to handle "self" lookups.
+        Processes a TLObject to find and cache the 'self' user entity.
+        This is the standard mechanism Telethon uses to update the session
+        with the current user's ID and access_hash after login.
         """
-        if key == 0:
-            return InputPeerSelf()
-        raise KeyError("Entity not found in DbSession cache (caching is not implemented).")
+        if not hasattr(tlo, '__iter__'):
+            tlo = (tlo,)
+        
+        for entity in tlo:
+            if isinstance(entity, User) and entity.is_self:
+                self._self_user_id = entity.id
+                self._self_access_hash = entity.access_hash
+                # No need to save immediately, Telethon will call .save() when it's appropriate.
+                break
 
     def cache_file(self, md5_digest: bytes, file_size: int, instance: Any) -> None: pass
-
-    def get_file(self, md5_digest: bytes, file_size: int, exact: bool = True) -> Optional[Any]:
-        return None
+    def get_file(self, md5_digest: bytes, file_size: int, exact: bool = True) -> Optional[Any]: return None