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

"""Mount a given filesystem image inside a user and mount namespace using
an unprivileged fuse driver and chroot into the mounted filesystem. Supported
drivers are fuse2fs and squashfuse.

This requires a fuse package with support for /dev/fd/N. Either with fuse2 via
https://bugs.debian.org/1055222 or fuse3. On Debian, squashfuse uses fuse3.
"""

import argparse
import os
import pathlib
import socket
import sys

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

import linuxnamespaces


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fstype", choices=["ext4", "squashfs"])
    parser.add_argument("fsimage", type=pathlib.Path)
    args = parser.parse_args()
    assert args.fsimage.exists()
    if args.fstype is None:
        args.fstype = args.fsimage.suffix.removeprefix(".")
    if args.fstype not in ("ext4", "squashfs"):
        print("Cannot determine filesystem type for image.")
        sys.exit(1)
    uidmap = linuxnamespaces.IDAllocation.loadsubid("uid").allocatemap(65536)
    gidmap = linuxnamespaces.IDAllocation.loadsubid("gid").allocatemap(65536)
    mainpid = os.getpid()
    mainsock, childsock = socket.socketpair()
    @linuxnamespaces.run_in_fork
    def setup() -> None:
        mainsock.close()
        linuxnamespaces.newidmaps(mainpid, [uidmap], [gidmap])
        childsock.send(b"\0")
        _, fds, _, _ = socket.recv_fds(childsock, 1, 1, 0)
        childsock.close()
        os.set_inheritable(fds[0], True)
        driver = {
            "ext4": "fuse2fs",
            "squashfs": "squashfuse",
        }[args.fstype]
        os.execlp(driver, driver, str(args.fsimage), "/dev/fd/%d" % fds[0])
    childsock.close()
    linuxnamespaces.unshare(
        linuxnamespaces.CloneFlags.NEWUSER | linuxnamespaces.CloneFlags.NEWNS
    )
    setup.start()
    mainsock.recv(1)
    os.setreuid(0, 0)
    os.setregid(0, 0)
    fusefd = os.open("/dev/fuse", os.O_RDWR)
    socket.send_fds(mainsock, [b"\0"], [fusefd])
    mainsock.close()
    readonly = {
        "ext4": False,
        "squashfs": True,
    }[args.fstype]
    linuxnamespaces.mount(
        str(args.fsimage),
        "/mnt",
        "fuse." + args.fstype,
        linuxnamespaces.MountFlags.RDONLY
        if readonly
        else linuxnamespaces.MountFlags.NONE,
        [
            "fd=%d" % fusefd,
            "rootmode=040755",
            "user_id=0",
            "group_id=0",
            "allow_other",
        ],
    )
    os.chdir("/mnt")
    linuxnamespaces.bind_mount("/proc", "proc", recursive=True)
    linuxnamespaces.bind_mount("/sys", "sys", recursive=True)
    linuxnamespaces.populate_dev("/", ".", pidns=False, tun=False)
    if readonly:
        linuxnamespaces.mount(
            "tmpfs", "tmp", "tmpfs", linuxnamespaces.MountFlags.NODEV
        )
        linuxnamespaces.mount(
            "tmpfs",
            "run",
            "tmpfs",
            linuxnamespaces.MountFlags.NODEV
            | linuxnamespaces.MountFlags.NOSUID
            | linuxnamespaces.MountFlags.NOEXEC,
            "mode=0755",
        )
    linuxnamespaces.pivot_root(".", ".")
    linuxnamespaces.umount(".", linuxnamespaces.UmountFlags.DETACH)
    os.close(fusefd)
    os.execlp(os.environ["SHELL"], os.environ["SHELL"])


if __name__ == "__main__":
    main()