Session

def __init__(self) -> None:
    # A runtime directory an image's own init would have made is no place to start from
    base = os.environ.get("XDG_RUNTIME_DIR") or ""
    self.runtime_dir = tempfile.mkdtemp(prefix="selkies-session-",
                                        dir=base if os.path.isdir(base) and os.access(base, os.W_OK) else None)
    self.env = {k: v for k, v in os.environ.items() if k not in HOST_SESSION_VARS}
    self.env["XDG_RUNTIME_DIR"] = self.runtime_dir
    self.children: List[subprocess.Popen] = []
    self.groups: Set[int] = set()
    self.selkies: Optional[subprocess.Popen] = None

def spawn(self, argv: List[str], env: Optional[Dict[str, str]] = None, **extra: Any) -> subprocess.Popen:
    group = next(iter(self.groups), 0)
    proc = subprocess.Popen(argv, env=env or self.env, stdin=subprocess.DEVNULL,
                            preexec_fn=lambda: join_group(group), **extra)
    try:
        self.groups.add(os.getpgid(proc.pid))
    except ProcessLookupError:
        pass
    self.children.append(proc)
    return proc

def sound(self) -> bool:
    """A sound server Selkies can capture: the one the environment names or
    runs, else the first of PipeWire and PulseAudio that comes up here. A
    named socket that nothing serves, as an image's own init would, is passed over."""
    named = self.env.get("PULSE_SERVER", "")
    socket_path = named[len("unix:"):] if named.startswith("unix:") else named if named.startswith("/") else ""
    if named and (not socket_path or answers(socket_path)):
        return True
    running = os.environ.get("PULSE_RUNTIME_PATH") or os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/nonexistent"), "pulse")
    if answers(os.path.join(running, "native")):
        self.env["PULSE_SERVER"] = f"unix:{os.path.join(running, 'native')}"
        return True
    own = os.path.join(self.runtime_dir, "pulse", "native")
    manager = next((n for n in ("wireplumber", "pipewire-media-session") if shutil.which(n)), None)
    candidates = [[["pipewire"], [manager], ["pipewire-pulse"]]] if manager else []
    candidates.append([["pulseaudio", "--daemonize=no", "--exit-idle-time=-1"]])
    for servers in candidates:
        if not all(shutil.which(argv[0]) for argv in servers):
            continue
        procs = [self.spawn(argv) for argv in servers]
        try:
            wait_for(lambda: answers(own), 20, servers[-1][0], procs[-1])
        except RuntimeError as err:
            log(str(err))
            for proc in procs:
                proc.kill()
            continue
        self.env["PULSE_SERVER"] = f"unix:{own}"
        log(f"{servers[0][0]} serves the session's audio")
        return True
    log("no sound server came up; the session has no audio")
    return False

def x11(self, render_dri: str) -> None:
    """An Xvfb on the first free display that admits this session's cookie
    alone, rendering on the GPU through glamor where the server offers it.

    A server whose GLX cannot load (a vendor's GL stack beside a stock
    Xvfb) is started again without it, so the desktop comes up and only
    GL applications on X11 go without.
    """
    if not shutil.which("Xvfb"):
        raise RuntimeError("the X11 backend needs Xvfb; install it, or set SELKIES_WAYLAND=true")
    auth = os.path.join(self.runtime_dir, "Xauthority")
    with open(auth, "wb") as fh:
        # One FamilyWild entry: the display number is the server's to pick
        fh.write(struct.pack(">HHHH18sH16s", 0xFFFF, 0, 0, 18, b"MIT-MAGIC-COOKIE-1", 16, os.urandom(16)))
    node = render_dri or next(iter(sorted(glob.glob("/dev/dri/renderD*"))), "")
    usage = subprocess.run(["Xvfb", "-help"], capture_output=True, text=True)
    glx = ["+extension", "GLX"]
    gpu = [glx + ["-glamor", "-dri", node]] if node and "-glamor" in usage.stdout + usage.stderr else []
    for extra in gpu + [glx, ["-extension", "GLX"]]:
        read, write = os.pipe()
        proc = self.spawn(["Xvfb", "-displayfd", str(write), "-auth", auth, *XVFB_ARGS, *extra],
                          pass_fds=(write,))
        os.close(write)
        with os.fdopen(read) as pipe:
            number = pipe.readline().strip() if select.select([pipe], [], [], 30)[0] else ""
        if number:
            self.env.update(DISPLAY=f":{number}", XAUTHORITY=auth)
            log(f"Xvfb serves display :{number}" + (f", rendering on {node}" if "-glamor" in extra else "")
                + (", without GLX, which it could not load" if "-extension" in extra else ""))
            return
        proc.kill()
    raise RuntimeError("Xvfb did not come up")

def capture_socket(self) -> str:
    """The socket of Selkies' compositor, the first in the runtime directory."""
    def found() -> str:
        return next((n for n in sorted(os.listdir(self.runtime_dir)) if n.startswith("wayland-")
                     and not n.endswith(".lock") and answers(os.path.join(self.runtime_dir, n))), "")
    wait_for(found, 60, "Selkies' compositor", self.selkies)
    return found()

def desktop(self, sid: str, entry: Dict[str, str], wayland: bool) -> None:
    """The session as a display manager starts one, on a session bus of its own."""
    env = dict(self.env, XDG_SESSION_TYPE="wayland" if wayland else "x11")
    if wayland:
        env["WAYLAND_DISPLAY"] = self.capture_socket()
    names = entry.get("DesktopNames", "").strip(";").replace(";", ":")
    if names:
        env["XDG_CURRENT_DESKTOP"] = names
    else:
        env.pop("XDG_CURRENT_DESKTOP", None)
    if sid:
        env.update(XDG_SESSION_DESKTOP=sid, DESKTOP_SESSION=sid)
    # The interposers serve the session's applications; Selkies keeps the real device nodes
    preload = [env.get(v, "") for v in ("SELKIES_INTERPOSER", "SELKIES_WEBCAM_INTERPOSER")]
    preload = [p for p in preload if os.path.isfile(p)] + [p for p in env.get("LD_PRELOAD", "").split(":") if p]
    if preload:
        env["LD_PRELOAD"] = ":".join(preload)
    bus = ["dbus-run-session", "--"] if shutil.which("dbus-run-session") else []
    self.spawn(bus + shlex.split(entry["Exec"]), env=env)
    log(f"desktop {sid or entry['Exec']} on {env.get('WAYLAND_DISPLAY') or env['DISPLAY']}")

def stragglers(self) -> List[int]:
    """The processes still carrying the session's runtime directory, such
    as the agents a desktop daemonizes out of its process group."""
    marker = f"XDG_RUNTIME_DIR={self.runtime_dir}".encode()
    found = []
    for entry in filter(str.isdigit, os.listdir("/proc")):
        try:
            with open(f"/proc/{entry}/environ", "rb") as fh:
                if marker in fh.read().split(b"\0"):
                    found.append(int(entry))
        except OSError:
            continue
    return found

def signal_groups(self, sig: int) -> None:
    for pgid in self.groups:
        try:
            os.killpg(pgid, sig)
        except OSError:
            pass

def stop(self) -> None:
    """Selkies first, so its capture closes before its display, then the
    session's process groups and whatever left them, and the runtime directory."""
    if self.selkies is not None and self.selkies.poll() is None:
        self.selkies.terminate()
        try:
            self.selkies.wait(timeout=10)
        except subprocess.TimeoutExpired:
            self.selkies.kill()
    self.signal_groups(signal.SIGTERM)
    deadline = time.monotonic() + 5
    for proc in self.children:
        try:
            proc.wait(timeout=max(0.1, deadline - time.monotonic()))
        except subprocess.TimeoutExpired:
            pass
    self.signal_groups(signal.SIGKILL)
    for pid in self.stragglers():
        try:
            os.kill(pid, signal.SIGKILL)
        except OSError:
            pass
    for proc in self.children:
        proc.wait()
    shutil.rmtree(self.runtime_dir, ignore_errors=True)