# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: GPL-3

import functools
import os
import pathlib
import socket
import unittest

import pytest

import linuxnamespaces


def allow_fork_exit(function):
    @functools.wraps(function)
    def wrapped(*args, **kwargs):
        mainpid = os.getpid()
        try:
            return function(*args, **kwargs)
        except SystemExit as sysexit:
            if sysexit.code or os.getpid() == mainpid:
                raise

            # We're supposed to successfully exit from a child process.  If we
            # were to return or raise here, pytest would record success or
            # failure. Instead we hide this process from pytest.
            os._exit(0)
    return pytest.mark.forked(wrapped)

class IDAllocationTest(unittest.TestCase):
    def test_idalloc(self) -> None:
        alloc = linuxnamespaces.IDAllocation()
        alloc.add_range(1, 2)
        alloc.add_range(5, 4)
        self.assertIn(alloc.find(3), (5, 6))
        self.assertIn(alloc.allocate(3), (5, 6))
        self.assertRaises(ValueError, alloc.find, 3)
        self.assertRaises(ValueError, alloc.allocate, 3)
        self.assertEqual(alloc.find(2), 1)

    def test_merge(self) -> None:
        alloc = linuxnamespaces.IDAllocation()
        alloc.add_range(1, 2)
        alloc.add_range(3, 2)
        self.assertIn(alloc.allocate(3), (1, 2))


class UnshareTest(unittest.TestCase):
    @pytest.mark.forked
    def test_unshare_user(self) -> None:
        overflowuid = int(pathlib.Path("/proc/sys/fs/overflowuid").read_text())
        idmap = linuxnamespaces.IDMapping(0, os.getuid(), 1)
        linuxnamespaces.unshare(linuxnamespaces.CloneFlags.NEWUSER)
        self.assertEqual(os.getuid(), overflowuid)
        linuxnamespaces.newuidmap(-1, [idmap], False)
        self.assertEqual(os.getuid(), 0)
        # UID 1 is not mapped.
        self.assertRaises(OSError, os.setuid, 1)

    @allow_fork_exit
    def test_mount_proc(self) -> None:
        idmap = linuxnamespaces.IDMapping(0, os.getuid(), 1)
        linuxnamespaces.unshare(
            linuxnamespaces.CloneFlags.NEWUSER
            | linuxnamespaces.CloneFlags.NEWNS
            | linuxnamespaces.CloneFlags.NEWPID
        )
        linuxnamespaces.newuidmap(-1, [idmap], False)
        @linuxnamespaces.run_in_fork
        def setup() -> None:
            self.assertEqual(os.getpid(), 1)
            linuxnamespaces.mount("proc", "/proc", "proc")
        setup()

    @pytest.mark.forked
    def test_sethostname(self) -> None:
        self.assertRaises(socket.error, socket.sethostname, "example")
        linuxnamespaces.unshare(
            linuxnamespaces.CloneFlags.NEWUSER
            | linuxnamespaces.CloneFlags.NEWUTS
        )
        socket.sethostname("example")

    @pytest.mark.forked
    def test_populate_dev(self) -> None:
        uidmap = linuxnamespaces.IDMapping(0, os.getuid(), 1)
        gidmap = linuxnamespaces.IDMapping(0, os.getgid(), 1)
        linuxnamespaces.unshare(
            linuxnamespaces.CloneFlags.NEWUSER
            | linuxnamespaces.CloneFlags.NEWNS
        )
        pathlib.Path("/proc/self/setgroups").write_text("deny")
        linuxnamespaces.newuidmap(-1, [uidmap], False)
        linuxnamespaces.newgidmap(-1, [gidmap], False)
        linuxnamespaces.mount("tmpfs", "/mnt", "tmpfs", data="mode=0755")
        os.mkdir("/mnt/dev")
        linuxnamespaces.populate_dev("/", "/mnt", pidns=False)
        self.assertTrue(os.access("/mnt/dev/null", os.W_OK))
        pathlib.Path("/mnt/dev/null").write_text("")


class UnshareIdmapTest(unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        self.uidalloc = linuxnamespaces.IDAllocation.loadsubid("uid")
        self.gidalloc = linuxnamespaces.IDAllocation.loadsubid("gid")
        try:
            self.uidalloc.find(65536)
            self.gidalloc.find(65536)
        except ValueError:
            self.skipTest("insufficient /etc/sub?id allocation")

    @allow_fork_exit
    def test_unshare_user_idmap(self) -> None:
        overflowuid = int(pathlib.Path("/proc/sys/fs/overflowuid").read_text())
        uidmap = linuxnamespaces.IDMapping(
            0, self.uidalloc.allocate(65536), 65536
        )
        self.assertNotEqual(os.getuid(), uidmap.outerstart)
        gidmap = linuxnamespaces.IDMapping(
            0, self.gidalloc.allocate(65536), 65536
        )
        pid = os.getpid()
        @linuxnamespaces.run_in_fork
        def setup() -> None:
            linuxnamespaces.newgidmap(pid, [gidmap])
            linuxnamespaces.newuidmap(pid, [uidmap])
        linuxnamespaces.unshare(linuxnamespaces.CloneFlags.NEWUSER)
        setup()
        self.assertEqual(os.getuid(), overflowuid)
        os.setuid(0)
        self.assertEqual(os.getuid(), 0)
        os.setuid(1)
        self.assertEqual(os.getuid(), 1)

    @allow_fork_exit
    def test_populate_dev(self) -> None:
        uidmap = linuxnamespaces.IDMapping(
            0, self.uidalloc.allocate(65536), 65536
        )
        self.assertNotEqual(os.getuid(), uidmap.outerstart)
        gidmap = linuxnamespaces.IDMapping(
            0, self.gidalloc.allocate(65536), 65536
        )
        pid = os.getpid()
        @linuxnamespaces.run_in_fork
        def setup() -> None:
            linuxnamespaces.newgidmap(pid, [gidmap])
            linuxnamespaces.newuidmap(pid, [uidmap])
        linuxnamespaces.unshare(
            linuxnamespaces.CloneFlags.NEWUSER
            | linuxnamespaces.CloneFlags.NEWNS
            | linuxnamespaces.CloneFlags.NEWPID
        )
        setup()
        os.setreuid(0, 0)
        os.setregid(0, 0)
        linuxnamespaces.mount("tmpfs", "/mnt", "tmpfs")
        os.mkdir("/mnt/dev")
        @linuxnamespaces.run_in_fork
        def test() -> None:
            linuxnamespaces.populate_dev("/", "/mnt")
        test()