Update bot/handlers.py

bot/handlers.py

@@ -1,5 +1,9 @@
 # PATH: bot/handlers.py
-import json, re, time
+import asyncio
+import json
+import os
+import re
+import time
 from typing import Dict
 
 from hydrogram import Client, filters
@@ -14,9 +18,23 @@ 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.core.speedtest import (
+    ping_ms,
+    net_download_test,
+    net_upload_test,
+    disk_total_free,
+    bytes_to_mb,
+    bytes_per_sec_to_mb_s,
+)
+
 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
+    allow_user,
+    disallow_user,
+    pick_profile,
+    access_token,
+    record_upload,
+    stats_today,
 )
 
 from bot.telegram.media import download_to_temp
@@ -32,6 +50,10 @@ from bot.integrations.http import fetch_status
 # ---------- in-memory auth state (simple FSM) ----------
 _AWAIT_AUTH: Dict[int, str] = {}
 
+# uptime + speedtest cooldown
+_STARTED_AT = time.time()
+_SPEED_COOLDOWN: dict[int, float] = {}
+
 
 def _pick_login_url(j: dict) -> str:
     """Accept login_url from either top-level OR nested {data:{...}} shapes."""
@@ -57,7 +79,6 @@ def _pick_login_url(j: dict) -> str:
 
 
 def setup_handlers(app: Client) -> None:
-
     @app.on_message(filters.command(["start"]) & filters.private)
     async def start_handler(_: Client, m: Message):
         await safe_reply(m, texts.START_TEXT)
@@ -71,6 +92,54 @@ 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}`")
 
+    # -------------------- SPEEDTEST (any user) --------------------
+    @app.on_message(filters.command(["speedtest"]) & filters.private)
+    async def speedtest_cmd(_: Client, m: Message):
+        uid = m.from_user.id if m.from_user else 0
+
+        now = time.time()
+        last = _SPEED_COOLDOWN.get(uid, 0.0)
+        if now - last < 30:
+            return await safe_reply(m, "⏳ Wait 30s then try again.")
+        _SPEED_COOLDOWN[uid] = now
+
+        msg = await safe_reply(m, "⚡ Running speedtest…")
+
+        # uptime
+        up_sec = int(max(0, time.time() - _STARTED_AT))
+        up_txt = human_eta(up_sec)
+
+        # disk
+        d = disk_total_free("/")
+        total_gb = d["total"] / (1024**3)
+        free_gb = d["free"] / (1024**3)
+
+        # run parallel
+        dl_task = asyncio.create_task(net_download_test())
+        up_task = asyncio.create_task(net_upload_test())
+        p_task = asyncio.create_task(ping_ms())
+
+        dl = await dl_task
+        up = await up_task
+        p = await p_task
+
+        ping_line = f"{p:.0f} ms" if p is not None else "N/A"
+
+        dl_mb_s = bytes_per_sec_to_mb_s(dl["bps"])
+        up_mb_s = bytes_per_sec_to_mb_s(up["bps"])
+
+        txt = (
+            "⚡ **Speedtest**\n\n"
+            f"🕒 Uptime: **{up_txt}**\n"
+            f"📶 Ping: **{ping_line}**\n"
+            f"💾 Disk: **{total_gb:.1f} GB total**, **{free_gb:.1f} GB free**\n\n"
+            f"⬇️ Download: **{dl_mb_s:.2f} MB/s** "
+            f"({bytes_to_mb(dl['bytes']):.2f} MB in {dl['seconds']:.2f}s)\n"
+            f"⬆️ Upload: **{up_mb_s:.2f} MB/s** "
+            f"({bytes_to_mb(up['bytes']):.2f} MB in {up['seconds']:.2f}s)\n"
+        )
+        await safe_edit(msg, txt)
+
     # -------- OWNER allow/disallow via forward OR id --------
     @app.on_message(filters.command(["allow"]) & filters.private)
     async def allow_cmd(_: Client, m: Message):
@@ -128,7 +197,7 @@ def setup_handlers(app: Client) -> None:
             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"
+            f"DNS google -> {dns_check('https://www.google.com')}\n",
         )
 
     # -------- AUTH UI --------
@@ -184,7 +253,25 @@ def setup_handlers(app: Client) -> None:
         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"]))
+    @app.on_message(
+        filters.private
+        & filters.text
+        & ~filters.command(
+            [
+                "start",
+                "ping",
+                "me",
+                "auth",
+                "cancel",
+                "profiles",
+                "speedtest",
+                "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)
@@ -294,38 +381,29 @@ def setup_handlers(app: Client) -> None:
 
         file_path = ""
         file_size = 0
+        file_name = "video.bin"
 
         try:
-            # ---------------- Download with ETA ----------------
+            # ---------------- Download (Telegram) ----------------
             set_task(task_id, "downloading", "")
 
-            dl_se = SpeedETA()
-            dl_last_edit = 0.0
+            dl_t0 = time.time()
+            file_path, file_size, file_name = await download_to_temp(app_, m)
+            dl_sec = max(0.001, time.time() - dl_t0)
 
-            async def dl_prog(done: int, total: int):
-                nonlocal dl_last_edit
-                snap = dl_se.update(done, total)
-                if time.time() - dl_last_edit < 2.0 and done < total:
-                    return
-                dl_last_edit = time.time()
-                txt = (
-                    f"⬇️ Downloading…\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)
+            # fallback size from disk if TG didn't give it
+            if not file_size and file_path and os.path.exists(file_path):
+                try:
+                    file_size = os.path.getsize(file_path)
+                except Exception:
+                    file_size = 0
 
-            file_path, file_size, file_name, dl_sec = await download_to_temp(
-                app_, m,
-                progress=dl_prog,
-                progress_args=(),
-            )
-            dl_avg_bps = int((file_size or 0) / max(dl_sec, 0.001))
+            dl_avg_bps = int((file_size or 0) / dl_sec) if file_size else 0
 
             # parse title/desc after download
             title, desc = extract_title_description(m, file_name)
 
-            # ---------------- Upload with ETA ----------------
+            # ---------------- Upload (YouTube) ----------------
             set_task(task_id, "uploading", "")
             up_t0 = time.time()
 
@@ -355,7 +433,7 @@ def setup_handlers(app: Client) -> None:
             )
 
             up_sec = max(0.001, time.time() - up_t0)
-            up_avg_bps = int((file_size or 0) / up_sec)
+            up_avg_bps = int((file_size or 0) / up_sec) if file_size else 0
 
             await record_upload(uid, use_profile_id)
 
@@ -366,7 +444,7 @@ def setup_handlers(app: Client) -> None:
                 f"🎬 {title}\n"
                 f"🔗 {yt_url}\n\n"
                 f"⬇️ Download: {human_eta(int(dl_sec))} | avg {human_bytes(dl_avg_bps)}/s\n"
-                f"⬆️ Upload: {human_eta(int(up_sec))} | avg {human_bytes(up_avg_bps)}/s"
+                f"⬆️ Upload: {human_eta(int(up_sec))} | avg {human_bytes(up_avg_bps)}/s",
             )
 
         except Exception as e: