sockkeeper.py

"""sockkeeper — keep your Unix sockets alive through /tmp cleaners.

Problem: services bind a Unix socket under /tmp, /tmp gets cleaned (systemd-
tmpfiles, custom cron jobs, aggressive Python scripts), and the service keeps
running with a valid file descriptor but an unreachable path. New clients
cannot connect — the service is invisibly broken until someone restarts it.

This library solves that two ways:

1. ``SockKeeper`` — a drop-in replacement for the usual bind-listen pattern
   that optionally runs a background thread rebinding the socket path if it
   disappears. Zero deps beyond the stdlib.

2. ``sockkeeper`` CLI — emits a systemd-tmpfiles.d snippet that tells
   ``systemd-tmpfiles --clean`` to leave your socket alone, and a
   recommendation for moving to ``$XDG_RUNTIME_DIR`` where sockets belong.

Basic usage:

    from sockkeeper import SockKeeper

    with SockKeeper("/tmp/myservice.sock", watch=True) as keeper:
        while True:
            conn, _ = keeper.sock.accept()
            ...

The ``watch=True`` mode spawns a daemon thread that checks every
``watch_interval`` seconds whether the socket path still exists. If not, it
rebinds to the same path (without interrupting existing connections, which
are held by fd on the service side).
"""

from __future__ import annotations

import os
import socket
import stat
import sys
import threading
import time
from pathlib import Path

__version__ = "0.1.0"
__all__ = ["SockKeeper", "recommend_path", "tmpfiles_snippet"]


class SockKeeper:
    """Bind a Unix SOCK_STREAM socket to ``path`` and keep it there.

    Parameters
    ----------
    path : str
        Filesystem path for the socket. Parent directory must exist.
    mode : int, default 0o660
        Permissions for the socket file.
    backlog : int, default 16
        ``listen()`` backlog.
    watch : bool, default False
        If True, start a background thread that rebinds the path if it
        disappears (e.g. deleted by /tmp cleaner).
    watch_interval : float, default 60.0
        Seconds between existence checks when watch=True.
    on_rebind : callable, optional
        Called with no arguments after each successful rebind. Useful for
        logging.
    """

    def __init__(
        self,
        path: str | os.PathLike,
        mode: int = 0o660,
        backlog: int = 16,
        watch: bool = False,
        watch_interval: float = 60.0,
        on_rebind=None,
    ):
        self.path = str(path)
        self.mode = mode
        self.backlog = backlog
        self.watch_flag = watch
        self.watch_interval = watch_interval
        self.on_rebind = on_rebind
        self.sock: socket.socket | None = None
        self._watcher: threading.Thread | None = None
        self._stop = threading.Event()

    # ── Lifecycle ──────────────────────────────────────────────────────────

    def bind(self) -> socket.socket:
        """Bind and listen. Returns the listening socket."""
        self._unlink_stale()
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.bind(self.path)
        os.chmod(self.path, self.mode)
        self.sock.listen(self.backlog)
        if self.watch_flag:
            self._start_watcher()
        return self.sock

    def close(self) -> None:
        self._stop.set()
        if self._watcher is not None:
            self._watcher.join(timeout=2)
        if self.sock is not None:
            try:
                self.sock.close()
            finally:
                self.sock = None
        try:
            os.unlink(self.path)
        except FileNotFoundError:
            pass

    def __enter__(self) -> "SockKeeper":
        self.bind()
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # ── Rebind watcher ─────────────────────────────────────────────────────

    def _start_watcher(self) -> None:
        self._watcher = threading.Thread(
            target=self._watch_loop, name="sockkeeper-watch", daemon=True
        )
        self._watcher.start()

    def _watch_loop(self) -> None:
        while not self._stop.wait(self.watch_interval):
            if not self._path_is_live():
                try:
                    self._rebind()
                    if self.on_rebind is not None:
                        try:
                            self.on_rebind()
                        except Exception:
                            pass
                except OSError:
                    # Try again next tick.
                    pass

    def _path_is_live(self) -> bool:
        """True when ``path`` exists and is a socket."""
        try:
            st = os.lstat(self.path)
        except FileNotFoundError:
            return False
        return stat.S_ISSOCK(st.st_mode)

    def _rebind(self) -> None:
        """Open a fresh listening socket at ``path``.

        The previous socket (self.sock) may still have active connections via
        fd; we keep it alive for those. New connections land on the fresh
        socket, which replaces ``self.sock`` so future ``accept()`` calls use it.
        """
        old = self.sock
        fresh = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        # Clean up any leftover at the path (might be a regular file etc.)
        try:
            os.unlink(self.path)
        except FileNotFoundError:
            pass
        fresh.bind(self.path)
        os.chmod(self.path, self.mode)
        fresh.listen(self.backlog)
        self.sock = fresh
        # Close the old listening socket — any accepted connections remain
        # alive on their own fd; only the listening path is replaced.
        if old is not None:
            try:
                old.close()
            except OSError:
                pass

    def _unlink_stale(self) -> None:
        try:
            st = os.lstat(self.path)
        except FileNotFoundError:
            return
        if stat.S_ISSOCK(st.st_mode):
            try:
                os.unlink(self.path)
            except FileNotFoundError:
                pass


# ── Recommendations ────────────────────────────────────────────────────────


def recommend_path(path: str) -> dict:
    """Return a dict explaining whether ``path`` is a safe socket location.

    Safe rankings:
    - ``ok`` — path is under $XDG_RUNTIME_DIR (per-user, cleaned on logout)
    - ``caution`` — path is under /run or /var/run (system-managed runtime)
    - ``risky`` — path is under /tmp (subject to tmp cleaners)
    - ``unknown`` — something else
    """
    p = Path(path).resolve()
    xdg = os.environ.get("XDG_RUNTIME_DIR")
    if xdg and str(p).startswith(xdg):
        rating = "ok"
        reason = "under $XDG_RUNTIME_DIR — the canonical runtime-state location"
    elif str(p).startswith(("/run/", "/var/run/")):
        rating = "caution"
        reason = "under /run — system-level runtime; OK for system services"
    elif str(p).startswith("/tmp/"):
        rating = "risky"
        reason = (
            "under /tmp — subject to systemd-tmpfiles cleaners and ad-hoc "
            "cleanup scripts. Consider $XDG_RUNTIME_DIR or add a tmpfiles.d "
            "exclude (see: sockkeeper tmpfiles-snippet)."
        )
    else:
        rating = "unknown"
        reason = "unusual location; no recommendation"
    return {"path": str(p), "rating": rating, "reason": reason}


def tmpfiles_snippet(path: str, user: str | None = None) -> str:
    """Return a systemd-tmpfiles.d snippet that protects ``path`` from cleaning."""
    user = user or os.environ.get("USER", "root")
    return (
        "# Protect a Unix socket from /tmp cleaners.\n"
        "# Drop this into /etc/tmpfiles.d/sockkeeper.conf (or ~/.config/user-tmpfiles.d/ "
        "for user-scoped cleanup).\n"
        "#\n"
        "# Format: Type Path Mode User Group Age Argument\n"
        "# The 'x' type excludes the path (and its children) from cleaning.\n"
        f"x {path}\n"
    )


# ── CLI ────────────────────────────────────────────────────────────────────


def _cli():
    import argparse

    ap = argparse.ArgumentParser(
        prog="sockkeeper",
        description="Keep Unix sockets alive through /tmp cleaners.",
    )
    sub = ap.add_subparsers(dest="cmd", required=True)

    p_rec = sub.add_parser("check", help="Check if a socket path is safe from cleaners.")
    p_rec.add_argument("path")

    p_snip = sub.add_parser(
        "tmpfiles-snippet",
        help="Emit a systemd-tmpfiles.d snippet that protects the path.",
    )
    p_snip.add_argument("path")

    p_demo = sub.add_parser(
        "demo",
        help="Bind a socket with watch=True; simulate deletion to show recovery.",
    )
    p_demo.add_argument("path")

    args = ap.parse_args()

    if args.cmd == "check":
        r = recommend_path(args.path)
        print(f"{r['rating'].upper()}: {r['reason']}")
        sys.exit(0 if r["rating"] == "ok" else 1)

    if args.cmd == "tmpfiles-snippet":
        print(tmpfiles_snippet(args.path), end="")
        sys.exit(0)

    if args.cmd == "demo":
        def on_rebind():
            print(f"[sockkeeper] rebound {args.path}")

        with SockKeeper(args.path, watch=True, watch_interval=1.0, on_rebind=on_rebind) as k:
            print(f"listening on {args.path} (watch=1s). unlink the file to see a rebind.")
            time.sleep(10)
            # Simulate external cleaner removing the socket.
            try:
                os.unlink(args.path)
                print(f"[demo] removed {args.path}; waiting for watcher to notice...")
            except FileNotFoundError:
                pass
            time.sleep(4)
        sys.exit(0)


if __name__ == "__main__":
    _cli()