tg-ws-proxy/linux.py

from __future__ import annotations

import os
import subprocess
import sys
import threading
import time
import webbrowser
from typing import Optional

import customtkinter as ctk
import pyperclip
import pystray
from PIL import Image, ImageTk

from proxy import __version__, get_link_host

from utils.tray_common import (
    APP_NAME, DEFAULT_CONFIG, FIRST_RUN_MARKER, IS_FROZEN, LOG_FILE,
    acquire_lock, bootstrap, check_ipv6_warning, check_update_async,
    ctk_run_dialog, ensure_ctk_thread, ensure_dirs, load_config, load_icon,
    log, quit_ctk, release_lock, restart_proxy,
    save_config, start_proxy, stop_proxy, tg_proxy_url,
)
from ui.ctk_tray_ui import (
    install_tray_config_buttons, install_tray_config_form,
    populate_first_run_window, tray_settings_scroll_and_footer,
)
from ui.ctk_theme import (
    CONFIG_DIALOG_FRAME_PAD, CONFIG_DIALOG_SIZE, FIRST_RUN_SIZE,
    create_ctk_toplevel, ctk_theme_for_platform, main_content_frame,
)
from ui.i18n import t
from ui.settings_dialog import SettingsDialog
from utils.update_check import get_status, relaunch_env

_tray_icon: Optional[object] = None
_config: dict = {}
_exiting = False
_update_flow_lock = threading.Lock()

# dialogs (tkinter messagebox)


def _msgbox(kind: str, text: str, title: str, **kw):
    import tkinter as _tk
    from tkinter import messagebox as _mb

    root = _tk.Tk()
    root.withdraw()
    try:
        root.attributes("-topmost", True)
    except Exception:
        pass
    result = getattr(_mb, kind)(title, text, parent=root, **kw)
    root.destroy()
    return result


def _show_error(text: str, title: Optional[str] = None) -> None:
    _msgbox("showerror", text, title or t("app.error_title"))


def _show_info(text: str, title: Optional[str] = None) -> None:
    _msgbox("showinfo", text, title or t("app.name"))


def _ask_yes_no(text: str, title: Optional[str] = None) -> bool:
    return bool(_msgbox("askyesno", text, title or t("app.name")))


def _apply_window_icon(root) -> None:
    icon_img = load_icon()
    if icon_img:
        root._ctk_icon_photo = ImageTk.PhotoImage(icon_img.resize((64, 64)))
        root.iconphoto(False, root._ctk_icon_photo)


# tray callbacks


def _on_open_in_telegram(icon=None, item=None) -> None:
    url = tg_proxy_url(_config)
    log.info("Copying %s", url)
    try:
        pyperclip.copy(url)
        _show_info(t("dialog.copy_ok", url=url))
    except Exception as exc:
        log.error("Clipboard copy failed: %s", exc)
        _show_error(t("dialog.copy_fail", error=exc))


def _on_copy_link(icon=None, item=None) -> None:
    url = tg_proxy_url(_config)
    log.info("Copying link: %s", url)
    try:
        pyperclip.copy(url)
    except Exception as exc:
        log.error("Clipboard copy failed: %s", exc)
        _show_error(t("dialog.copy_fail", error=exc))


def _on_restart(icon=None, item=None) -> None:
    threading.Thread(
        target=lambda: restart_proxy(_config, _show_error), daemon=True
    ).start()


def _on_edit_config(icon=None, item=None) -> None:
    threading.Thread(target=_edit_config_dialog, daemon=True).start()


def _on_open_logs(icon=None, item=None) -> None:
    log.info("Opening log file: %s", LOG_FILE)
    if LOG_FILE.exists():
        env = {k: v for k, v in os.environ.items() if k not in ("VIRTUAL_ENV", "PYTHONPATH", "PYTHONHOME")}
        subprocess.Popen(
            ["xdg-open", str(LOG_FILE)], env=env,
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL, start_new_session=True,
        )
    else:
        _show_info(t("dialog.log_not_found"))


def _on_exit(icon=None, item=None) -> None:
    global _exiting
    if _exiting:
        os._exit(0)
        return
    _exiting = True
    log.info("User requested exit")
    quit_ctk()
    threading.Thread(target=lambda: (time.sleep(3), os._exit(0)), daemon=True, name="force-exit").start()
    if icon:
        icon.stop()


# settings dialog


def _edit_config_dialog() -> None:
    if not ensure_ctk_thread(ctk, _config.get("appearance", "auto")):
        _show_error(t("dialog.ctk_missing"))
        return

    cfg = dict(_config)

    def _build(done: threading.Event) -> None:
        theme = ctk_theme_for_platform()
        w, h = CONFIG_DIALOG_SIZE
        root = create_ctk_toplevel(
            ctk, title=t("app.settings_title"), width=w, height=h, 
            theme=theme, topmost=False, after_create=_apply_window_icon
        )
        fpx, fpy = CONFIG_DIALOG_FRAME_PAD
        frame = main_content_frame(ctk, root, theme, padx=fpx, pady=fpy)
        scroll, footer = tray_settings_scroll_and_footer(ctk, frame, theme)

        def _refresh_tray_menu() -> None:
            if _tray_icon is not None:
                _tray_icon.menu = _build_menu()

        widgets = install_tray_config_form(
            ctk, scroll, theme, cfg, DEFAULT_CONFIG,
            show_autostart=False,
            on_update_click=_trigger_update_flow,
        )

        def _finish() -> None:
            root.destroy()
            done.set()

        dialog = SettingsDialog(
            ctk=ctk, root=root, widgets=widgets, config=_config,
            defaults=DEFAULT_CONFIG, persist=save_config,
            refresh_menu=_refresh_tray_menu, finish=_finish,
            restart=lambda config: restart_proxy(config, _show_error),
        )

        root.protocol("WM_DELETE_WINDOW", dialog.cancel)
        install_tray_config_buttons(ctk, footer, theme, on_save=dialog.save, on_cancel=dialog.cancel)

    ctk_run_dialog(_build)


# first run


def _show_first_run() -> None:
    ensure_dirs()
    if FIRST_RUN_MARKER.exists():
        return
    if not ensure_ctk_thread(ctk, _config.get("appearance", "auto")):
        FIRST_RUN_MARKER.touch()
        return

    host = _config.get("host", DEFAULT_CONFIG["host"])
    port = _config.get("port", DEFAULT_CONFIG["port"])
    secret = _config.get("secret", DEFAULT_CONFIG["secret"])

    def _build(done: threading.Event) -> None:
        theme = ctk_theme_for_platform()
        w, h = FIRST_RUN_SIZE
        root = create_ctk_toplevel(
            ctk, title=t("app.name"), width=w, height=h, theme=theme,
            topmost=False, after_create=_apply_window_icon,
        )

        def on_done(open_tg: bool) -> None:
            FIRST_RUN_MARKER.touch()
            root.destroy()
            done.set()
            if open_tg:
                _on_open_in_telegram()

        populate_first_run_window(ctk, root, theme, host=host, port=port, secret=secret, on_done=on_done)

    ctk_run_dialog(_build)


# update


def _restart_after_update() -> None:
    from utils.linux_update import exe_path

    stop_proxy()
    release_lock()
    exe = str(exe_path())
    try:
        subprocess.Popen(
            [exe, *sys.argv[1:]], env=relaunch_env(),
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL, start_new_session=True,
        )
    except OSError as exc:
        log.error("Failed to launch updated binary: %s", exc)
        if acquire_lock():
            start_proxy(_config, _show_error)
        _show_error(t("update.error", msg=exc))
        return
    os._exit(0)


def _trigger_update_flow(icon=None, item=None) -> None:
    if not _update_flow_lock.acquire(blocking=False):
        return

    def _work() -> None:
        from utils.linux_update import asset_name, detect_install_kind
        from utils.update_check import find_asset, RELEASES_PAGE_URL

        try:
            st = get_status()
            if _exiting or not st["has_update"]:
                return
            if _tray_icon is not None:
                _tray_icon.menu = _build_menu()
            kind = detect_install_kind() if IS_FROZEN else None
            asset = find_asset(asset_name(kind)) if kind else None
            page_url = st["html_url"] or RELEASES_PAGE_URL
            if not asset:
                if _ask_yes_no(t("update.ask_open", version=st["latest"]), t("app.update_title")):
                    webbrowser.open(page_url)
                return

            if _ask_yes_no(t("update.ask_install", version=st["latest"]), t("app.update_title")):
                _install_update(kind, asset, page_url)
        except Exception:
            log.exception("Update flow failed")
        finally:
            _update_flow_lock.release()

    threading.Thread(target=_work, daemon=True, name="manual-update").start()


def _install_update(kind: str, asset: dict, page_url: str) -> None:
    from utils.linux_update import apply_update

    err = apply_update(
        kind, asset["url"], digest=asset.get("digest", ""), log=log,
    )
    if err is None:
        log.info("Update installed, restarting")
        _restart_after_update()
        return

    log.error("Update failed: %s", err)
    if _ask_yes_no(t("update.install_fail", error=err), t("app.update_title")):
        webbrowser.open(page_url)


# tray menu


def _build_menu():
    host = _config.get("host", DEFAULT_CONFIG["host"])
    port = _config.get("port", DEFAULT_CONFIG["port"])
    link_host = get_link_host(host)
    items = [
        pystray.MenuItem(t("tray.open_telegram", host=link_host, port=port), _on_open_in_telegram, default=True),
        pystray.MenuItem(t("tray.copy_link"), _on_copy_link),
        pystray.Menu.SEPARATOR,
        pystray.MenuItem(t("tray.restart"), _on_restart),
        pystray.MenuItem(t("tray.settings"), _on_edit_config),
        pystray.MenuItem(t("tray.logs"), _on_open_logs),
        pystray.Menu.SEPARATOR,
        pystray.MenuItem(t("tray.exit"), _on_exit),
    ]
    st = get_status()
    if st["has_update"]:
        items.insert(-2, pystray.MenuItem(
            t("tray.update", current=__version__, new=st["latest"]), _trigger_update_flow,
        ))
    return pystray.Menu(*items)


# entry point


def run_tray() -> None:
    global _tray_icon, _config

    _config = load_config()
    bootstrap(_config)

    if pystray is None or Image is None:
        log.error("pystray or Pillow not installed; running in console mode")
        start_proxy(_config, _show_error)
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            stop_proxy()
        return

    start_proxy(_config, _show_error)
    check_update_async(_config, lambda: _exiting, lambda version, url: _trigger_update_flow())
    _show_first_run()
    check_ipv6_warning(_show_info)

    _tray_icon = pystray.Icon(APP_NAME, load_icon(), t("app.name"), menu=_build_menu())
    log.info("Tray icon running")
    _tray_icon.run()

    stop_proxy()
    log.info("Tray app exited")


def _smoke_test() -> None:
    """Exercise the frozen GUI stack without starting the proxy or writing config."""
    import gi

    gi.require_version('AppIndicator3', '0.1')
    gi.require_version('AyatanaAppIndicator3', '0.1')
    from gi.repository import AppIndicator3, AyatanaAppIndicator3

    # Namespace imports alone do not resolve app_indicator_new: call it too.
    indicators = [
        namespace.Indicator.new(
            'tg-ws-proxy-smoke-test', '', namespace.IndicatorCategory.APPLICATION_STATUS,
        )
        for namespace in (AppIndicator3, AyatanaAppIndicator3)
    ]
    icon = pystray.Icon('tg-ws-proxy-smoke-test', Image.new('RGB', (16, 16)))
    root = ctk.CTk()
    root.withdraw()
    root.update_idletasks()
    root.destroy()

    # A successful constructor on the build host can hide missing bundled
    # libraries. Check where the dynamic loader actually obtained them.
    bundle_dir = os.path.realpath(sys._MEIPASS) + os.sep
    prefixes = ('libglib-', 'libgobject-', 'libgio-', 'libgtk-',
                'libappindicator', 'libayatana-', 'libdbusmenu-')
    with open('/proc/self/maps', encoding='utf-8') as maps:
        paths = {line.split(maxsplit=5)[-1].strip() for line in maps}
    for path in sorted(paths):
        if os.path.basename(path).startswith(prefixes):
            if not os.path.realpath(path).startswith(bundle_dir):
                raise RuntimeError('GUI library loaded outside the bundle: ' + path)
            print(path)
    assert icon is not None and all(indicators)
    print('Linux bundle smoke test passed')


def main() -> None:
    if sys.argv[1:] == ['--smoke-test']:
        _smoke_test()
        return
    if not acquire_lock():
        _show_info(t("dialog.already_running"), os.path.basename(sys.argv[0]))
        return
    try:
        run_tray()
    finally:
        release_lock()


if __name__ == "__main__":
    main()