nixos-config/pkgs/bwrap-helper/bwrap-helper.py

#!/usr/bin/env python3
from functools import partial
from itertools import chain
import argparse
import os
import shutil
import subprocess
import time


def flat_map(f, iterable):
    return list(chain.from_iterable(map(f, iterable)))


def add_switch(name: str, default=False):
    if default:
        parser.add_argument(f"--no-{name}", dest=name, action="store_false")
    else:
        parser.add_argument(f"--{name}", dest=name, action="store_true", default=default)


def tmp_file(name: str):
    tmpdir = f"/tmp/bwrap-helper-{os.getpid()}"
    os.makedirs(tmpdir, exist_ok=True)
    if name is None:
        return tmpdir
    else:
        return f"{tmpdir}/{name}"


def generate_tmp_file(name: str, content: str):
    file_path = tmp_file(name)
    with open(file_path, "w") as f:
        f.write(content)
    return file_path


def bind(src: str, dest=None, type=None, required=True):
    arg = "--"
    if type is not None:
        arg += type + "-"
    arg += "bind"
    if not required:
        arg += "-try"

    if dest is None:
        dest = src

    return [arg, src, dest]


def setenv(name: str, value: str):
    return ["--setenv", name, value]


def parse_passthrough_arg(name: str, nargs: int):
    parser.add_argument(f"--{name}", action="append", nargs=nargs)


def assemble_passthrough_arg(name: str):
    for value in getattr(args, name.replace("-", "_")) or []:
        assembled_args.extend([f"--{name}", *value])


dev_bind = partial(bind, type="dev")
ro_bind = partial(bind, type="ro")
ro_bind_try = partial(bind, type="ro", required=False)

username = os.getenv("USER")
uid = os.getuid()
gid = os.getgid()
home = os.getenv("HOME")

path_entries = [
    f"/etc/profiles/per-user/{username}/bin",
    "/run/current-system/sw/bin",
]

argument_groups = {
    "base": (True, [
        "--tmpfs", "/tmp",
        "--proc", "/proc",
        "--dev", "/dev",
        "--dir", home,
        "--dir", f"/run/user/{uid}",
        *ro_bind("/etc/localtime"),
        "--unshare-all",
        "--die-with-parent",
    ]),
    "nix-store": (True, [
        *flat_map(ro_bind, [
            "/nix/store",
            "/etc/static",
        ]),
    ]),
    "path": (True, [
        *flat_map(ro_bind_try, path_entries),
        *setenv("PATH", ":".join(path_entries)),
        *ro_bind_try("/run/current-system/sw")  # not really path, but also libraries etc.
    ]),
    "gui": (False, [
        *dev_bind("/dev/dri"),
        *flat_map(ro_bind, [
            "/sys/dev/char",
            "/sys/devices/pci0000:00",
            f"/run/user/{uid}/{os.getenv('WAYLAND_DISPLAY')}",
            "/run/opengl-driver",
            "/etc/fonts",
        ]),
        *ro_bind_try("/run/opengl-driver-32"),
    ]),
    "x11": (False, [
        *ro_bind("/tmp/.X11-unix"),
    ]),
    "audio": (False, [
        *ro_bind(f"/run/user/{uid}/pulse"),
        # should in theory autodetect, but sometimes it does not work
        *setenv("PULSE_SERVER", f"/run/user/{uid}/pulse/native"),
        # some programs need the cookie
        *ro_bind(f"{home}/.config/pulse/cookie"),
        *setenv("PULSE_COOKIE", f"{home}/.config/pulse/cookie"),
        # ALSA compat
        *ro_bind("/etc/asound.conf"),
    ]),
    "passwd": (False, [
        *ro_bind(
            generate_tmp_file(
                "passwd",
                f"{username}:x:{uid}:{gid}::{home}:/run/current-system/sw/bin/bash\n"
            ),
            "/etc/passwd"
        )
    ]),
    "network": (False, [
        "--share-net",
        *flat_map(ro_bind, [
            "/etc/resolv.conf",
            "/etc/ssl/certs",
        ]),
    ]),
    "dbus": (False, [
        *ro_bind(tmp_file("dbus"), "/run/dbus/system_bus_socket"),
        "--unsetenv", "DBUS_SESSION_BUS_ADDRESS",
        *ro_bind(generate_tmp_file("machine-id", "0" * 32), "/etc/machine-id"),
    ]),
    "new-session": (True, [
        "--new-session",
    ]),
    "pwd": (False, [
        *ro_bind(os.getcwd()),
        "--chdir", os.getcwd(),
    ]),
    "pwd-rw": (False, [
        *bind(os.getcwd()),
        "--chdir", os.getcwd(),
    ]),
}

passthrough_args = [
    ("bind", 2),
    ("ro-bind", 2),
    ("symlink", 2),
]

for (_, arguments) in argument_groups.values():
    for argument in arguments:
        assert type(argument) == str

parser = argparse.ArgumentParser()
parser.add_argument("--show-cmdline", action="store_true")
for name, (default, _) in argument_groups.items():
    add_switch(name, default)
parser.add_argument("program")
parser.add_argument("args", nargs="*")

for (arg, nargs) in passthrough_args:
    parse_passthrough_arg(arg, nargs)

args = parser.parse_args()

assembled_args = ["bwrap"]

for name, (_, arguments) in argument_groups.items():
    if getattr(args, name):
        assembled_args.extend(arguments)

for (arg, _) in passthrough_args:
    assemble_passthrough_arg(arg)

if args.show_cmdline:
    for idx, assembled_arg in enumerate(assembled_args):
        if idx == 0:
            print(assembled_arg, end="")
            continue
        if assembled_arg.startswith("--"):
            print("\n    ", end="")
        else:
            print(end=" ")
        print(assembled_arg, end="")
    print()

assembled_args.append(args.program)
assembled_args.extend(args.args)

children = []
if args.dbus:
    children.append(
        subprocess.Popen(
            [
                "xdg-dbus-proxy",
                os.getenv("DBUS_SESSION_BUS_ADDRESS"),
                tmp_file("dbus"),
                "--filter"
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    )
    while not os.path.exists(tmp_file("dbus")):
        time.sleep(0.1)

try:
    subprocess.run(assembled_args)
finally:
    shutil.rmtree(tmp_file(None))
    for child in children:
        child.terminate()
        child.wait()