Update bot/handlers.py

bot/handlers.py

@@ -14,6 +14,7 @@ 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
@@ -30,6 +31,7 @@ from bot.youtube.uploader import upload_video
 # values: "json" | "ci"
 _AWAIT_AUTH: Dict[int, str] = {}
 
+
 def setup_handlers(app: Client) -> None:
 
     @app.on_message(filters.command(["start"]) & filters.private)
@@ -45,7 +47,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 forward OR id --------
+    # -------- OWNER allow/disallow via reply OR 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):
@@ -62,10 +64,15 @@ 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}")
+        await safe_reply(m, f"✅ allowed: `{target}`\n`{j}`")
 
     @app.on_message(filters.command(["disallow"]) & filters.private)
     async def disallow_cmd(_: Client, m: Message):
@@ -75,9 +82,10 @@ 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}")
+        await safe_reply(m, f"✅ disallowed: `{target}`\n`{j}`")
 
     # -------- Owner stats + diag --------
     @app.on_message(filters.command(["stats"]) & filters.private)
@@ -91,21 +99,40 @@ def setup_handlers(app: Client) -> None:
     async def diag_cmd(_: Client, m: Message):
         if not is_owner_id(m.from_user.id if m.from_user else None):
             return await safe_reply(m, texts.OWNER_ONLY)
-        # quick DNS/connect check
+
         from bot.integrations.http import get_client
+
         c = get_client()
-        res1 = res2 = None
-        try:
-            res1 = await c.get((Workers.WORKER1_URL or "") + "/")
-            r1 = f"{res1.status_code}"
-        except Exception as e:
-            r1 = f"ERR:{type(e).__name__}:{e}"
-        try:
-            res2 = await c.get((Workers.WORKER2_URL or "") + "/")
-            r2 = f"{res2.status_code}"
-        except Exception as e:
-            r2 = f"ERR:{type(e).__name__}:{e}"
-        await safe_reply(m, f"🔎 DIAG\nW1={Workers.WORKER1_URL} -> {r1}\nW2={Workers.WORKER2_URL} -> {r2}")
+        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}"
+
+        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}"
+
+        await safe_reply(
+            m,
+            "🔎 DIAG\n"
+            f"W1={Workers.WORKER1_URL} -> {r1}\n"
+            f"W2={Workers.WORKER2_URL} -> {r2}\n\n"
+            f"DNS W1 -> {dns_check(Workers.WORKER1_URL)}\n"
+            f"DNS W2 -> {dns_check(Workers.WORKER2_URL)}\n"
+            f"DNS google -> {dns_check('https://www.google.com')}\n"
+        )
 
     # -------- AUTH UI --------
     @app.on_message(filters.command(["auth"]) & filters.private)
@@ -119,46 +146,54 @@ def setup_handlers(app: Client) -> None:
 
     @app.on_message(filters.command(["cancel"]) & filters.private)
     async def cancel_cmd(_: Client, m: Message):
-        _AWAIT_AUTH.pop(m.from_user.id, None)
+        _AWAIT_AUTH.pop(m.from_user.id if m.from_user else 0, None)
         await safe_reply(m, texts.CANCELLED)
 
     @app.on_callback_query()
     async def cb(_: Client, q: CallbackQuery):
         uid = q.from_user.id if q.from_user else 0
         action, value = parse_cb(q.data or "")
+
         if action in (AUTH_JSON, AUTH_CI, CANCEL):
             if action == CANCEL:
                 _AWAIT_AUTH.pop(uid, None)
-                return await q.answer("Cancelled")
+                await q.answer("Cancelled")
+                return
             if action == AUTH_JSON:
                 _AWAIT_AUTH[uid] = "json"
                 await q.message.edit_text(texts.ASK_JSON)
-                return await q.answer("Paste JSON")
+                await q.answer("Paste JSON")
+                return
             if action == AUTH_CI:
                 _AWAIT_AUTH[uid] = "ci"
                 await q.message.edit_text(texts.ASK_ID_SECRET)
-                return await q.answer("Send ID & Secret")
-            return
+                await q.answer("Send ID & Secret")
+                return
 
-        # handle set default profile callbacks too
+        # set default profile callback
         if action == "setdef" and value:
-            # owner/allowed only
-            from bot.core.auth import require_allowed
             if not await require_allowed(uid):
                 return await q.answer("Not allowed", show_alert=True)
+
             j = await profile_set_default(uid, value)
             if not j.get("ok"):
                 return await q.answer("Failed", show_alert=True)
+
             await q.answer("Default set ✅", show_alert=False)
             try:
                 await q.message.edit_text(f"✅ Default set: `{value}`")
             except Exception:
                 pass
             return
+
         await q.answer("OK")
 
-    # message sink for auth states
-    @app.on_message(filters.private & filters.text & ~filters.command(["start","ping","me","auth","cancel","profiles","stats","allow","disallow","diag"]))
+    # -------- auth message sink --------
+    @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)
@@ -167,17 +202,14 @@ def setup_handlers(app: Client) -> None:
 
         try:
             if mode == "json":
-                # parse JSON
                 data = json.loads(m.text)
-                # accept both shapes
                 root = data.get("installed") or data.get("web") or {}
                 client_id = root.get("client_id") or data.get("client_id")
                 client_secret = root.get("client_secret") or data.get("client_secret")
                 if not client_id or not client_secret:
                     return await safe_reply(m, texts.PARSE_FAIL)
             else:
-                # expect "id\nsecret" or "id | secret"
-                txt = m.text.strip()
+                txt = (m.text or "").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)
@@ -186,6 +218,7 @@ def setup_handlers(app: Client) -> None:
             j = await profile_add(uid, client_id, client_secret, label="main", ttl_sec=600)
             if not j.get("ok"):
                 return await safe_reply(m, texts.PROFILE_ADD_FAIL.format(j))
+
             _AWAIT_AUTH.pop(uid, None)
             await safe_reply(m, texts.SENT_AUTH_LINK + j["login_url"])
 
@@ -206,7 +239,10 @@ 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')} | refresh={p.get('has_refresh')} | ch={p.get('channel_title')}\n"
+            txt += (
+                f"- `{p['profile_id']}` | {p.get('label')} | "
+                f"refresh={p.get('has_refresh')} | ch={p.get('channel_title')}\n"
+            )
 
         kb = profiles_keyboard(profiles)
         await safe_reply(m, txt, reply_markup=kb)
@@ -230,6 +266,7 @@ 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)
 
@@ -240,11 +277,13 @@ 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)
@@ -277,7 +316,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)
+            yt_url = await upload_video(
+                access_tok,
+                file_path,
+                title,
+                desc,
+                privacy="private",
+                progress_cb=prog
+            )
 
             await record_upload(uid, use_profile_id)
 
@@ -287,6 +333,7 @@ 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)