fix: access_hash gathering

An0nX · An0nX · commit 7806dd05e420 · 2025-09-21T20:20:42.000+03:00
diff --git a/scripts/manage_account.py b/scripts/manage_account.py
@@ -4,7 +4,7 @@
 import getpass
 import argparse
 import logging
-from typing import Dict
+from typing import Dict, Optional
 
 # Add project root to sys.path
 project_root = Path(__file__).resolve().parent.parent
@@ -19,25 +19,29 @@
 import userbot.src.db_manager as db_manager
 from userbot.src.encrypt import encryption_manager
 
-logging.basicConfig(level=logging.WARNING)
+logging.basicConfig(level=logging.WARNING, format='%(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger('manage_account_cli')
 logger.setLevel(logging.INFO)
 
-async def add_account_logic(args):
+async def add_account_logic(args: argparse.Namespace) -> None:
     """
-    Logic to add a new account with a full interactive Telethon login session.
+    Handles the logic for adding a new account with a full interactive session.
+    Allows for detailed configuration of device, proxy, and other settings.
+
+    Args:
+        args (argparse.Namespace): The command-line arguments, expecting `args.name`.
     """
     logger.info(f"--- Adding new account: {args.name} ---")
     
-    temp_client = None
+    temp_client: Optional[TelegramClient] = None
     try:
         async with get_db() as db:
             if await db_manager.get_account(db, args.name):
                 logger.error(f"Error: Account '{args.name}' already exists.")
                 return
 
-            api_id = input("Enter API ID: ").strip()
-            api_hash = getpass.getpass("Enter API Hash: ").strip()
+            api_id: str = input("Enter API ID: ").strip()
+            api_hash: str = getpass.getpass("Enter API Hash: ").strip()
 
             if not api_id.isdigit() or not api_hash:
                 logger.error("Error: API ID must be a number and API Hash cannot be empty.")
@@ -48,33 +52,80 @@ async def add_account_logic(args):
             await temp_client.connect()
 
             if not await temp_client.is_user_authorized():
-                phone_number = input("Session not found. Please enter your phone number (e.g., +1234567890): ").strip()
+                phone_number: str = input("Session not found. Please enter your phone number (e.g., +1234567890): ").strip()
                 await temp_client.send_code_request(phone_number)
-                code = input("Please enter the code you received: ").strip()
+                code: str = input("Please enter the code you received: ").strip()
                 try:
                     await temp_client.sign_in(phone_number, code)
                 except SessionPasswordNeededError:
-                    two_fa_password = getpass.getpass("2FA Password required: ").strip()
+                    two_fa_password: str = getpass.getpass("2FA Password required: ").strip()
                     await temp_client.sign_in(password=two_fa_password)
 
             me = await temp_client.get_me(input_peer=True)
-            user_id = me.user_id
-            access_hash = me.access_hash
+            user_id: int = me.user_id
+            access_hash: int = me.access_hash
             logger.info(f"Successfully logged in as user ID: {user_id}.")
 
             existing_by_id = await db_manager.get_account_by_user_id(db, user_id)
             if existing_by_id:
                 logger.error(f"Error: This Telegram account (ID: {user_id}) already exists as '{existing_by_id.account_name}'.")
                 return
 
-            device_details = _generate_random_device()
-            lang_code = input(f"Enter language code (e.g., ru, en) [ru]: ").strip() or 'ru'
-            is_enabled = (input("Enable this account now? (yes/no) [yes]: ").strip().lower() or 'yes').startswith('y')
-            
+            # --- Interactive Configuration ---
+            print("\n--- Account Configuration ---")
+            lang_code: str = input(f"Enter language code (e.g., ru, en) [ru]: ").strip() or 'ru'
+            is_enabled: bool = (input("Enable this account now? (yes/no) [yes]: ").strip().lower() or 'yes').startswith('y')
+
+            # Device Configuration
+            if (input("Configure custom device details? (yes/no) [no]: ").strip().lower()).startswith('y'):
+                device_model: str = input("Enter device model: ").strip()
+                system_version: str = input("Enter system version: ").strip()
+                app_version: str = input("Enter app version: ").strip()
+            else:
+                device_details: Dict[str, str] = _generate_random_device()
+                device_model = device_details['device_model']
+                system_version = device_details['system_version']
+                app_version = device_details['app_version']
+                logger.info("Generated random device details.")
+
+            # Proxy Configuration
+            proxy_type: Optional[str] = None
+            proxy_ip: Optional[str] = None
+            proxy_port: Optional[int] = None
+            proxy_username: Optional[str] = None
+            proxy_password: Optional[str] = None
+            if (input("Configure a proxy for this account? (yes/no) [no]: ").strip().lower()).startswith('y'):
+                proxy_type_input = ""
+                while proxy_type_input not in ["http", "socks4", "socks5"]:
+                    proxy_type_input = input("Enter proxy type (http, socks4, socks5): ").strip().lower()
+                proxy_type = proxy_type_input
+                proxy_ip = input("Enter proxy IP address: ").strip()
+                proxy_port_str = ""
+                while not proxy_port_str.isdigit():
+                    proxy_port_str = input("Enter proxy port: ").strip()
+                proxy_port = int(proxy_port_str)
+                if (input("Does the proxy require authentication? (yes/no) [no]: ").strip().lower()).startswith('y'):
+                    proxy_username = input("Enter proxy username: ").strip()
+                    proxy_password = getpass.getpass("Enter proxy password: ").strip()
+
+            logger.info("\nSaving account to the database...")
             new_account = await db_manager.add_account(
-                db, args.name, api_id, api_hash, lang_code, is_enabled,
-                device_details['device_model'], device_details['system_version'], device_details['app_version'], 
-                user_id, access_hash
+                db,
+                account_name=args.name,
+                api_id=api_id,
+                api_hash=api_hash,
+                lang_code=lang_code,
+                is_enabled=is_enabled,
+                device_model=device_model,
+                system_version=system_version,
+                app_version=app_version,
+                user_telegram_id=user_id,
+                access_hash=access_hash,
+                proxy_type=proxy_type,
+                proxy_ip=proxy_ip,
+                proxy_port=proxy_port,
+                proxy_username=proxy_username,
+                proxy_password=proxy_password
             )
             
             if new_account:
@@ -89,7 +140,7 @@ async def add_account_logic(args):
             await temp_client.disconnect()
 
 
-async def edit_account_logic(args):
+async def edit_account_logic(args: argparse.Namespace) -> None:
     """Logic to edit an account's properties."""
     logger.info(f"--- Editing account: {args.name} ---")
     async with get_db() as db:
@@ -108,7 +159,8 @@ async def edit_account_logic(args):
 
 # ... (other logic functions like toggle, delete)
 
-async def main():
+async def main() -> None:
+    """The main entry point for the CLI script."""
     parser = argparse.ArgumentParser(description="DeBot Account Management CLI")
     subparsers = parser.add_subparsers(dest="command", required=True)
 
@@ -118,7 +170,7 @@ async def main():
     
     # ... (full parser setup as in previous correct answer)
 
-    args = parser.parse_args()
+    args: argparse.Namespace = parser.parse_args()
     
     await initialize_database()
     if hasattr(args, 'func'):
diff --git a/userbot/src/core_handlers.py b/userbot/src/core_handlers.py
@@ -84,7 +84,7 @@ async def add_account_handler(event: events.NewMessage.Event):
             
             await conv.send_message(await event.client.get_string("verifying_creds"))
             temp_client = TelegramClient(StringSession(), int(api_id), api_hash)
-            user_id = None
+            
             try:
                 await temp_client.connect()
                 if not await temp_client.is_user_authorized():
@@ -102,7 +102,9 @@ async def add_account_handler(event: events.NewMessage.Event):
                 await pass_resp.delete()
                 await temp_client.sign_in(password=two_fa_pass)
             
-            me = await temp_client.get_me(); user_id = me.id
+            me = await temp_client.get_me(input_peer=True)
+            user_id = me.user_id
+            access_hash = me.access_hash
             await temp_client.disconnect()
 
             async with get_db() as db:
@@ -119,8 +121,17 @@ async def add_account_handler(event: events.NewMessage.Event):
             await conv.send_message(await event.client.get_string("saving_to_db"))
             async with get_db() as db:
                 new_acc = await db_manager.add_account(
-                    db, account_name, api_id, api_hash, lang_code, is_enabled,
-                    device_details['device_model'], device_details['system_version'], device_details['app_version'], user_id
+                    db,
+                    account_name=account_name,
+                    api_id=api_id,
+                    api_hash=api_hash,
+                    lang_code=lang_code,
+                    is_enabled=is_enabled,
+                    device_model=device_details['device_model'],
+                    system_version=device_details['system_version'],
+                    app_version=device_details['app_version'],
+                    user_telegram_id=user_id,
+                    access_hash=access_hash
                 )
             
             if new_acc:
diff --git a/userbot/src/db_manager.py b/userbot/src/db_manager.py
@@ -9,6 +9,7 @@
 
 from userbot.src.encrypt import encryption_manager
 from userbot.src.db.models import Account, Session, Module, AccountModule, Log, ModuleData
+from userbot.src.db.session import get_db
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -24,21 +25,62 @@ async def add_account(
     system_version: str,
     app_version: str,
     user_telegram_id: int,
-    access_hash: int
+    access_hash: int,
+    proxy_type: Optional[str] = None,
+    proxy_ip: Optional[str] = None,
+    proxy_port: Optional[int] = None,
+    proxy_username: Optional[str] = None,
+    proxy_password: Optional[str] = None
 ) -> Optional[Account]:
-    """Adds a new account, including its access_hash, to the database."""
+    """
+    Adds a new account with its full configuration to the database.
+
+    Args:
+        db (AsyncSession): The database session.
+        account_name (str): The unique name for the account.
+        api_id (str): The user's API ID.
+        api_hash (str): The user's API Hash.
+        lang_code (str): The language code for the account.
+        is_enabled (bool): Whether the account should be active.
+        device_model (str): The device model for the session.
+        system_version (str): The system version for the session.
+        app_version (str): The app version for the session.
+        user_telegram_id (int): The Telegram user ID.
+        access_hash (int): The user's access hash.
+        proxy_type (Optional[str]): The type of proxy (e.g., 'http', 'socks5').
+        proxy_ip (Optional[str]): The proxy IP address.
+        proxy_port (Optional[int]): The proxy port.
+        proxy_username (Optional[str]): The username for proxy authentication.
+        proxy_password (Optional[str]): The password for proxy authentication.
+
+    Returns:
+        Optional[Account]: The newly created Account object, or None on failure.
+    """
     try:
+        encrypted_proxy_user = None
+        if proxy_username:
+            encrypted_proxy_user = encryption_manager.encrypt(proxy_username.encode('utf-8'))
+
+        encrypted_proxy_pass = None
+        if proxy_password:
+            encrypted_proxy_pass = encryption_manager.encrypt(proxy_password.encode('utf-8'))
+
         new_account = Account(
             account_name=account_name,
-            api_id=encryption_manager.encrypt(str(api_id).encode('utf-8')),
-            api_hash=encryption_manager.encrypt(str(api_hash).encode('utf-8')),
+            api_id=encryption_manager.encrypt(api_id.encode('utf-8')),
+            api_hash=encryption_manager.encrypt(api_hash.encode('utf-8')),
             lang_code=lang_code,
             is_enabled=is_enabled,
             device_model=device_model,
             system_version=system_version,
             app_version=app_version,
             user_telegram_id=user_telegram_id,
-            access_hash=access_hash
+            access_hash=access_hash,
+            proxy_type=proxy_type,
+            proxy_ip=proxy_ip,
+            proxy_port=proxy_port,
+            proxy_username=encrypted_proxy_user,
+            proxy_password=encrypted_proxy_pass
         )
         db.add(new_account)
         await db.flush()
@@ -99,6 +141,24 @@ async def update_account_lang(db: AsyncSession, account_id: int, lang_code: str)
     result = await db.execute(stmt)
     return result.rowcount > 0
 
+async def update_account_self_info(account_id: int, user_id: int, access_hash: int) -> None:
+    """
+    Updates the user_telegram_id and access_hash for an account.
+
+    This is useful if the information changes or was missing initially.
+
+    Args:
+        account_id (int): The primary key of the account to update.
+        user_id (int): The Telegram user ID.
+        access_hash (int): The user's access hash.
+    """
+    async with get_db() as db:
+        stmt = update(Account).where(Account.account_id == account_id).values(
+            user_telegram_id=user_id,
+            access_hash=access_hash
+        )
+        await db.execute(stmt)
+
 # --- 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."""
diff --git a/userbot/src/db_session.py b/userbot/src/db_session.py
@@ -128,7 +128,7 @@ 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. Self-user ID and access hash were not available on session creation.")
+        raise KeyError("Self-user ID and/or access hash were not found in the database for this account.")
 
     def process_entities(self, tlo: object) -> None:
         if not hasattr(tlo, '__iter__'):
@@ -140,6 +140,10 @@ def process_entities(self, tlo: object) -> None:
                     logger.info(f"Self-entity info for account {self.account_id} updated in-session.")
                     self._self_user_id = entity.id
                     self._self_access_hash = entity.access_hash
+                    # Persist the newly discovered access_hash to the database
+                    # This is a good place to ensure it's always up-to-date
+                    from asyncio import create_task
+                    create_task(db_manager.update_account_self_info(self.account_id, entity.id, entity.access_hash))
                 break
     
     def cache_file(self, md5_digest: bytes, file_size: int, instance: Any) -> None: pass
