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 size_screen(self, settings: Any) -> None:
    """Start the screen at the manual resolution where one is set, else at
    `START_SIZE`, through `selkies resize`. A server that cannot be resized
    keeps its whole framebuffer, as before."""
    size = START_SIZE
    if settings.manual_resolution[0] and settings.manual_width and settings.manual_height:
        size = (int(settings.manual_width), int(settings.manual_height))
    try:
        done = subprocess.run([sys.executable, "-m", "selkies", "resize", "%dx%d" % size],
                              env=self.env, stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=30)
        realized = re.fullmatch(r"\((\d+), (\d+)\)", done.stdout.strip())
    except subprocess.TimeoutExpired:
        realized = None
    if realized:
        log(f"the screen starts at {realized[1]}x{realized[2]}")
    else:
        log("the screen could not be sized; it starts at the whole framebuffer")

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()
        x11_socket_dir()
    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 outside the session's process groups: those carrying
    its runtime directory, and the orphans the launcher adopted, since an
    agent that makes itself undumpable (ssh-agent) hides its environment."""
    marker = f"XDG_RUNTIME_DIR={self.runtime_dir}".encode()
    launcher, ours = str(os.getpid()).encode(), {proc.pid for proc in self.children}
    found = []
    for entry in filter(str.isdigit, os.listdir("/proc")):
        try:
            with open(f"/proc/{entry}/stat", "rb") as fh:
                adopted = fh.read().rsplit(b")", 1)[1].split()[1] == launcher
            if adopted and int(entry) not in ours:
                found.append(int(entry))
                continue
            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 reap(self) -> None:
    """Collect the children that exited: the session's own through their
    Popen, which keeps their status, and the orphans it adopted."""
    for proc in self.children:
        proc.poll()
    ours = {proc.pid for proc in self.children}
    while True:
        try:
            exited = os.waitid(os.P_ALL, 0, os.WEXITED | os.WNOHANG | os.WNOWAIT)
        except ChildProcessError:
            return
        if exited is None or exited.si_pid in ours:
            return
        os.waitpid(exited.si_pid, 0)

def signal_session(self, sig: int) -> None:
    """The session's process groups, then whatever left them."""
    for pgid in self.groups:
        try:
            os.killpg(pgid, sig)
        except OSError:
            pass
    for pid in self.stragglers():
        try:
            os.kill(pid, 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_session(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_session(signal.SIGKILL)
    for proc in self.children:
        proc.wait()
    self.reap()
    shutil.rmtree(self.runtime_dir, ignore_errors=True)