#!/usr/bin/env python3
"""Apply a Landlock + seccomp policy, drop privileges, then exec one shell command."""

from __future__ import annotations

import argparse
import ctypes
import ctypes.util
import errno
import json
import os
import platform
import sys
from pathlib import Path


LANDLOCK_CREATE_RULESET_VERSION = 1
LANDLOCK_RULE_PATH_BENEATH = 1

ACCESS_FS_EXECUTE = 1 << 0
ACCESS_FS_WRITE_FILE = 1 << 1
ACCESS_FS_READ_FILE = 1 << 2
ACCESS_FS_READ_DIR = 1 << 3
ACCESS_FS_REMOVE_DIR = 1 << 4
ACCESS_FS_REMOVE_FILE = 1 << 5
ACCESS_FS_MAKE_CHAR = 1 << 6
ACCESS_FS_MAKE_DIR = 1 << 7
ACCESS_FS_MAKE_REG = 1 << 8
ACCESS_FS_MAKE_SOCK = 1 << 9
ACCESS_FS_MAKE_FIFO = 1 << 10
ACCESS_FS_MAKE_BLOCK = 1 << 11
ACCESS_FS_MAKE_SYM = 1 << 12
ACCESS_FS_REFER = 1 << 13
ACCESS_FS_TRUNCATE = 1 << 14

ACCESS_NET_CONNECT_TCP = 1 << 0
ACCESS_NET_BIND_TCP = 1 << 1

READ_ACCESS = ACCESS_FS_EXECUTE | ACCESS_FS_READ_FILE | ACCESS_FS_READ_DIR
WRITE_ACCESS = (
    ACCESS_FS_WRITE_FILE
    | ACCESS_FS_REMOVE_DIR
    | ACCESS_FS_REMOVE_FILE
    | ACCESS_FS_MAKE_CHAR
    | ACCESS_FS_MAKE_DIR
    | ACCESS_FS_MAKE_REG
    | ACCESS_FS_MAKE_SOCK
    | ACCESS_FS_MAKE_FIFO
    | ACCESS_FS_MAKE_BLOCK
    | ACCESS_FS_MAKE_SYM
    | ACCESS_FS_REFER
    | ACCESS_FS_TRUNCATE
)
HANDLED_FS_ACCESS = READ_ACCESS | WRITE_ACCESS

PR_SET_NO_NEW_PRIVS = 38
AF_INET = 2
AF_INET6 = 10
SCMP_ACT_ALLOW = 0x7FFF0000
SCMP_ACT_ERRNO = 0x00050000 | errno.EPERM
SCMP_CMP_EQ = 4


class RulesetAttr(ctypes.Structure):
    _fields_ = [
        ("handled_access_fs", ctypes.c_uint64),
        ("handled_access_net", ctypes.c_uint64),
    ]


class PathBeneathAttr(ctypes.Structure):
    _fields_ = [
        ("allowed_access", ctypes.c_uint64),
        ("parent_fd", ctypes.c_int32),
    ]


class ScmpArgCmp(ctypes.Structure):
    _fields_ = [
        ("arg", ctypes.c_uint32),
        ("op", ctypes.c_uint32),
        ("datum_a", ctypes.c_uint64),
        ("datum_b", ctypes.c_uint64),
    ]


def syscall_numbers() -> tuple[int, int, int]:
    machine = platform.machine().lower()
    if machine not in {"x86_64", "amd64", "aarch64", "arm64"}:
        raise RuntimeError(f"unsupported architecture for Landlock syscalls: {machine}")
    return 444, 445, 446


def checked_syscall(libc: ctypes.CDLL, number: int, *args: object) -> int:
    result = int(libc.syscall(number, *args))
    if result < 0:
        error_number = ctypes.get_errno()
        raise OSError(error_number, os.strerror(error_number))
    return result


def get_landlock_abi(libc: ctypes.CDLL) -> int:
    create_ruleset, _, _ = syscall_numbers()
    return checked_syscall(
        libc,
        create_ruleset,
        ctypes.c_void_p(),
        ctypes.c_size_t(0),
        ctypes.c_uint32(LANDLOCK_CREATE_RULESET_VERSION),
    )


def load_seccomp() -> ctypes.CDLL:
    library_name = ctypes.util.find_library("seccomp") or "libseccomp.so.2"
    library = ctypes.CDLL(library_name, use_errno=True)
    library.seccomp_init.argtypes = [ctypes.c_uint32]
    library.seccomp_init.restype = ctypes.c_void_p
    library.seccomp_release.argtypes = [ctypes.c_void_p]
    library.seccomp_syscall_resolve_name.argtypes = [ctypes.c_char_p]
    library.seccomp_syscall_resolve_name.restype = ctypes.c_int
    library.seccomp_rule_add_array.argtypes = [
        ctypes.c_void_p,
        ctypes.c_uint32,
        ctypes.c_int,
        ctypes.c_uint,
        ctypes.POINTER(ScmpArgCmp),
    ]
    library.seccomp_rule_add_array.restype = ctypes.c_int
    library.seccomp_load.argtypes = [ctypes.c_void_p]
    library.seccomp_load.restype = ctypes.c_int
    return library


def add_path_rule(
    libc: ctypes.CDLL,
    add_rule_number: int,
    ruleset_fd: int,
    path: str,
    access: int,
) -> None:
    resolved_path = os.path.realpath(path)
    if not os.path.exists(resolved_path):
        return
    path_fd = os.open(resolved_path, os.O_PATH | os.O_CLOEXEC)
    try:
        allowed_access = access
        if not Path(resolved_path).is_dir():
            allowed_access &= ACCESS_FS_EXECUTE | ACCESS_FS_READ_FILE | ACCESS_FS_WRITE_FILE | ACCESS_FS_TRUNCATE
        attribute = PathBeneathAttr(
            allowed_access=allowed_access,
            parent_fd=path_fd,
        )
        checked_syscall(
            libc,
            add_rule_number,
            ctypes.c_int(ruleset_fd),
            ctypes.c_int(LANDLOCK_RULE_PATH_BENEATH),
            ctypes.byref(attribute),
            ctypes.c_uint32(0),
        )
    finally:
        os.close(path_fd)


def build_landlock_ruleset(
    libc: ctypes.CDLL,
    abi: int,
    workspace: str,
    read_only_paths: list[str],
) -> int:
    create_ruleset, add_rule, _ = syscall_numbers()
    handled_network = (
        ACCESS_NET_CONNECT_TCP | ACCESS_NET_BIND_TCP if abi >= 4 else 0
    )
    ruleset_attribute = RulesetAttr(
        handled_access_fs=HANDLED_FS_ACCESS,
        handled_access_net=handled_network,
    )
    ruleset_fd = checked_syscall(
        libc,
        create_ruleset,
        ctypes.byref(ruleset_attribute),
        ctypes.sizeof(ruleset_attribute),
        ctypes.c_uint32(0),
    )
    try:
        add_path_rule(
            libc,
            add_rule,
            ruleset_fd,
            workspace,
            HANDLED_FS_ACCESS,
        )
        for path in read_only_paths:
            add_path_rule(libc, add_rule, ruleset_fd, path, READ_ACCESS)
        for device_path, access in (
            ("/dev/null", ACCESS_FS_READ_FILE | ACCESS_FS_WRITE_FILE),
            ("/dev/zero", ACCESS_FS_READ_FILE | ACCESS_FS_WRITE_FILE),
            ("/dev/urandom", ACCESS_FS_READ_FILE),
            ("/dev/random", ACCESS_FS_READ_FILE),
        ):
            add_path_rule(libc, add_rule, ruleset_fd, device_path, access)
    except Exception:
        os.close(ruleset_fd)
        raise
    return ruleset_fd


def install_seccomp_network_filter(library: ctypes.CDLL) -> None:
    context = library.seccomp_init(SCMP_ACT_ALLOW)
    if not context:
        raise RuntimeError("seccomp_init failed")
    try:
        socket_syscall = library.seccomp_syscall_resolve_name(b"socket")
        if socket_syscall < 0:
            raise RuntimeError("could not resolve socket syscall")
        for domain in (AF_INET, AF_INET6):
            comparison = ScmpArgCmp(
                arg=0,
                op=SCMP_CMP_EQ,
                datum_a=domain,
                datum_b=0,
            )
            result = library.seccomp_rule_add_array(
                context,
                SCMP_ACT_ERRNO,
                socket_syscall,
                1,
                ctypes.byref(comparison),
            )
            if result != 0:
                raise OSError(-result, os.strerror(-result))
        result = library.seccomp_load(context)
        if result != 0:
            raise OSError(-result, os.strerror(-result))
    finally:
        library.seccomp_release(context)


def set_no_new_privileges(libc: ctypes.CDLL) -> None:
    result = libc.prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0)
    if result != 0:
        error_number = ctypes.get_errno()
        raise OSError(error_number, os.strerror(error_number))


def restrict_process(
    workspace: str,
    read_only_paths: list[str],
    uid: int,
    gid: int,
) -> int:
    libc = ctypes.CDLL(None, use_errno=True)
    libc.syscall.restype = ctypes.c_long
    libc.prctl.restype = ctypes.c_int
    seccomp = load_seccomp()
    abi = get_landlock_abi(libc)
    if abi < 4:
        raise RuntimeError(f"Landlock ABI 4 or newer is required; detected ABI {abi}")

    ruleset_fd = build_landlock_ruleset(
        libc,
        abi,
        workspace,
        read_only_paths,
    )
    try:
        if os.geteuid() == 0:
            os.setgroups([])
            os.setgid(gid)
            os.setuid(uid)
        set_no_new_privileges(libc)
        _, _, restrict_self = syscall_numbers()
        checked_syscall(
            libc,
            restrict_self,
            ctypes.c_int(ruleset_fd),
            ctypes.c_uint32(0),
        )
        install_seccomp_network_filter(seccomp)
    finally:
        os.close(ruleset_fd)
    return abi


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--probe", action="store_true")
    parser.add_argument("--workspace")
    parser.add_argument("--read-only", action="append", default=[])
    parser.add_argument("--uid", type=int, default=10_001)
    parser.add_argument("--gid", type=int, default=10_001)
    parser.add_argument("--command")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    libc = ctypes.CDLL(None, use_errno=True)
    libc.syscall.restype = ctypes.c_long
    if args.probe:
        abi = get_landlock_abi(libc)
        load_seccomp()
        print(json.dumps({"ok": True, "landlock_abi": abi, "seccomp": True}))
        return 0 if abi >= 4 else 1

    if not args.workspace or args.command is None:
        raise RuntimeError("workspace and command are required")
    workspace = os.path.realpath(args.workspace)
    if not os.path.isdir(workspace):
        raise RuntimeError("workspace must be an existing directory")
    os.chdir(workspace)
    restrict_process(workspace, args.read_only, args.uid, args.gid)
    os.execve("/bin/bash", ["bash", "-c", args.command], dict(os.environ))
    return 127


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:
        print(
            json.dumps(
                {
                    "ok": False,
                    "error": "SANDBOX_UNAVAILABLE",
                    "message": str(error),
                }
            ),
            file=sys.stderr,
        )
        raise SystemExit(125)