#!/usr/bin/python3
"""Restart running Moodtop shells whose code is newer than the process (run as root by moodtop-core's postinst).

The new shell is started as a transient *user* service with the old one's Wayland/Wayfire/D-Bus environment,
so it lives in the user's session instead of dying with the updater's systemd unit."""
import os
import pwd
import signal
import subprocess
import time

SHELL = "/usr/lib/mood/shell/mood-shell.py"
KEEP = ("WAYLAND_DISPLAY", "WAYFIRE_SOCKET", "XDG_RUNTIME_DIR", "DBUS_SESSION_BUS_ADDRESS", "HOME", "USER", "LANG", "LANGUAGE",
        "PATH", "DISPLAY", "XDG_SESSION_TYPE", "XDG_CURRENT_DESKTOP", "XDG_SESSION_DESKTOP", "XDG_DATA_DIRS", "XDG_CONFIG_DIRS",
        "GIO_USE_PROXY_RESOLVER", "GTK_THEME", "XCURSOR_THEME", "XCURSOR_SIZE", "MOZ_ENABLE_WAYLAND", "QT_QPA_PLATFORM")


def shells():
    # ctime, not mtime: dpkg keeps the (reproducible, 2022) mtimes from the .deb, but ctime is when it was installed
    newest = max(os.stat(os.path.join(os.path.dirname(SHELL), f)).st_ctime for f in os.listdir(os.path.dirname(SHELL)))
    try:
        boot = time.time() - float(open("/proc/uptime").read().split()[0])
    except OSError:  # no /proc (chroot / image build): nothing is running anyway
        return
    tick = os.sysconf("SC_CLK_TCK")
    for pid in filter(str.isdigit, os.listdir("/proc")):
        try:
            argv = open(f"/proc/{pid}/cmdline", "rb").read().split(b"\0")
            if not any(a.decode(errors="replace") == SHELL for a in argv[:3]):
                continue
            started = boot + int(open(f"/proc/{pid}/stat").read().rsplit(")", 1)[1].split()[19]) / tick
            if started > newest:
                continue
            uid = os.stat(f"/proc/{pid}").st_uid
            try:
                raw = open(f"/proc/{pid}/environ", "rb").read()
            except PermissionError:  # hardened kernels/sandboxes: the owner can always read its own
                raw = subprocess.run(["runuser", "-u", pwd.getpwuid(uid).pw_name, "--", "cat", f"/proc/{pid}/environ"],
                                     capture_output=True, timeout=5).stdout
            env = dict(kv.split("=", 1) for kv in raw.decode(errors="replace").split("\0") if "=" in kv)
            yield int(pid), uid, env
        except (OSError, ValueError, IndexError, KeyError, subprocess.SubprocessError):
            continue


def main():
    for pid, uid, env in list(shells()):
        if uid == 0 or not env.get("WAYLAND_DISPLAY"):
            continue
        user = pwd.getpwuid(uid).pw_name
        run = env.get("XDG_RUNTIME_DIR") or f"/run/user/{uid}"
        try:
            os.kill(pid, signal.SIGTERM)
            for _ in range(30):
                time.sleep(0.1)
                os.kill(pid, 0)
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        sets = [f"--setenv={k}={env[k]}" for k in KEEP if k in env]
        subprocess.run(["runuser", "-u", user, "--", "env", f"XDG_RUNTIME_DIR={run}", f"DBUS_SESSION_BUS_ADDRESS={env.get('DBUS_SESSION_BUS_ADDRESS', f'unix:path={run}/bus')}",
                        "systemd-run", "--user", "--collect", "--quiet", f"--unit=mood-shell-{int(time.time())}-{pid}", *sets,
                        "/usr/bin/python3", SHELL], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=20, check=False)
        print(f"mood-shell-restart: restarted the shell for {user}")


if __name__ == "__main__":
    main()
