SpeedTray.pyw

"""SpeedTray: a tiny, dependency-free Windows taskbar network meter."""

from __future__ import annotations

import ctypes
from ctypes import wintypes
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import tkinter as tk
import winreg


APP_NAME = "SpeedTray"
DISPLAY_NAME = "SpeedTray"
WINDOW_WIDTH = 128
WINDOW_HEIGHT = 44
UPDATE_MS = 1000
DARK_COLORS = ("#1f1f1f", "#ffffff")
LIGHT_COLORS = ("#f3f3f3", "#111111")
IS_FROZEN = bool(getattr(sys, "frozen", False))

LOCAL_APP_DIR = Path(os.environ.get("LOCALAPPDATA", Path.home())) / APP_NAME
INSTALLED_SCRIPT = LOCAL_APP_DIR / ("SpeedTray.exe" if IS_FROZEN else "SpeedTray.pyw")
SETTINGS_FILE = LOCAL_APP_DIR / "settings.json"
RUN_KEY = r"Software\Microsoft\Windows\CurrentVersion\Run"


def pythonw_path() -> Path:
    candidate = Path(sys.executable).with_name("pythonw.exe")
    return candidate if candidate.exists() else Path(sys.executable)


def install_and_relaunch_if_needed() -> None:
    """Copy this file to a stable location before registering auto-start."""
    source = Path(sys.executable if IS_FROZEN else sys.argv[0]).resolve()
    try:
        already_installed = source.samefile(INSTALLED_SCRIPT)
    except (FileNotFoundError, OSError):
        already_installed = source == INSTALLED_SCRIPT.resolve()

    if already_installed:
        return

    LOCAL_APP_DIR.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, INSTALLED_SCRIPT)
    flags = getattr(subprocess, "DETACHED_PROCESS", 0) | getattr(
        subprocess, "CREATE_NEW_PROCESS_GROUP", 0
    )
    launch_command = (
        [str(INSTALLED_SCRIPT)]
        if IS_FROZEN
        else [str(pythonw_path()), str(INSTALLED_SCRIPT)]
    )
    subprocess.Popen(
        launch_command,
        close_fds=True,
        creationflags=flags,
    )
    raise SystemExit


def load_settings() -> dict:
    try:
        data = json.loads(SETTINGS_FILE.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else {}
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return {}


def save_settings(settings: dict) -> None:
    LOCAL_APP_DIR.mkdir(parents=True, exist_ok=True)
    temporary = SETTINGS_FILE.with_suffix(".tmp")
    temporary.write_text(json.dumps(settings, indent=2), encoding="utf-8")
    os.replace(temporary, SETTINGS_FILE)


def set_autostart(enabled: bool) -> None:
    with winreg.CreateKey(winreg.HKEY_CURRENT_USER, RUN_KEY) as key:
        if enabled:
            command = (
                f'"{INSTALLED_SCRIPT}"'
                if IS_FROZEN
                else f'"{pythonw_path()}" "{INSTALLED_SCRIPT}"'
            )
            winreg.SetValueEx(key, APP_NAME, 0, winreg.REG_SZ, command)
        else:
            try:
                winreg.DeleteValue(key, APP_NAME)
            except FileNotFoundError:
                pass


def taskbar_colors() -> tuple[str, str]:
    """Match the Windows system/taskbar theme and keep the text readable."""
    theme_key = r"Software\Microsoft\Windows\CurrentVersion\Themes\Personalize"
    try:
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, theme_key) as key:
            light_theme, _ = winreg.QueryValueEx(key, "SystemUsesLightTheme")
        return LIGHT_COLORS if light_theme else DARK_COLORS
    except (FileNotFoundError, OSError):
        return DARK_COLORS


def make_single_instance() -> object:
    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    create_mutex = kernel32.CreateMutexW
    create_mutex.argtypes = [ctypes.c_void_p, wintypes.BOOL, wintypes.LPCWSTR]
    create_mutex.restype = wintypes.HANDLE
    handle = create_mutex(None, False, "Local\\SpeedTray.Instance")
    if handle and ctypes.get_last_error() == 183:  # ERROR_ALREADY_EXISTS
        raise SystemExit
    return handle


# Structures used by the Windows IP Helper API (GetIfTable2).
IF_MAX_STRING_SIZE = 256
IF_MAX_PHYS_ADDRESS_LENGTH = 32


class GUID(ctypes.Structure):
    _fields_ = [
        ("Data1", wintypes.DWORD),
        ("Data2", wintypes.WORD),
        ("Data3", wintypes.WORD),
        ("Data4", ctypes.c_ubyte * 8),
    ]


class MIB_IF_ROW2(ctypes.Structure):
    _fields_ = [
        ("InterfaceLuid", ctypes.c_uint64),
        ("InterfaceIndex", wintypes.ULONG),
        ("InterfaceGuid", GUID),
        ("Alias", wintypes.WCHAR * (IF_MAX_STRING_SIZE + 1)),
        ("Description", wintypes.WCHAR * (IF_MAX_STRING_SIZE + 1)),
        ("PhysicalAddressLength", wintypes.ULONG),
        ("PhysicalAddress", ctypes.c_ubyte * IF_MAX_PHYS_ADDRESS_LENGTH),
        ("PermanentPhysicalAddress", ctypes.c_ubyte * IF_MAX_PHYS_ADDRESS_LENGTH),
        ("Mtu", wintypes.ULONG),
        ("Type", wintypes.ULONG),
        ("TunnelType", ctypes.c_int),
        ("MediaType", ctypes.c_int),
        ("PhysicalMediumType", ctypes.c_int),
        ("AccessType", ctypes.c_int),
        ("DirectionType", ctypes.c_int),
        ("InterfaceAndOperStatusFlags", ctypes.c_ubyte),
        ("OperStatus", ctypes.c_int),
        ("AdminStatus", ctypes.c_int),
        ("MediaConnectState", ctypes.c_int),
        ("NetworkGuid", GUID),
        ("ConnectionType", ctypes.c_int),
        ("TransmitLinkSpeed", ctypes.c_uint64),
        ("ReceiveLinkSpeed", ctypes.c_uint64),
        ("InOctets", ctypes.c_uint64),
        ("InUcastPkts", ctypes.c_uint64),
        ("InNUcastPkts", ctypes.c_uint64),
        ("InDiscards", ctypes.c_uint64),
        ("InErrors", ctypes.c_uint64),
        ("InUnknownProtos", ctypes.c_uint64),
        ("InUcastOctets", ctypes.c_uint64),
        ("InMulticastOctets", ctypes.c_uint64),
        ("InBroadcastOctets", ctypes.c_uint64),
        ("OutOctets", ctypes.c_uint64),
        ("OutUcastPkts", ctypes.c_uint64),
        ("OutNUcastPkts", ctypes.c_uint64),
        ("OutDiscards", ctypes.c_uint64),
        ("OutErrors", ctypes.c_uint64),
        ("OutUcastOctets", ctypes.c_uint64),
        ("OutMulticastOctets", ctypes.c_uint64),
        ("OutBroadcastOctets", ctypes.c_uint64),
        ("OutQLen", ctypes.c_uint64),
    ]


class MIB_IF_TABLE2(ctypes.Structure):
    _fields_ = [("NumEntries", wintypes.ULONG), ("Table", MIB_IF_ROW2 * 1)]


class NetworkCounters:
    def __init__(self) -> None:
        self.dll = ctypes.WinDLL("iphlpapi.dll")
        self.get_table = self.dll.GetIfTable2
        self.get_table.argtypes = [ctypes.POINTER(ctypes.POINTER(MIB_IF_TABLE2))]
        self.get_table.restype = wintypes.ULONG
        self.free_table = self.dll.FreeMibTable
        self.free_table.argtypes = [ctypes.c_void_p]

    def read(self) -> dict[int, tuple[int, int]]:
        pointer = ctypes.POINTER(MIB_IF_TABLE2)()
        result = self.get_table(ctypes.byref(pointer))
        if result:
            raise OSError(result, "GetIfTable2 failed")

        physical: dict[int, tuple[int, int]] = {}
        fallback: dict[int, tuple[int, int]] = {}
        try:
            count = pointer.contents.NumEntries
            rows = ctypes.cast(
                ctypes.addressof(pointer.contents.Table),
                ctypes.POINTER(MIB_IF_ROW2),
            )
            for index in range(count):
                row = rows[index]
                alias = row.Alias
                # Up, non-loopback interfaces only. The -0000 rows are Windows
                # filter layers whose byte counts duplicate their parent adapter.
                if row.OperStatus != 1 or row.Type == 24 or alias.endswith("-0000"):
                    continue
                values = (int(row.InOctets), int(row.OutOctets))
                fallback[int(row.InterfaceIndex)] = values
                if row.InterfaceAndOperStatusFlags & 1:  # HardwareInterface
                    physical[int(row.InterfaceIndex)] = values
            return physical or fallback
        finally:
            self.free_table(pointer)


class RECT(ctypes.Structure):
    _fields_ = [
        ("left", ctypes.c_long),
        ("top", ctypes.c_long),
        ("right", ctypes.c_long),
        ("bottom", ctypes.c_long),
    ]


def taskbar_rects() -> tuple[RECT | None, RECT | None]:
    """Return the taskbar and notification-area rectangles from Explorer."""
    user32 = ctypes.WinDLL("user32")
    find_window = user32.FindWindowW
    find_window.argtypes = [wintypes.LPCWSTR, wintypes.LPCWSTR]
    find_window.restype = wintypes.HWND
    find_child = user32.FindWindowExW
    find_child.argtypes = [wintypes.HWND, wintypes.HWND, wintypes.LPCWSTR, wintypes.LPCWSTR]
    find_child.restype = wintypes.HWND
    get_rect = user32.GetWindowRect
    get_rect.argtypes = [wintypes.HWND, ctypes.POINTER(RECT)]
    get_rect.restype = wintypes.BOOL

    taskbar_hwnd = find_window("Shell_TrayWnd", None)
    bar = RECT()
    if not taskbar_hwnd or not get_rect(taskbar_hwnd, ctypes.byref(bar)):
        return None, None

    tray_hwnd = find_child(taskbar_hwnd, None, "TrayNotifyWnd", None)
    tray = RECT()
    if not tray_hwnd or not get_rect(tray_hwnd, ctypes.byref(tray)):
        return bar, None
    return bar, tray


def default_position(root: tk.Tk) -> tuple[int, int]:
    bar, tray = taskbar_rects()
    if bar:
        bar_width = bar.right - bar.left
        bar_height = bar.bottom - bar.top
        if bar_width >= bar_height:  # Normal horizontal taskbar.
            if tray and tray.left > bar.left + bar_width // 2:
                x = tray.left - WINDOW_WIDTH - 8
            elif tray:
                x = tray.right + 8
            else:
                x = bar.right - WINDOW_WIDTH - 245
            x = max(bar.left, min(x, bar.right - WINDOW_WIDTH))
            y = bar.top + max(0, (bar_height - WINDOW_HEIGHT) // 2)
        else:  # Windows 10 can also use a vertical taskbar.
            x = bar.left + max(0, (bar_width - WINDOW_WIDTH) // 2)
            if tray and tray.top > bar.top + bar_height // 2:
                y = tray.top - WINDOW_HEIGHT - 8
            elif tray:
                y = tray.bottom + 8
            else:
                y = bar.bottom - WINDOW_HEIGHT - 170
            y = max(bar.top, min(y, bar.bottom - WINDOW_HEIGHT))
        return x, y
    return root.winfo_screenwidth() - WINDOW_WIDTH - 245, root.winfo_screenheight() - WINDOW_HEIGHT


def format_rate(bytes_per_second: float) -> str:
    bits_per_second = max(0.0, bytes_per_second * 8.0)
    if bits_per_second >= 1_000_000:
        return f"{bits_per_second / 1_000_000:.1f} Mb/s"
    return f"{bits_per_second / 1_000:.1f} Kb/s"


class SpeedMeter:
    def __init__(self, settings: dict) -> None:
        self.settings = settings
        self.counters = NetworkCounters()
        self.previous = self.counters.read()
        self.previous_time = time.perf_counter()
        self.drag_offset = (0, 0)
        self.dragging = False
        self.auto_position = not bool(settings.get("manual_position", False))

        self.bg, self.fg = taskbar_colors()

        self.root = tk.Tk(className=APP_NAME)
        self.root.title(DISPLAY_NAME)
        self.root.configure(bg=self.bg)
        self.root.overrideredirect(True)
        self.root.attributes("-topmost", True)

        saved_x = settings.get("x")
        saved_y = settings.get("y")
        if not self.auto_position and isinstance(saved_x, int) and isinstance(saved_y, int):
            x, y = saved_x, saved_y
        else:
            x, y = default_position(self.root)
        self.root.geometry(f"{WINDOW_WIDTH}x{WINDOW_HEIGHT}+{x}+{y}")

        self.up_var = tk.StringVar(value="↑: 0.0 Kb/s")
        self.down_var = tk.StringVar(value="↓: 0.0 Kb/s")
        common = {
            "bg": self.bg,
            "fg": self.fg,
            "font": ("Segoe UI", 9),
            "anchor": "e",
            "padx": 7,
            "borderwidth": 0,
        }
        self.up_label = tk.Label(self.root, textvariable=self.up_var, **common)
        self.down_label = tk.Label(self.root, textvariable=self.down_var, **common)
        self.up_label.pack(fill="both", expand=True)
        self.down_label.pack(fill="both", expand=True)

        self.autostart_var = tk.BooleanVar(value=bool(settings.get("autostart", True)))
        self.auto_position_var = tk.BooleanVar(value=self.auto_position)
        self.menu = tk.Menu(self.root, tearoff=False)
        self.menu.add_checkbutton(
            label="Start with Windows",
            variable=self.autostart_var,
            command=self.toggle_autostart,
        )
        self.menu.add_checkbutton(
            label="Auto-position on taskbar",
            variable=self.auto_position_var,
            command=self.toggle_auto_position,
        )
        self.menu.add_command(label="Reset position", command=self.reset_position)
        self.menu.add_separator()
        self.menu.add_command(label="Exit", command=self.root.destroy)

        for widget in (self.root, self.up_label, self.down_label):
            widget.bind("<ButtonPress-1>", self.start_drag)
            widget.bind("<B1-Motion>", self.drag)
            widget.bind("<ButtonRelease-1>", self.end_drag)
            widget.bind("<Button-3>", self.show_menu)

        self.root.update_idletasks()
        self.hide_from_alt_tab()
        self.root.after(UPDATE_MS, self.update_speed)
        self.root.after(100, self.keep_above_taskbar)
        self.root.after(2000, self.refresh_theme)
        self.root.after(1000, self.follow_taskbar)

    def native_window_handle(self) -> int:
        """Return Tk's outer Windows handle, not its inner drawing window."""
        inner = self.root.winfo_id()
        user32 = ctypes.WinDLL("user32")
        get_parent = user32.GetParent
        get_parent.argtypes = [wintypes.HWND]
        get_parent.restype = wintypes.HWND
        outer = get_parent(inner)
        return outer or inner

    def hide_from_alt_tab(self) -> None:
        hwnd = self.native_window_handle()
        user32 = ctypes.windll.user32
        get_style = user32.GetWindowLongW
        set_style = user32.SetWindowLongW
        GWL_EXSTYLE = -20
        WS_EX_TOOLWINDOW = 0x00000080
        WS_EX_APPWINDOW = 0x00040000
        style = get_style(hwnd, GWL_EXSTYLE)
        set_style(hwnd, GWL_EXSTYLE, (style | WS_EX_TOOLWINDOW) & ~WS_EX_APPWINDOW)

    def keep_above_taskbar(self) -> None:
        try:
            HWND_TOPMOST = -1
            SWP_NOMOVE = 0x0002
            SWP_NOSIZE = 0x0001
            SWP_NOACTIVATE = 0x0010
            user32 = ctypes.WinDLL("user32", use_last_error=True)
            set_window_pos = user32.SetWindowPos
            set_window_pos.argtypes = [
                wintypes.HWND,
                wintypes.HWND,
                ctypes.c_int,
                ctypes.c_int,
                ctypes.c_int,
                ctypes.c_int,
                wintypes.UINT,
            ]
            set_window_pos.restype = wintypes.BOOL
            set_window_pos(
                self.native_window_handle(),
                wintypes.HWND(HWND_TOPMOST),
                0,
                0,
                0,
                0,
                SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE,
            )
            self.root.after(2000, self.keep_above_taskbar)
        except tk.TclError:
            pass

    def refresh_theme(self) -> None:
        try:
            bg, fg = taskbar_colors()
            if (bg, fg) != (self.bg, self.fg):
                self.bg, self.fg = bg, fg
                self.root.configure(bg=bg)
                self.up_label.configure(bg=bg, fg=fg)
                self.down_label.configure(bg=bg, fg=fg)
            self.root.after(2000, self.refresh_theme)
        except tk.TclError:
            pass

    def follow_taskbar(self) -> None:
        try:
            if self.auto_position and not self.dragging:
                x, y = default_position(self.root)
                if (x, y) != (self.root.winfo_x(), self.root.winfo_y()):
                    self.root.geometry(f"+{x}+{y}")
            self.root.after(2000, self.follow_taskbar)
        except tk.TclError:
            pass

    def update_speed(self) -> None:
        try:
            now = time.perf_counter()
            current = self.counters.read()
            elapsed = max(now - self.previous_time, 0.001)
            downloaded = 0
            uploaded = 0
            for interface, (received, sent) in current.items():
                old = self.previous.get(interface)
                if old:
                    downloaded += max(0, received - old[0])
                    uploaded += max(0, sent - old[1])
            self.up_var.set(f"↑: {format_rate(uploaded / elapsed)}")
            self.down_var.set(f"↓: {format_rate(downloaded / elapsed)}")
            self.previous = current
            self.previous_time = now
        except (OSError, tk.TclError):
            self.up_var.set("↑: --")
            self.down_var.set("↓: --")
        finally:
            try:
                self.root.after(UPDATE_MS, self.update_speed)
            except tk.TclError:
                pass

    def start_drag(self, event: tk.Event) -> None:
        self.dragging = True
        self.drag_offset = (event.x_root - self.root.winfo_x(), event.y_root - self.root.winfo_y())

    def drag(self, event: tk.Event) -> None:
        x = event.x_root - self.drag_offset[0]
        y = event.y_root - self.drag_offset[1]
        self.root.geometry(f"+{x}+{y}")

    def end_drag(self, _event: tk.Event) -> None:
        self.dragging = False
        self.auto_position = False
        self.auto_position_var.set(False)
        self.settings["x"] = self.root.winfo_x()
        self.settings["y"] = self.root.winfo_y()
        self.settings["manual_position"] = True
        save_settings(self.settings)

    def show_menu(self, event: tk.Event) -> None:
        self.menu.tk_popup(event.x_root, event.y_root)

    def toggle_autostart(self) -> None:
        enabled = self.autostart_var.get()
        set_autostart(enabled)
        self.settings["autostart"] = enabled
        save_settings(self.settings)

    def toggle_auto_position(self) -> None:
        self.auto_position = self.auto_position_var.get()
        self.settings["manual_position"] = not self.auto_position
        if self.auto_position:
            x, y = default_position(self.root)
            self.root.geometry(f"+{x}+{y}")
        save_settings(self.settings)

    def reset_position(self) -> None:
        self.auto_position = True
        self.auto_position_var.set(True)
        x, y = default_position(self.root)
        self.root.geometry(f"+{x}+{y}")
        self.settings.update({"x": x, "y": y, "manual_position": False})
        save_settings(self.settings)

    def run(self) -> None:
        self.root.mainloop()


def enable_dpi_awareness() -> None:
    try:
        ctypes.windll.shcore.SetProcessDpiAwareness(2)
    except (AttributeError, OSError):
        try:
            ctypes.windll.user32.SetProcessDPIAware()
        except (AttributeError, OSError):
            pass


def main() -> None:
    if sys.platform != "win32":
        raise SystemExit("SpeedTray requires Windows.")

    if "--test-counters" in sys.argv:
        print(NetworkCounters().read())
        return

    install_and_relaunch_if_needed()
    enable_dpi_awareness()
    _mutex = make_single_instance()

    settings = load_settings()
    if "autostart" not in settings:
        settings["autostart"] = True
        save_settings(settings)
    set_autostart(bool(settings["autostart"]))
    SpeedMeter(settings).run()
    _ = _mutex  # Keep the Windows mutex handle alive until the app exits.


if __name__ == "__main__":
    main()