Update bot/handlers.py

bot/handlers.py

@@ -1,19 +1,23 @@
 # PATH: bot/handlers.py
-import asyncio
+import asyncio, json, re
+from typing import Dict
+
 from hydrogram import Client, filters
 from hydrogram.types import Message, CallbackQuery
 
-from bot.config import Telegram
+from bot.config import Telegram, Workers
 from bot.ui import texts
-from bot.ui.keyboards import profiles_keyboard
-from bot.ui.callbacks import parse_cb
+from bot.ui.keyboards import profiles_keyboard, auth_menu_keyboard
+from bot.ui.callbacks import parse_cb, AUTH_JSON, AUTH_CI, CANCEL
 
 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.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.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
@@ -22,6 +26,10 @@ from bot.temp.files import cleanup_file
 
 from bot.youtube.uploader import upload_video
 
+# ---------- in-memory auth state (simple FSM) ----------
+# values: "json" | "ci"
+_AWAIT_AUTH: Dict[int, str] = {}
+
 def setup_handlers(app: Client) -> None:
 
     @app.on_message(filters.command(["start"]) & filters.private)
@@ -43,7 +51,6 @@ def setup_handlers(app: Client) -> None:
         if not is_owner_id(m.from_user.id if m.from_user else None):
             return await safe_reply(m, texts.OWNER_ONLY)
 
-        # If forward present -> use forwarded user id
         target = None
         if m.reply_to_message and m.reply_to_message.from_user:
             target = m.reply_to_message.from_user.id
@@ -68,12 +75,11 @@ 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}")
 
-    # -------- Owner stats --------
+    # -------- Owner stats + diag --------
     @app.on_message(filters.command(["stats"]) & filters.private)
     async def stats_cmd(_: Client, m: Message):
         if not is_owner_id(m.from_user.id if m.from_user else None):
@@ -81,27 +87,110 @@ def setup_handlers(app: Client) -> None:
         j = await stats_today()
         await safe_reply(m, f"📊 Stats:\n`{j}`")
 
-    # -------- Auth flow (Worker1) --------
+    @app.on_message(filters.command(["diag"]) & filters.private)
+    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}")
+
+    # -------- AUTH UI --------
     @app.on_message(filters.command(["auth"]) & filters.private)
     async def auth_cmd(_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
         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, None)
+        await safe_reply(m, texts.CANCELLED)
 
-        # expects: /auth <client_id> <client_secret>
-        parts = (m.text or "").split(maxsplit=2)
-        if len(parts) < 3:
-            return await safe_reply(m, "Usage:\n/auth <CLIENT_ID> <CLIENT_SECRET>\n(baad me JSON method add kar denge)")
+    @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")
+            if action == AUTH_JSON:
+                _AWAIT_AUTH[uid] = "json"
+                await q.message.edit_text(texts.ASK_JSON)
+                return await q.answer("Paste JSON")
+            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
+
+        # handle set default profile callbacks too
+        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")
 
-        client_id = parts[1].strip()
-        client_secret = parts[2].strip()
+    # 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
 
-        j = await profile_add(uid, client_id, client_secret, label="main", ttl_sec=600)
-        if not j.get("ok"):
-            return await safe_reply(m, f"❌ profile_add failed: `{j}`")
+        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()
+                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)
+                client_id, client_secret = parts[0], parts[1]
+
+            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"])
 
-        await safe_reply(m, texts.SENT_AUTH_LINK + j["login_url"])
+        except Exception as e:
+            await safe_reply(m, f"{texts.PARSE_FAIL}\n`{type(e).__name__}: {e}`")
 
     @app.on_message(filters.command(["profiles"]) & filters.private)
     async def profiles_cmd(_: Client, m: Message):
@@ -112,7 +201,7 @@ def setup_handlers(app: Client) -> None:
 
         j = await profile_list(uid)
         if not j.get("ok"):
-            return await safe_reply(m, f"❌ profile_list failed: `{j}`")
+            return await safe_reply(m, texts.PROFILE_LIST_FAIL.format(j))
 
         profiles = j.get("profiles") or []
         txt = f"Default: `{j.get('default_profile_id')}`\n\n"
@@ -122,26 +211,7 @@ def setup_handlers(app: Client) -> None:
         kb = profiles_keyboard(profiles)
         await safe_reply(m, txt, reply_markup=kb)
 
-    @app.on_callback_query()
-    async def cb(_: Client, q: CallbackQuery):
-        uid = q.from_user.id if q.from_user else 0
-        if not await require_allowed(uid):
-            return await q.answer("Not allowed", show_alert=True)
-
-        action, value = parse_cb(q.data or "")
-        if action == "setdef" and value:
-            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
-        else:
-            await q.answer("OK")
-
-    # -------- Upload handler (send video/document) --------
+    # -------- Upload handler --------
     @app.on_message(filters.private & (filters.video | filters.document))
     async def upload_handler(app_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
@@ -149,14 +219,13 @@ def setup_handlers(app: Client) -> None:
         if not ok:
             return await safe_reply(m, texts.NOT_ALLOWED)
 
-        media, kind = extract_media_message(m)
+        media = m.video or m.document
         if not media:
             return await safe_reply(m, "Send a video/document.")
 
-        # find default profile + channel_id from worker1 list
         pl = await profile_list(uid)
         if not pl.get("ok"):
-            return await safe_reply(m, f"❌ profile_list failed: `{pl}`")
+            return await safe_reply(m, texts.PROFILE_LIST_FAIL.format(pl))
 
         default_id = pl.get("default_profile_id")
         profiles = pl.get("profiles") or []
@@ -168,18 +237,16 @@ def setup_handlers(app: Client) -> None:
         if not channel_id:
             return await safe_reply(m, "❌ channel_id missing. Re-auth profile.")
 
-        # pick profile (rotation same channel only)
         pick = await pick_profile(uid, channel_id)
         if not pick.get("ok"):
-            return await safe_reply(m, f"❌ pick_profile failed: `{pick}`")
+            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, f"❌ access_token failed: `{tok}`")
+            return await safe_reply(m, texts.TOKEN_FAIL.format(tok))
         access_tok = tok["access_token"]
 
-        # status message
         status = await safe_reply(m, texts.UPLOAD_START)
 
         task_id = f"{uid}:{m.id}"
@@ -188,30 +255,23 @@ def setup_handlers(app: Client) -> None:
 
         file_path = ""
         try:
-            # download
             file_path, file_size, file_name = await download_to_temp(app_, m)
 
-            # metadata rules
             title, desc = extract_title_description(m, file_name)
 
-            # upload progress updates
             se = SpeedETA()
             last_edit = 0.0
 
             async def prog(done: int, total: int):
                 nonlocal last_edit
                 snap = se.update(done, total)
-                now = snap["done"]
-                if not status:
-                    return
-                # throttle edit
                 import time
                 if time.time() - last_edit < 2.0 and done < total:
                     return
                 last_edit = time.time()
                 txt = (
                     f"{texts.UPLOAD_TO_YT}\n"
-                    f"{human_bytes(now)}/{human_bytes(total)}\n"
+                    f"{human_bytes(done)}/{human_bytes(total)}\n"
                     f"speed: {human_bytes(snap['speed_bps'])}/s | eta: {human_eta(snap['eta_sec'])}"
                 )
                 await safe_edit(status, txt)
@@ -219,7 +279,6 @@ def setup_handlers(app: Client) -> None:
             set_task(task_id, "uploading", "")
             yt_url = await upload_video(access_tok, file_path, title, desc, privacy="private", progress_cb=prog)
 
-            # record upload counter
             await record_upload(uid, use_profile_id)
 
             set_task(task_id, "done", "", yt_url=yt_url)