#!/usr/bin/python3
# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: GPL-3

"""Unshare a cgroup (and user) namespace such that the entire cgroup hierarchy
(inside the namespace) becomes writeable to the user.
"""

import asyncio
import os
import pathlib
import sys

if __file__.split("/")[-2:-1] == ["examples"]:
    sys.path.insert(0, "/".join(__file__.split("/")[:-2]))

import linuxnamespaces
import linuxnamespaces.systemd


def get_cgroup(pid: int = -1) -> pathlib.PurePath:
    """Look up the cgroup that the given pid or the running process belongs
    to.
    """
    return pathlib.PurePath(
        pathlib.Path(
            f"/proc/{pid}/cgroup" if pid > 0 else "/proc/self/cgroup"
        ).read_text().split(":", 2)[2].strip()
    )


def main() -> None:
    mycgroup = get_cgroup()
    if not os.access(
        pathlib.Path("/sys/fs/cgroup") / mycgroup.relative_to("/"),
        os.W_OK,
    ):
        # For some shells - notably from graphical desktop environments, the
        # hierarchy is immediately writeable. For others, we may create a scope
        # unit.
        try:
            asyncio.run(
                linuxnamespaces.systemd.start_transient_unit(
                    f"cgroup-{os.getpid()}.scope",
                    properties={"Delegate": True},
                ),
            )
            mycgroup = get_cgroup()
        except NotImplementedError:
            linuxnamespaces.systemd.reexec_as_transient_unit(
                properties={"Delegate": True}
            )
    namespaces = (
        linuxnamespaces.CloneFlags.NEWUSER
        | linuxnamespaces.CloneFlags.NEWNS
        | linuxnamespaces.CloneFlags.NEWCGROUP
    )
    linuxnamespaces.unshare_user_idmap(
        [linuxnamespaces.IDMapping(os.getuid(), os.getuid(), 1)],
        [linuxnamespaces.IDMapping(os.getgid(), os.getgid(), 1)],
        namespaces,
    )
    linuxnamespaces.populate_sys("/", "/", namespaces, mycgroup)
    os.execlp(os.environ["SHELL"], os.environ["SHELL"])


if __name__ == "__main__":
    main()