#!/usr/bin/python3
# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: LGPL-2.0-or-later

"""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 re
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(
        "--same-uid",
        action="store_true",
        help="map the current user to itself in the namespace",
    )
    parser.add_argument(
        "--bind",
        action="append",
        help="bind mount the given location after extraction",
        default=[],
    )
    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()
    myuid = os.getuid()
    mygid = os.getgid()
    uidrange = linuxnamespaces.IDAllocation.loadsubid("uid").allocatemap(65536)
    gidrange = linuxnamespaces.IDAllocation.loadsubid("gid").allocatemap(65536)
    if args.same_uid:
        uidmaps = [
            uidrange[:myuid],
            linuxnamespaces.IDMapping(uidrange.innerstart + myuid, myuid, 1),
            uidrange[myuid + 1:],
        ]
        gidmaps = [
            gidrange[:mygid],
            linuxnamespaces.IDMapping(gidrange.innerstart + mygid, mygid, 1),
            gidrange[mygid + 1:],
        ]
    else:
        uidmaps = [uidrange]
        gidmaps = [gidrange]
    bindmounts = []
    for conf in args.bind:
        confparts = conf.split(":")
        if not os.path.exists(confparts[0]):
            raise ValueError(
                f"cannot bind mount {confparts[0]}: does not exist"
            )
        if len(confparts) > 2:
            raise ValueError(f"bind mount {conf} not understood")
        if len(confparts) < 2:
            confparts.append(confparts[0])
        bindmounts.append(
            (
                os.path.abspath(confparts[0]),
                os.path.normpath(
                    os.path.join("/", confparts[1])
                ).removeprefix("/"),
            ),
        )
    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:
                    name = re.sub(r"^/*(\.{1,2}/+)*", "", tmem.name)
                    if name.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("/", ".", pts="host", tun=False)
            for source, target in bindmounts:
                os.makedirs(target, exist_ok=True)
                linuxnamespaces.bind_mount(source, target, recursive=True)
            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, uidmaps, gidmaps)
        # 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.
        if not args.same_uid:
            uidmaps.append(linuxnamespaces.IDMapping(65536, myuid, 1))
            gidmaps.append(linuxnamespaces.IDMapping(65536, mygid, 1))
        linuxnamespaces.unshare_user_idmap(uidmaps, gidmaps)
        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()