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

"""Extract a given tarball into a temporary location and chroot into it inside
a user and mount namespace.
"""

import argparse
import os
import pathlib
import socket
import sys
import tempfile

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

import linuxnamespaces
import linuxnamespaces.tarutils


class TarFile(
    linuxnamespaces.tarutils.ZstdTarFile, linuxnamespaces.tarutils.XAttrTarFile
):
    pass


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--save",
        action="store_true",
        help="save and replace the tarball at the end of the session",
    )
    parser.add_argument(
        "basetar",
        type=pathlib.Path,
        action="store",
        help="location of the tarball containing the chroot",
    )
    parser.add_argument(
        "command",
        nargs=argparse.REMAINDER,
        help="command to run inside the chroot",
    )
    args = parser.parse_args()
    assert args.basetar.exists()
    uidmap = linuxnamespaces.IDAllocation.loadsubid("uid").allocatemap(65536)
    gidmap = linuxnamespaces.IDAllocation.loadsubid("gid").allocatemap(65536)
    with tempfile.TemporaryDirectory() as tdir:
        parentsock, childsock = socket.socketpair()
        pid = os.fork()
        if pid == 0:
            parentsock.close()
            # Once we drop privileges via setreuid and friends, we may become
            # unable to open basetar or to chdir to tdir, so do those early.
            with TarFile.open(args.basetar, "r:*") as tarf:
                os.chdir(tdir)
                linuxnamespaces.unshare(
                    linuxnamespaces.CloneFlags.NEWUSER
                    | linuxnamespaces.CloneFlags.NEWNS
                )
                childsock.send(
                    linuxnamespaces.tarutils.get_comptype(
                        tarf
                    ).encode("ascii") + b"\0",
                )
                childsock.recv(1)
                childsock.close()
                # The other process will now have set up our id mapping and
                # will have changed ownership of our working directory.
                os.setreuid(0, 0)
                os.setregid(0, 0)
                os.setgroups([])
                for tmem in tarf:
                    if tmem.name.removeprefix("./").startswith("dev/"):
                        continue
                    tarf.extract(tmem, numeric_owner=True)
            linuxnamespaces.bind_mount(".", "/mnt", recursive=True)
            os.chdir("/mnt")
            linuxnamespaces.bind_mount("/proc", "proc", recursive=True)
            linuxnamespaces.bind_mount("/sys", "sys", recursive=True)
            linuxnamespaces.populate_dev("/", ".", pidns=False, tun=False)
            linuxnamespaces.pivot_root(".", ".")
            linuxnamespaces.umount(".", linuxnamespaces.UmountFlags.DETACH)
            if args.command:
                os.execvp(args.command[0], args.command)
            else:
                os.execlp(os.environ["SHELL"], os.environ["SHELL"])
            os._exit(1)

        childsock.close()
        comptype = parentsock.recv(10).split(b"\0", 1)[0].decode("ascii")
        linuxnamespaces.newidmaps(pid, [uidmap], [gidmap])
        # We still had to be in the initial namespace to call newidmaps and
        # now we transition to a namespace that can access both the container
        # and the files of the invoking user.
        linuxnamespaces.unshare_user_idmap(
            [uidmap, linuxnamespaces.IDMapping(65536, os.getuid(), 1)],
            [gidmap, linuxnamespaces.IDMapping(65536, os.getgid(), 1)],
        )
        os.chown(tdir, 0, 0)
        os.chmod(tdir, 0o755)
        parentsock.send(b"\0")
        parentsock.close()
        _, ret = os.waitpid(pid, 0)
        if args.save and ret == 0:
            tmptar = f"{args.basetar}.new"
            try:
                with TarFile.open(tmptar, "x:" + comptype) as tout:
                    tout.add(tdir, ".")
                os.rename(tmptar, args.basetar)
            except:
                os.unlink(tmptar)
                raise
    sys.exit(ret)


if __name__ == "__main__":
    main()