Blame for bot/speed.py

"""/speed: an Ookla Speedtest run from the bot's own network.

The official Speedtest CLI is used when it is installed.  If it is missing or fails, the ``speedtest-cli`` Python
package (also an Ookla client) is tried.  Both give the same figures; only the parsing differs.
"""
from __future__ import annotations

import asyncio
import html
import json
import logging
import os
import shutil
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from .utils import esc

log = logging.getLogger(__name__)

CLI_TIMEOUT = 120  # seconds for one run of either client
CANDIDATES = ("ookla-speedtest", "speedtest")  # the Dockerfile installs the official binary as ookla-speedtest


class SpeedError(RuntimeError):
    """The speed test could not be run; the message is safe to show to the user."""


@dataclass
class SpeedResult:
    download_mbps: float
    upload_mbps: float
    ping_ms: float
    jitter_ms: Optional[float] = None
    packet_loss: Optional[float] = None
    isp: str = ""
    server: str = ""
    location: str = ""
    url: str = ""  # the speedtest.net page of this result
    source: str = ""


def _num(value: Any) -> Optional[float]:
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def parse_cli(data: Dict[str, Any]) -> SpeedResult:
    """The JSON the official CLI prints (bandwidth is in bytes per second)."""
    download = _num((data.get("download") or {}).get("bandwidth"))
    upload = _num((data.get("upload") or {}).get("bandwidth"))
    if download is None or upload is None:
        raise SpeedError("the Speedtest CLI gave no download or upload figure")
    ping = data.get("ping") or {}
    server = data.get("server") or {}
    location = ", ".join(part for part in (server.get("location"), server.get("country")) if part)
    return SpeedResult(
        download_mbps=download * 8 / 1e6,
        upload_mbps=upload * 8 / 1e6,
        ping_ms=_num(ping.get("latency")) or 0.0,
        jitter_ms=_num(ping.get("jitter")),
        packet_loss=_num(data.get("packetLoss")),
        isp=str(data.get("isp") or ""),
        server=str(server.get("name") or server.get("host") or ""),
        location=location,
        url=str((data.get("result") or {}).get("url") or ""),
        source="Ookla Speedtest CLI",
    )


def parse_python(data: Dict[str, Any]) -> SpeedResult:
    """``Speedtest.results.dict()`` of the speedtest-cli package (speeds are in bits per second)."""
    download, upload = _num(data.get("download")), _num(data.get("upload"))
    if download is None or upload is None:
        raise SpeedError("speedtest-cli gave no download or upload figure")
    server = data.get("server") or {}
    client = data.get("client") or {}
    location = ", ".join(part for part in (server.get("name"), server.get("country")) if part)
    return SpeedResult(
        download_mbps=download / 1e6,
        upload_mbps=upload / 1e6,
        ping_ms=_num(data.get("ping")) or 0.0,
        isp=str(client.get("isp") or ""),
        server=str(server.get("sponsor") or ""),
        location=location,
        url=str(data.get("share") or ""),
        source="speedtest-cli (Ookla)",
    )


def _last_result(stdout: str) -> Optional[Dict[str, Any]]:
    """The ``{"type": "result", ...}`` line of the CLI's JSON output."""
    found = None
    for line in stdout.splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            obj = json.loads(line)
        except ValueError:
            continue
        if isinstance(obj, dict) and obj.get("type") == "result":
            found = obj
    return found


def _log_message(stdout: str, stderr: str) -> str:
    """Why the CLI failed, as far as it says."""
    for text in (stdout, stderr):
        for line in reversed(text.splitlines()):
            line = line.strip()
            if line.startswith("{"):
                try:
                    obj = json.loads(line)
                except ValueError:
                    continue
                if isinstance(obj, dict) and obj.get("message"):
                    return str(obj["message"])
            elif line:
                return line
    return "no output"


async def _run(cmd: List[str], timeout: float) -> "tuple[int, str, str]":
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
    except (FileNotFoundError, PermissionError, OSError) as exc:
        raise SpeedError(f"{cmd[0]} cannot be started ({exc})") from exc
    try:
        out, err = await asyncio.wait_for(proc.communicate(), timeout)
    except asyncio.TimeoutError:
        proc.kill()
        await proc.wait()
        raise SpeedError("the speed test took too long") from None
    return proc.returncode or 0, out.decode("utf-8", "replace"), err.decode("utf-8", "replace")


async def find_ookla_cli() -> Optional[str]:
    """The official CLI, if installed.  ``speedtest`` can also be the Python package's script, so ``--version`` decides."""
    names = [os.environ.get("SPEEDTEST_BIN", "")] + list(CANDIDATES)
    for name in dict.fromkeys(n for n in names if n):
        path = shutil.which(name)
        if not path:
            continue
        try:
            code, out, _err = await _run([path, "--version"], 15)
        except SpeedError:
            continue
        if code == 0 and "ookla" in out.lower():
            return path
    return None


async def run_cli(path: str, timeout: float = CLI_TIMEOUT) -> SpeedResult:
    code, out, err = await _run([path, "--accept-license", "--accept-gdpr", "--format=json"], timeout)
    result = _last_result(out)
    if result is None:
        raise SpeedError(f"the Speedtest CLI failed: {_log_message(out, err)}")
    return parse_cli(result)


def _run_python() -> Dict[str, Any]:
    """Blocking: one test with the speedtest-cli package."""
    try:
        import speedtest  # type: ignore
    except ImportError as exc:
        raise SpeedError("the speedtest-cli package is not installed") from exc
    try:
        test = speedtest.Speedtest(secure=True)
        test.get_best_server()
        test.download()
        test.upload()
        return dict(test.results.dict())
    except Exception as exc:  # noqa: BLE001 - the package raises its own family of errors
        raise SpeedError(f"speedtest-cli failed: {exc}") from exc


async def run_speedtest() -> SpeedResult:
    """Run a test with the official CLI, or with the Python client when the CLI is missing or fails."""
    problems: List[str] = []
    cli = await find_ookla_cli()
    if cli is not None:
        try:
            return await run_cli(cli)
        except SpeedError as exc:
            log.warning("Speedtest CLI failed: %s", exc)
            problems.append(str(exc))
    else:
        problems.append("the Speedtest CLI is not installed")
    try:
        data = await asyncio.wait_for(asyncio.to_thread(_run_python), CLI_TIMEOUT)
        return parse_python(data)
    except asyncio.TimeoutError:
        problems.append("speedtest-cli took too long")
    except SpeedError as exc:
        problems.append(str(exc))
    raise SpeedError("; ".join(problems))


def _mbps(value: float) -> str:
    return f"{value:,.2f} Mbps"


def format_result(r: SpeedResult) -> str:
    lines = [
        "<b>⚡ SPEEDTEST</b> · Ookla",
        "",
        f"◈ <b>Download</b> ▸ <code>{_mbps(r.download_mbps)}</code>",
        f"◈ <b>Upload</b> ▸ <code>{_mbps(r.upload_mbps)}</code>",
        f"◈ <b>Ping</b> ▸ <code>{r.ping_ms:.2f} ms</code>" + (f" · jitter <code>{r.jitter_ms:.2f} ms</code>" if r.jitter_ms is not None else ""),
    ]
    if r.packet_loss is not None:
        lines.append(f"◈ <b>Packet loss</b> ▸ <code>{r.packet_loss:g}%</code>")
    if r.isp:
        lines.append(f"◈ <b>ISP</b> ▸ {esc(r.isp)}")
    server = " · ".join(part for part in (esc(r.server), esc(r.location)) if part)
    if server:
        lines.append(f"◈ <b>Server</b> ▸ {server}")
    if r.url.startswith(("http://", "https://")):
        lines.append(f"◈ <b>Result</b> ▸ <a href=\"{html.escape(r.url, quote=True)}\">speedtest.net</a>")
    lines.append(f"\n<i>Measured from the machine the bot runs on · {esc(r.source)}</i>")
    return "\n".join(lines)