Update bot/handlers.py

bot/handlers.py

@@ -1,5 +1,5 @@
 # PATH: bot/handlers.py
-import asyncio, json, re
+import asyncio, json, re, time
 from typing import Dict
 
 from hydrogram import Client, filters
@@ -14,19 +14,21 @@ from bot.core.auth import is_owner_id, require_allowed
 from bot.core.progress import SpeedETA, human_bytes, human_eta
 from bot.core.tasks import create_task, set_task
 
-from bot.integrations.diag_extra import dns_check
 from bot.integrations.cf_worker1 import profile_add, profile_list, profile_set_default
 from bot.integrations.cf_worker2 import (
     allow_user, disallow_user, pick_profile, access_token, record_upload, stats_today
 )
 
 from bot.telegram.media import download_to_temp
-from bot.telegram.parse import extract_media_message, extract_title_description
+from bot.telegram.parse import extract_title_description
 from bot.telegram.replies import safe_reply, safe_edit
 from bot.temp.files import cleanup_file
 
 from bot.youtube.uploader import upload_video
 
+# ✅ NEW: extra dns helper for /diag
+from bot.integrations.diag_extra import dns_check
+
 # ---------- in-memory auth state (simple FSM) ----------
 # values: "json" | "ci"
 _AWAIT_AUTH: Dict[int, str] = {}
@@ -47,7 +49,7 @@ def setup_handlers(app: Client) -> None:
         uid = m.from_user.id if m.from_user else None
         await safe_reply(m, f"👤 Your ID: `{uid}`")
 
-    # -------- OWNER allow/disallow via reply OR forward OR id --------
+    # -------- OWNER allow/disallow via forward OR id --------
     @app.on_message(filters.command(["allow"]) & filters.private)
     async def allow_cmd(_: Client, m: Message):
         if not is_owner_id(m.from_user.id if m.from_user else None):
@@ -64,12 +66,7 @@ def setup_handlers(app: Client) -> None:
                 target = int(parts[1])
 
         if not target:
-            return await safe_reply(
-                m,
-                "Usage:\n"
-                "• Reply/forward user msg then /allow\n"
-                "• or /allow <user_id>"
-            )
+            return await safe_reply(m, "Usage:\n• Reply/forward user msg then /allow\n• or /allow <user_id>")
 
         j = await allow_user(target)
         await safe_reply(m, f"✅ allowed: `{target}`\n`{j}`")
@@ -82,7 +79,6 @@ def setup_handlers(app: Client) -> None:
         parts = (m.text or "").split()
         if len(parts) < 2 or not parts[1].isdigit():
             return await safe_reply(m, "Usage: /disallow <user_id>")
-
         target = int(parts[1])
         j = await disallow_user(target)
         await safe_reply(m, f"✅ disallowed: `{target}`\n`{j}`")
@@ -101,28 +97,19 @@ def setup_handlers(app: Client) -> None:
             return await safe_reply(m, texts.OWNER_ONLY)
 
         from bot.integrations.http import get_client
-
         c = get_client()
-        headers = {"User-Agent": "hf-diag/1.0", "Accept": "*/*"}
 
-        # HTTP checks
-        if not Workers.WORKER1_URL:
-            r1 = "ERR:missing_WORKER1_URL"
-        else:
-            try:
-                res1 = await c.get(Workers.WORKER1_URL + "/", headers=headers)
-                r1 = f"{res1.status_code}"
-            except Exception as e:
-                r1 = f"ERR:{type(e).__name__}:{e}"
+        try:
+            r = await c.get((Workers.WORKER1_URL or "") + "/")
+            r1 = f"{r.status_code}"
+        except Exception as e:
+            r1 = f"ERR:{type(e).__name__}:{e}"
 
-        if not Workers.WORKER2_URL:
-            r2 = "ERR:missing_WORKER2_URL"
-        else:
-            try:
-                res2 = await c.get(Workers.WORKER2_URL + "/", headers=headers)
-                r2 = f"{res2.status_code}"
-            except Exception as e:
-                r2 = f"ERR:{type(e).__name__}:{e}"
+        try:
+            r = await c.get((Workers.WORKER2_URL or "") + "/")
+            r2 = f"{r.status_code}"
+        except Exception as e:
+            r2 = f"ERR:{type(e).__name__}:{e}"
 
         await safe_reply(
             m,
@@ -141,12 +128,13 @@ def setup_handlers(app: Client) -> None:
         ok = await require_allowed(uid)
         if not ok:
             return await safe_reply(m, texts.NOT_ALLOWED)
+
         _AWAIT_AUTH.pop(uid, None)
         await safe_reply(m, texts.AUTH_MENU, reply_markup=auth_menu_keyboard())
 
     @app.on_message(filters.command(["cancel"]) & filters.private)
     async def cancel_cmd(_: Client, m: Message):
-        _AWAIT_AUTH.pop(m.from_user.id if m.from_user else 0, None)
+        _AWAIT_AUTH.pop(m.from_user.id, None)
         await safe_reply(m, texts.CANCELLED)
 
     @app.on_callback_query()
@@ -157,20 +145,17 @@ def setup_handlers(app: Client) -> None:
         if action in (AUTH_JSON, AUTH_CI, CANCEL):
             if action == CANCEL:
                 _AWAIT_AUTH.pop(uid, None)
-                await q.answer("Cancelled")
-                return
+                return await q.answer("Cancelled")
             if action == AUTH_JSON:
                 _AWAIT_AUTH[uid] = "json"
                 await q.message.edit_text(texts.ASK_JSON)
-                await q.answer("Paste JSON")
-                return
+                return await q.answer("Paste JSON")
             if action == AUTH_CI:
                 _AWAIT_AUTH[uid] = "ci"
                 await q.message.edit_text(texts.ASK_ID_SECRET)
-                await q.answer("Send ID & Secret")
-                return
+                return await q.answer("Send ID & Secret")
+            return
 
-        # set default profile callback
         if action == "setdef" and value:
             if not await require_allowed(uid):
                 return await q.answer("Not allowed", show_alert=True)
@@ -188,17 +173,13 @@ def setup_handlers(app: Client) -> None:
 
         await q.answer("OK")
 
-    # -------- auth message sink --------
-    @app.on_message(
-        filters.private
-        & filters.text
-        & ~filters.command(["start", "ping", "me", "auth", "cancel", "profiles", "stats", "allow", "disallow", "diag"])
-    )
+    # message sink for auth states
+    @app.on_message(filters.private & filters.text & ~filters.command(["start","ping","me","auth","cancel","profiles","stats","allow","disallow","diag"]))
     async def auth_sink(_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
         mode = _AWAIT_AUTH.get(uid)
         if not mode:
-            return  # not in auth flow
+            return
 
         try:
             if mode == "json":
@@ -209,7 +190,7 @@ def setup_handlers(app: Client) -> None:
                 if not client_id or not client_secret:
                     return await safe_reply(m, texts.PARSE_FAIL)
             else:
-                txt = (m.text or "").strip()
+                txt = m.text.strip()
                 parts = [p.strip() for p in re.split(r"[|\n]+", txt) if p.strip()]
                 if len(parts) < 2:
                     return await safe_reply(m, texts.PARSE_FAIL)
@@ -239,10 +220,7 @@ def setup_handlers(app: Client) -> None:
         profiles = j.get("profiles") or []
         txt = f"Default: `{j.get('default_profile_id')}`\n\n"
         for p in profiles:
-            txt += (
-                f"- `{p['profile_id']}` | {p.get('label')} | "
-                f"refresh={p.get('has_refresh')} | ch={p.get('channel_title')}\n"
-            )
+            txt += f"- `{p['profile_id']}` | {p.get('label')} | refresh={p.get('has_refresh')} | ch={p.get('channel_title')}\n"
 
         kb = profiles_keyboard(profiles)
         await safe_reply(m, txt, reply_markup=kb)
@@ -266,7 +244,6 @@ def setup_handlers(app: Client) -> None:
         default_id = pl.get("default_profile_id")
         profiles = pl.get("profiles") or []
         dp = next((p for p in profiles if p.get("profile_id") == default_id), None)
-
         if not dp or not dp.get("has_refresh"):
             return await safe_reply(m, texts.NEED_AUTH)
 
@@ -277,13 +254,11 @@ def setup_handlers(app: Client) -> None:
         pick = await pick_profile(uid, channel_id)
         if not pick.get("ok"):
             return await safe_reply(m, texts.PICK_FAIL.format(pick))
-
         use_profile_id = pick["profile_id"]
 
         tok = await access_token(uid, use_profile_id)
         if not tok.get("ok"):
             return await safe_reply(m, texts.TOKEN_FAIL.format(tok))
-
         access_tok = tok["access_token"]
 
         status = await safe_reply(m, texts.UPLOAD_START)
@@ -295,7 +270,6 @@ def setup_handlers(app: Client) -> None:
         file_path = ""
         try:
             file_path, file_size, file_name = await download_to_temp(app_, m)
-
             title, desc = extract_title_description(m, file_name)
 
             se = SpeedETA()
@@ -304,7 +278,6 @@ def setup_handlers(app: Client) -> None:
             async def prog(done: int, total: int):
                 nonlocal last_edit
                 snap = se.update(done, total)
-                import time
                 if time.time() - last_edit < 2.0 and done < total:
                     return
                 last_edit = time.time()
@@ -316,13 +289,14 @@ def setup_handlers(app: Client) -> None:
                 await safe_edit(status, txt)
 
             set_task(task_id, "uploading", "")
+
             yt_url = await upload_video(
                 access_tok,
                 file_path,
                 title,
                 desc,
                 privacy="private",
-                progress_cb=prog
+                progress_cb=prog,
             )
 
             await record_upload(uid, use_profile_id)
@@ -333,7 +307,6 @@ def setup_handlers(app: Client) -> None:
         except Exception as e:
             set_task(task_id, "error", str(e))
             await safe_reply(m, f"❌ Upload failed: `{type(e).__name__}: {e}`")
-
         finally:
             if file_path:
                 cleanup_file(file_path)