#!/usr/bin/env python3

# SPDX-FileCopyrightText: (c) 2026 Daniel Hast
#
# SPDX-License-Identifier: Apache-2.0 OR MIT

from pathlib import Path
from tempfile import TemporaryFile
from typing import IO, Final
from unittest.mock import patch

from src.bwrap_restricted import Config, run_with_config

TEST_CONFIG: Final[Config] = Config(
    bwrap_path=Path("/usr/bin/bwrap"),
    forbidden_options=[b"--cap-add"],
    extra_options=[b"--cap-drop", b"ALL", b"--new-session", b"--disable-userns"],
)


def split_cmdline(cmdline: str) -> list[bytes]:
    return [bytes(arg, encoding="utf8") for arg in cmdline.split()]


def assert_bwrap_args(input_args: list[bytes], config: Config, expected_args: list[bytes]) -> None:
    argv = [b"bwrap-restricted", *input_args]
    with patch("os.execv") as mock_exec:
        run_with_config(argv, config)
        assert mock_exec.call_args == ((expected_args[0], expected_args), {})


def assert_valid_options_end_at_index(
    input_args: str, index: int, config: Config = TEST_CONFIG
) -> None:
    in_args = split_cmdline(input_args)
    expected_args = [
        bytes(config.bwrap_path),
        *in_args[:index],
        *config.extra_options,
        *in_args[index:],
    ]
    assert_bwrap_args(in_args, config, expected_args)


def assert_rejected(input_args: str, config: Config = TEST_CONFIG) -> None:
    in_args = split_cmdline(input_args)
    argv = [b"bwrap-restricted", *in_args]
    with patch("os.execv") as mock_exec:
        ret = run_with_config(argv, config)
        assert mock_exec.call_args_list == []
        assert ret == 1


def assert_error(message: str, stderr: str) -> None:
    assert stderr == f"bwrap-restricted: Error: {message}\n"


def write_args_to_file(file: IO[bytes], args: str) -> None:
    arg_list = args.split()
    for arg in arg_list:
        file.write(bytes(arg, encoding="utf8") + b"\0")
    file.seek(0)


def test_basic():
    assert_valid_options_end_at_index("--ro-bind /usr /usr -- foo bar", 3)


def test_no_double_dash():
    assert_valid_options_end_at_index("--ro-bind /usr /usr foo bar", 3)


def test_empty():
    assert_valid_options_end_at_index("", 0)


def test_forbidden_opt(capsys):
    assert_rejected("--ro-bind /usr /usr --cap-add CAP -- foo bar")
    assert_error("Forbidden bwrap option: --cap-add", capsys.readouterr().err)


def test_arg_named_like_forbidden():
    assert_valid_options_end_at_index("--ro-bind /usr --cap-add CAP -- foo bar", 3)


def test_unknown_option(capsys):
    assert_rejected("--ro-bind foo bar --uwu -- foo")
    assert_error("Unknown bwrap option: --uwu", capsys.readouterr().err)


def test_forbidden_opt_in_command_args():
    assert_valid_options_end_at_index("--ro-bind /usr /usr -- --cap-add CAP_SYS_ADMIN", 3)


def test_too_few_args_after_option(capsys):
    assert_rejected("--overlay one two")
    assert_error(
        "Too few arguments after option --overlay (needed 3, found 2)", capsys.readouterr().err
    )


def test_args_fd_valid():
    with TemporaryFile() as temp:
        write_args_to_file(temp, "--ro-bind /usr /usr --unshare-all")
        args = f"--args {temp.fileno()} foo bar"
        assert_valid_options_end_at_index(args, 2)


def test_args_fd_reject_forbidden(capsys):
    with TemporaryFile() as temp:
        write_args_to_file(temp, "--ro-bind /usr /usr --cap-add ALL")
        fd = temp.fileno()
        assert_rejected(f"--args {fd} -- foo bar")
    assert_error(f"Forbidden bwrap option: --cap-add (in --args {fd})", capsys.readouterr().err)


def test_too_few_args_after_option_in_fd(capsys):
    with TemporaryFile() as temp:
        write_args_to_file(temp, "--overlay one two")
        fd = temp.fileno()
        assert_rejected(f"--args {fd} --unshare-all -- foo")
    assert_error(
        f"Too few arguments after option --overlay (needed 3, found 2) (in --args {fd})",
        capsys.readouterr().err,
    )


def test_reject_fd_in_fd(capsys):
    with TemporaryFile() as temp1, TemporaryFile() as temp2:
        write_args_to_file(temp1, f"--args {temp2.fileno()}")
        write_args_to_file(temp2, "--unshare-net")
        fd1 = temp1.fileno()
        assert_rejected(f"--args {fd1} -- foo bar")
    assert_error(
        f"--args option not accepted inside file descriptor (in --args {fd1})",
        capsys.readouterr().err,
    )


def test_reject_invalid_args_fd(capsys):
    assert_rejected("--args 9001 -- foo bar")
    assert_error("Bad file descriptor: --args 9001", capsys.readouterr().err)


def test_reject_cmd_args_in_fd(capsys):
    with TemporaryFile() as temp:
        write_args_to_file(temp, "--ro-bind /usr /usr extra")
        fd = temp.fileno()
        assert_rejected(f"--args {fd} -- foo bar")
    assert_error(
        f"Non-option arguments in file descriptor (in --args {fd})", capsys.readouterr().err
    )
