Upload 2 files

utils/clients.py

@@ -1,144 +1,144 @@
-import asyncio, config
-from pathlib import Path
-from pyrogram import Client
-from utils.directoryHandler import backup_drive_data, loadDriveData
-from utils.logger import Logger
-import os
-import signal
-
-logger = Logger(__name__)
-
-multi_clients = {}
-premium_clients = {}
-work_loads = {}
-premium_work_loads = {}
-main_bot = None
-
-
-async def initialize_clients():
-    global multi_clients, work_loads, premium_clients, premium_work_loads
-    logger.info("Initializing Clients")
-
-    session_cache_path = Path(f"./cache")
-    session_cache_path.parent.mkdir(parents=True, exist_ok=True)
-
-    all_tokens = dict((i, t) for i, t in enumerate(config.BOT_TOKENS, start=1))
-    all_sessions = dict(
-        (i, s) for i, s in enumerate(config.STRING_SESSIONS, start=len(all_tokens) + 1)
-    )
-
-    # Scan for existing session files to allow starting without secrets
-    if session_cache_path.exists():
-        for file_path in session_cache_path.glob("*.session"):
-            if file_path.stem.isdigit():
-                client_id = int(file_path.stem)
-                # Only add if not already in config lists
-                if client_id not in all_tokens and client_id not in all_sessions:
-                    logger.info(f"Found existing session file for Client {client_id}")
-                    all_tokens[client_id] = None
-
-    async def start_client(client_id, token, type):
-        name = str(client_id)
-        for attempt in range(2):
-            try:
-                if attempt == 0:
-                    logger.info(f"Starting - {type.title()} Client {client_id}")
-                else:
-                    logger.info(f"Retrying - {type.title()} Client {client_id}")
-
-                if type == "bot":
-                    client = Client(
-                        name=name,
-                        api_id=config.API_ID,
-                        api_hash=config.API_HASH,
-                        bot_token=token,
-                        workdir=session_cache_path,
-                    )
-                    client.loop = asyncio.get_running_loop()
-                    await client.start()
-                    await client.send_message(
-                        config.STORAGE_CHANNEL,
-                        f"Started - {type.title()} Client {client_id}",
-                    )
-                    multi_clients[client_id] = client
-                    work_loads[client_id] = 0
-                elif type == "user":
-                    client = await Client(
-                        name=name,
-                        api_id=config.API_ID,
-                        api_hash=config.API_HASH,
-                        session_string=token,
-                        sleep_threshold=config.SLEEP_THRESHOLD,
-                        workdir=session_cache_path,
-                        no_updates=True,
-                    ).start()
-                    await client.send_message(
-                        config.STORAGE_CHANNEL,
-                        f"Started - {type.title()} Client {client_id}",
-                    )
-                    premium_clients[client_id] = client
-                    premium_work_loads[client_id] = 0
-
-                logger.info(f"Started - {type.title()} Client {client_id}")
-                return
-            except Exception as e:
-                logger.error(
-                    f"Failed To Start {type.title()} Client - {client_id} Error: {e}"
-                )
-                if attempt == 0:
-                    # Try to close client to release file locks
-                    if 'client' in locals():
-                        try:
-                            await client.stop()
-                        except:
-                            pass
-                            
-                    session_path = session_cache_path / f"{name}.session"
-                    if session_path.exists():
-                        logger.info(f"Deleting invalid session file: {session_path}")
-                        try:
-                            session_path.unlink()
-                        except Exception as delete_error:
-                            logger.error(f"Could not delete session file: {delete_error}")
-                else:
-                    logger.error(f"Client {client_id} failed to start after retry.")
-
-    await asyncio.gather(
-        *(
-            [
-                start_client(client_id, client, "bot")
-                for client_id, client in all_tokens.items()
-            ]
-            + [
-                start_client(client_id, client, "user")
-                for client_id, client in all_sessions.items()
-            ]
-        )
-    )
-    if len(multi_clients) == 0:
-        logger.error("No Clients Were Initialized")
-        raise RuntimeError("No Clients Were Initialized - Please add BOT_TOKENS to config or ensure valid .session files are in the cache folder.")
-
-    if len(premium_clients) == 0:
-        logger.info("No Premium Clients Were Initialized")
-
-    logger.info("Clients Initialized")
-
-    # Load the drive data
-    await loadDriveData()
-
-    # Start the backup drive data task
-    asyncio.create_task(backup_drive_data())
-
-
-def get_client(premium_required=False) -> Client:
-    global multi_clients, work_loads, premium_clients, premium_work_loads
-
-    if premium_required:
-        index = min(premium_work_loads, key=premium_work_loads.get)
-        premium_work_loads[index] += 1
-        return premium_clients[index]
-
-    index = min(work_loads, key=work_loads.get)
-    work_loads[index] += 1
-    return multi_clients[index]
+import asyncio, config
+from pathlib import Path
+from pyrogram import Client
+from utils.directoryHandler import backup_drive_data, loadDriveData
+from utils.logger import Logger
+import os
+import signal
+
+logger = Logger(__name__)
+
+multi_clients = {}
+premium_clients = {}
+work_loads = {}
+premium_work_loads = {}
+main_bot = None
+
+
+async def initialize_clients():
+    global multi_clients, work_loads, premium_clients, premium_work_loads
+    logger.info("Initializing Clients")
+
+    session_cache_path = Path(f"./cache")
+    session_cache_path.parent.mkdir(parents=True, exist_ok=True)
+
+    all_tokens = dict((i, t) for i, t in enumerate(config.BOT_TOKENS, start=1))
+    all_sessions = dict(
+        (i, s) for i, s in enumerate(config.STRING_SESSIONS, start=len(all_tokens) + 1)
+    )
+
+    # Scan for existing session files to allow starting without secrets
+    if session_cache_path.exists():
+        for file_path in session_cache_path.glob("*.session"):
+            if file_path.stem.isdigit():
+                client_id = int(file_path.stem)
+                # Only add if not already in config lists
+                if client_id not in all_tokens and client_id not in all_sessions:
+                    logger.info(f"Found existing session file for Client {client_id}")
+                    all_tokens[client_id] = None
+
+    async def start_client(client_id, token, type):
+        name = str(client_id)
+        for attempt in range(2):
+            try:
+                if attempt == 0:
+                    logger.info(f"Starting - {type.title()} Client {client_id}")
+                else:
+                    logger.info(f"Retrying - {type.title()} Client {client_id}")
+
+                if type == "bot":
+                    client = Client(
+                        name=name,
+                        api_id=config.API_ID,
+                        api_hash=config.API_HASH,
+                        bot_token=token,
+                        workdir=session_cache_path,
+                    )
+                    client.loop = asyncio.get_running_loop()
+                    await client.start()
+                    await client.send_message(
+                        config.STORAGE_CHANNEL,
+                        f"Started - {type.title()} Client {client_id}",
+                    )
+                    multi_clients[client_id] = client
+                    work_loads[client_id] = 0
+                elif type == "user":
+                    client = await Client(
+                        name=name,
+                        api_id=config.API_ID,
+                        api_hash=config.API_HASH,
+                        session_string=token,
+                        sleep_threshold=config.SLEEP_THRESHOLD,
+                        workdir=session_cache_path,
+                        no_updates=True,
+                    ).start()
+                    await client.send_message(
+                        config.STORAGE_CHANNEL,
+                        f"Started - {type.title()} Client {client_id}",
+                    )
+                    premium_clients[client_id] = client
+                    premium_work_loads[client_id] = 0
+
+                logger.info(f"Started - {type.title()} Client {client_id}")
+                return
+            except Exception as e:
+                logger.error(
+                    f"Failed To Start {type.title()} Client - {client_id} Error: {e}"
+                )
+                if attempt == 0:
+                    # Try to close client to release file locks
+                    if 'client' in locals():
+                        try:
+                            await client.stop()
+                        except:
+                            pass
+                            
+                    session_path = session_cache_path / f"{name}.session"
+                    if session_path.exists():
+                        logger.info(f"Deleting invalid session file: {session_path}")
+                        try:
+                            session_path.unlink()
+                        except Exception as delete_error:
+                            logger.error(f"Could not delete session file: {delete_error}")
+                else:
+                    logger.error(f"Client {client_id} failed to start after retry.")
+
+    await asyncio.gather(
+        *(
+            [
+                start_client(client_id, client, "bot")
+                for client_id, client in all_tokens.items()
+            ]
+            + [
+                start_client(client_id, client, "user")
+                for client_id, client in all_sessions.items()
+            ]
+        )
+    )
+    if len(multi_clients) == 0:
+        logger.error("No Clients Were Initialized")
+        raise RuntimeError("No Clients Were Initialized - Please add BOT_TOKENS to config or ensure valid .session files are in the cache folder.")
+
+    if len(premium_clients) == 0:
+        logger.info("No Premium Clients Were Initialized")
+
+    logger.info("Clients Initialized")
+
+    # Load the drive data
+    await loadDriveData()
+
+    # Start the backup drive data task
+    asyncio.create_task(backup_drive_data())
+
+
+def get_client(premium_required=False) -> Client:
+    global multi_clients, work_loads, premium_clients, premium_work_loads
+
+    if premium_required:
+        index = min(premium_work_loads, key=premium_work_loads.get)
+        premium_work_loads[index] += 1
+        return premium_clients[index]
+
+    index = min(work_loads, key=work_loads.get)
+    work_loads[index] += 1
+    return multi_clients[index]

utils/extra.py

@@ -66,21 +66,10 @@ def reset_cache_dir():
     shutil.rmtree(downloads_dir, ignore_errors=True)
     downloads_dir.mkdir(parents=True, exist_ok=True)
 
-    if cache_dir.exists():
-        for item in cache_dir.iterdir():
-            if item.is_dir():
-                shutil.rmtree(item, ignore_errors=True)
-            elif item.name.endswith(".session") or item.name == "drive.data":
-                continue
-            else:
-                try:
-                    item.unlink()
-                except Exception:
-                    pass
-    else:
+    if not cache_dir.exists():
         cache_dir.mkdir(parents=True, exist_ok=True)
 
-    logger.info("Cache and downloads directory reset")
+    logger.info("Cache and downloads directory checked")
 
 
 def parse_content_disposition(content_disposition):