xidle-inhibit/xautolock-inhibit

#!/usr/bin/env python3
"""Own org.freedesktop.ScreenSaver and turn inhibit requests into xautolock -disable/-enable.

While an inhibit is active, X screen blanking is also paused (like mpv does).
"""

import argparse
import ctypes
import json
import os
import signal
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import timedelta
from typing import Any

from gi.repository import Gio, GLib

try:
    from gi.repository import GLibUnix  # type: ignore[attr-defined]  # GLib >= 2.80, not in pygobject-stubs
    signal_add = GLibUnix.signal_add
except ImportError:
    signal_add = GLib.unix_signal_add

BUS_NAME = "org.freedesktop.ScreenSaver"
OBJECT_PATHS = ("/org/freedesktop/ScreenSaver", "/ScreenSaver")  # Vivaldi uses the first, Zen the second
EXIT_NAME_TAKEN = 2
ACTIVITY_QUIET = 300  # seconds without SimulateUserActivity that end a logged burst
STATE_FILE = "xautolock-inhibit.json"  # in $XDG_RUNTIME_DIR
STATE_HEARTBEAT = 10  # seconds between state file touches while inhibits are held

Cookie = tuple[str, str, str, bool]  # sender, application_name, reason, counts

INTROSPECTION = """
<node>
  <interface name="org.freedesktop.ScreenSaver">
    <method name="Inhibit">
      <arg name="application_name" type="s" direction="in"/>
      <arg name="reason_for_inhibit" type="s" direction="in"/>
      <arg name="cookie" type="u" direction="out"/>
    </method>
    <method name="UnInhibit">
      <arg name="cookie" type="u" direction="in"/>
    </method>
    <method name="SimulateUserActivity"/>
    <method name="Lock"/>
    <method name="GetActive">
      <arg type="b" direction="out"/>
    </method>
    <method name="SetActive">
      <arg type="b" direction="in"/>
      <arg type="b" direction="out"/>
    </method>
    <method name="GetActiveTime">
      <arg type="u" direction="out"/>
    </method>
    <method name="GetSessionIdleTime">
      <arg type="u" direction="out"/>
    </method>
  </interface>
</node>
"""


def log(msg: str) -> None:
    print(msg, flush=True)


def xautolock(arg: str) -> None:
    try:
        subprocess.run(["xautolock", arg], check=True, timeout=10)
        log(f"xautolock {arg}")
    except (OSError, subprocess.SubprocessError) as e:
        log(f"xautolock {arg} failed: {e}")


class XBlanking:
    """Pause X screen blanking with XScreenSaverSuspend (MIT-SCREEN-SAVER extension).

    Unlike `xset s off -dpms` this changes no global setting: the X server ends the
    pause by itself when this client releases it or its connection closes (exit, crash).
    """

    def __init__(self) -> None:
        self.dpy: int | None = None
        try:
            self.x11 = ctypes.CDLL("libX11.so.6")
            self.xss = ctypes.CDLL("libXss.so.1")
        except OSError as e:
            log(f"Cannot load X11 libraries ({e}); X blanking will not be paused")
            return
        self.x11.XOpenDisplay.argtypes = [ctypes.c_char_p]
        self.x11.XOpenDisplay.restype = ctypes.c_void_p
        self.x11.XFlush.argtypes = [ctypes.c_void_p]
        self.xss.XScreenSaverQueryExtension.argtypes = [
            ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int)]
        self.xss.XScreenSaverSuspend.argtypes = [ctypes.c_void_p, ctypes.c_int]

        dpy = self.x11.XOpenDisplay(None)
        if not dpy:
            log("Cannot open the X display; X blanking will not be paused")
            return
        event_base, error_base = ctypes.c_int(), ctypes.c_int()
        if not self.xss.XScreenSaverQueryExtension(dpy, ctypes.byref(event_base), ctypes.byref(error_base)):
            log("X server lacks the MIT-SCREEN-SAVER extension; X blanking will not be paused")
            return
        self.dpy = dpy

    def suspend(self, on: bool) -> None:
        if self.dpy:
            self.xss.XScreenSaverSuspend(self.dpy, int(on))
            self.x11.XFlush(self.dpy)
            log(f"X blanking {'paused' if on else 'resumed'}")


@dataclass
class ActivityBurst:
    first: float
    last: float
    calls: int = 0
    timer: int = 0


class Inhibitor:
    def __init__(
        self, ignored_reasons: set[str], release_delay: int, blanking: XBlanking,
        state_path: str | None, restore_window: int,
    ) -> None:
        self.ignored_reasons = ignored_reasons
        self.blanking = blanking
        self.release_delay = release_delay
        self.state_path = state_path
        self.restore_window = restore_window
        self.state_ready = False  # set once this instance owns the name and may write the state file
        self.heartbeat_timer = 0
        self.cookies: dict[int, Cookie] = {}
        self.last_cookie = 0
        self.proc_names: dict[str, str] = {}  # sender -> "comm[pid]", for logging only
        self.activity: dict[str, ActivityBurst] = {}  # sender -> SimulateUserActivity burst
        self.disabled = False
        self.release_timer = 0
        self.conn: Gio.DBusConnection | None = None

    # --- lock state -------------------------------------------------------

    def active_count(self) -> int:
        return sum(1 for *_, counts in self.cookies.values() if counts)

    def update(self) -> None:
        """Run xautolock only on 0->1 and 1->0 transitions; delay re-enabling."""
        active = self.active_count()
        if active and self.release_timer:
            GLib.source_remove(self.release_timer)
            self.release_timer = 0
        if active and not self.disabled:
            self.block()
        elif not active and self.disabled and not self.release_timer:
            # Browsers send quick Inhibit/UnInhibit bursts when a video starts.
            self.release_timer = GLib.timeout_add_seconds(self.release_delay, self.on_release_timeout)

    def on_release_timeout(self) -> bool:
        self.release_timer = 0
        if not self.active_count() and self.disabled:
            self.unblock()
        return GLib.SOURCE_REMOVE

    def block(self) -> None:
        xautolock("-disable")
        self.blanking.suspend(True)
        self.disabled = True

    def unblock(self) -> None:
        xautolock("-enable")
        self.blanking.suspend(False)
        self.disabled = False

    def shutdown(self) -> None:
        for sender in list(self.activity):
            self.end_activity(sender)
        if self.release_timer:
            GLib.source_remove(self.release_timer)
            self.release_timer = 0
        self.save(stopped=True)
        if self.disabled:
            self.unblock()

    # --- state file -------------------------------------------------------

    def start(self) -> None:
        """Take over inhibits saved by a previous instance, then set xautolock to match."""
        self.restore()
        self.state_ready = True
        self.save()
        if self.active_count():
            self.block()
        else:
            # Start from a known state, in case a previous instance died while disabled.
            xautolock("-enable")

    def restore(self) -> None:
        if not self.state_path:
            log("XDG_RUNTIME_DIR is not set; inhibits will not survive a restart")
            return
        assert self.conn is not None
        try:
            with open(self.state_path) as f:
                state = json.load(f)
            age = time.time() - (state["stopped_at"] or os.stat(self.state_path).st_mtime)
            bus_id = state["bus_id"]
            last_cookie = int(state["last_cookie"])
            saved = [(int(c), s, a, r) for c, s, a, r in state["cookies"]]
        except FileNotFoundError:
            return
        except (OSError, ValueError, KeyError, TypeError) as e:
            log(f"Ignoring unreadable state file {self.state_path}: {e}")
            return
        if bus_id != self.conn.get_guid():
            # Unique names like :1.42 are reused by a new bus, so they would mean other clients.
            if saved:
                log(f"Not restoring {len(saved)} inhibit(s): saved on another session bus")
            return
        # Keep numbering after old cookies, so a late UnInhibit can't hit a new one.
        self.last_cookie = last_cookie
        if not saved:
            return
        if age > self.restore_window:
            # An UnInhibit sent while no instance ran is lost; don't trust an old file.
            log(f"Not restoring {len(saved)} inhibit(s): saved {age:.0f}s ago, "
                f"over the {self.restore_window}s window")
            return
        for cookie, sender, app, reason in saved:
            if not self.has_owner(sender):
                log(f"Not restoring {cookie}: '{app}' '{reason}' ({sender} left the bus)")
                continue
            counts = reason not in self.ignored_reasons
            self.cookies[cookie] = (sender, app, reason, counts)
            note = "" if counts else " (reason ignored, does not block locking)"
            log(f"Restored {cookie}: '{app}' '{reason}' from {self.proc_name(sender)}{note}")

    def save(self, stopped: bool = False) -> None:
        """Write the held inhibits atomically; while any are held, touch the file regularly.

        The touches keep the mtime recent, so after a crash the next instance can tell how
        long ago it happened. A clean stop records the time in `stopped_at` instead.
        """
        if not self.state_ready or not self.state_path:
            return
        assert self.conn is not None
        state = {
            "bus_id": self.conn.get_guid(),
            "last_cookie": self.last_cookie,
            "stopped_at": time.time() if stopped else None,
            "cookies": [[c, s, a, r] for c, (s, a, r, _) in self.cookies.items()],
        }
        tmp = f"{self.state_path}.tmp"
        try:
            with open(tmp, "w") as f:
                json.dump(state, f)
            os.replace(tmp, self.state_path)
        except OSError as e:
            log(f"Cannot write state file {self.state_path}: {e}")
        if self.cookies and not stopped:
            if not self.heartbeat_timer:
                self.heartbeat_timer = GLib.timeout_add_seconds(STATE_HEARTBEAT, self.on_heartbeat)
        elif self.heartbeat_timer:
            GLib.source_remove(self.heartbeat_timer)
            self.heartbeat_timer = 0

    def on_heartbeat(self) -> bool:
        assert self.state_path is not None
        try:
            os.utime(self.state_path)
        except OSError as e:
            log(f"Cannot touch state file {self.state_path}: {e}")
        return GLib.SOURCE_CONTINUE

    # --- D-Bus ------------------------------------------------------------

    def bus_call(self, method: str, name: str, reply_type: str) -> Any:
        """Call an org.freedesktop.DBus method that takes a bus name; return its single result."""
        assert self.conn is not None  # set in on_bus_acquired, before any method call
        return self.conn.call_sync(
            "org.freedesktop.DBus", "/org/freedesktop/DBus", "org.freedesktop.DBus",
            method, GLib.Variant("(s)", (name,)),
            GLib.VariantType(reply_type), Gio.DBusCallFlags.NONE, 1000, None,
        ).unpack()[0]

    def has_owner(self, name: str) -> bool:
        try:
            return bool(self.bus_call("NameHasOwner", name, "(b)"))
        except GLib.Error:
            return False

    def proc_name(self, sender: str) -> str:
        if sender not in self.proc_names:
            try:
                pid = self.bus_call("GetConnectionUnixProcessID", sender, "(u)")
                with open(f"/proc/{pid}/comm") as f:
                    self.proc_names[sender] = f"{f.read().strip()}[{pid}]"
            except (GLib.Error, OSError):
                self.proc_names[sender] = "?"
        return f"{sender} {self.proc_names[sender]}"

    def inhibit(self, sender: str, app: str, reason: str) -> int:
        self.last_cookie += 1
        cookie = self.last_cookie
        counts = reason not in self.ignored_reasons
        self.cookies[cookie] = (sender, app, reason, counts)
        note = "" if counts else " (reason ignored, does not block locking)"
        log(f"Inhibit {cookie}: '{app}' '{reason}' from {self.proc_name(sender)}{note}")
        self.save()
        self.update()
        return cookie

    def uninhibit(self, sender: str, cookie: int) -> None:
        entry = self.cookies.get(cookie)
        if entry is None or entry[0] != sender:
            log(f"UnInhibit {cookie} from {self.proc_name(sender)}: unknown cookie, ignored")
            return
        del self.cookies[cookie]
        _, app, reason, _ = entry
        log(f"UnInhibit {cookie}: '{app}' '{reason}'")
        self.save()
        self.update()

    def simulate_activity(self, sender: str) -> None:
        """Log a sender's calls once per burst: Steam sends thousands per hour."""
        now = time.monotonic()
        burst = self.activity.get(sender)
        if burst is None:
            log(f"SimulateUserActivity from {self.proc_name(sender)} (not implemented)")
            burst = self.activity[sender] = ActivityBurst(first=now, last=now)
        else:
            GLib.source_remove(burst.timer)
        burst.calls += 1
        burst.last = now
        burst.timer = GLib.timeout_add_seconds(ACTIVITY_QUIET, self.on_activity_quiet, sender)

    def on_activity_quiet(self, sender: str) -> bool:
        self.activity[sender].timer = 0
        self.end_activity(sender)
        return GLib.SOURCE_REMOVE

    def end_activity(self, sender: str) -> None:
        burst = self.activity.pop(sender, None)
        if burst is None:
            return
        if burst.timer:
            GLib.source_remove(burst.timer)
        if burst.calls == 1:
            return
        span = timedelta(seconds=round(burst.last - burst.first))
        log(f"SimulateUserActivity from {self.proc_name(sender)}: {burst.calls} calls over {span}")

    def release_sender(self, sender: str) -> None:
        gone = [c for c, (s, *_) in self.cookies.items() if s == sender]
        for c in gone:
            _, app, reason, _ = self.cookies.pop(c)
            log(f"Released {c}: '{app}' '{reason}' ({sender} left the bus)")
        self.end_activity(sender)
        self.proc_names.pop(sender, None)
        if gone:
            self.save()
            self.update()

    def on_method_call(
        self, conn: Gio.DBusConnection, sender: str, path: str, iface: str, method: str,
        params: GLib.Variant, invocation: Gio.DBusMethodInvocation,
    ) -> None:
        if method == "Inhibit":
            app, reason = params.unpack()
            invocation.return_value(GLib.Variant("(u)", (self.inhibit(sender, app, reason),)))
        elif method == "UnInhibit":
            self.uninhibit(sender, params.unpack()[0])
            invocation.return_value(None)
        elif method == "SimulateUserActivity":
            self.simulate_activity(sender)
            invocation.return_value(None)
        elif method == "Lock":
            log(f"Lock from {self.proc_name(sender)}")
            xautolock("-locknow")
            invocation.return_value(None)
        elif method in ("GetActive", "SetActive"):
            invocation.return_value(GLib.Variant("(b)", (False,)))
        elif method in ("GetActiveTime", "GetSessionIdleTime"):
            invocation.return_value(GLib.Variant("(u)", (0,)))

    def on_name_owner_changed(
        self, conn: Gio.DBusConnection, sender: str, path: str, iface: str, signal_name: str,
        params: GLib.Variant,
    ) -> None:
        name, _old, new = params.unpack()
        if not new and name.startswith(":"):
            self.release_sender(name)

    def on_bus_acquired(self, conn: Gio.DBusConnection, name: str) -> None:
        self.conn = conn
        iface = Gio.DBusNodeInfo.new_for_xml(INTROSPECTION).interfaces[0]
        for path in OBJECT_PATHS:
            conn.register_object(path, iface, self.on_method_call, None, None)
        conn.signal_subscribe(
            "org.freedesktop.DBus", "org.freedesktop.DBus", "NameOwnerChanged",
            "/org/freedesktop/DBus", None, Gio.DBusSignalFlags.NONE, self.on_name_owner_changed,
        )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "-i", "--ignore-reason", action="append", metavar="REASON",
        help="inhibit reason that does not block locking (repeatable; default: 'Playing audio')",
    )
    parser.add_argument(
        "--count-all", action="store_true",
        help="let every inhibit block locking, ignoring no reasons",
    )
    parser.add_argument(
        "-d", "--release-delay", type=int, default=5, metavar="SECONDS",
        help="wait this long after the last release before re-enabling xautolock (default: 5)",
    )
    parser.add_argument(
        "--restore-window", type=int, default=30, metavar="SECONDS",
        help="take over inhibits from a previous instance that stopped at most this long ago "
             "(default: 30; 0 disables)",
    )
    args = parser.parse_args()

    ignored: set[str]
    if args.count_all:
        ignored = set()
    else:
        ignored = set(args.ignore_reason or ["Playing audio"])
    runtime_dir = os.environ.get("XDG_RUNTIME_DIR")
    state_path = os.path.join(runtime_dir, STATE_FILE) if runtime_dir else None
    inhibitor = Inhibitor(
        ignored, max(args.release_delay, 0), XBlanking(), state_path, max(args.restore_window, 0),
    )
    loop = GLib.MainLoop()
    exit_code = 0
    acquired = False

    def on_name_acquired(conn: Gio.DBusConnection, name: str) -> None:
        nonlocal acquired
        acquired = True
        log(f"Owning {name}; ignored reasons: {sorted(ignored) or 'none'}")
        # Only after owning the name, so a second instance can't undo an active inhibit
        # or overwrite the state file.
        inhibitor.start()

    def on_name_lost(conn: Gio.DBusConnection | None, name: str) -> None:
        nonlocal exit_code
        if conn is None:
            log("Could not connect to the session bus")
            exit_code = 1
        elif acquired:
            log(f"Lost {name} (bus connection closed?)")
            exit_code = 1
        else:
            log(f"{name} is owned by another process (a desktop environment?); exiting")
            exit_code = EXIT_NAME_TAKEN
        loop.quit()

    def on_signal() -> bool:
        log("Signal received, exiting")
        loop.quit()
        return GLib.SOURCE_REMOVE

    for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal_add(GLib.PRIORITY_HIGH, sig, on_signal)

    owner_id = Gio.bus_own_name(
        Gio.BusType.SESSION, BUS_NAME, Gio.BusNameOwnerFlags.DO_NOT_QUEUE,
        inhibitor.on_bus_acquired, on_name_acquired, on_name_lost,
    )
    try:
        loop.run()
    finally:
        inhibitor.shutdown()
        Gio.bus_unown_name(owner_id)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()