// am-sandbox —— 给一个命令套上 **Landlock** 内核边界，然后 exec 它。
//
// # 为什么需要它
//
// 「工作区档」在 pi 这一路上一直是名不副实的：档位表写的是「本目录内可动、越界要问人」，
// 而实现只能按**工具名**判（bash/write/edit 一律问人）—— 因为命令的影响范围无法从文本
// 静态判定（`cd /工作区 && rm -rf /opt/x` 以"进工作区"开头）。pi 自己**故意**不内置沙箱
// （docs/security.md §No Built-in Sandbox：进程内的部分沙箱会被误解成安全边界，
// "Real isolation needs to come from the operating system or a container boundary"）。
// dsh 的做法是把沙箱放进自己的运行时；这个二进制就是补上 pi 这一侧的宿主部分。
//
// # 语义（与档位表对齐）
//
//   - **写**：只有 `--rw` 列出的目录/文件可写（及其子树），其余一律 EACCES —— 内核判，
//     不是提示词里"请你遵守"。
//   - **读与执行**：不限制（Landlock 的规则是白名单式的加法，无法表达"除了 X 都能读"；
//     真要把凭证也挡在外面需要容器/挂载命名空间那一层，见文件末尾的"已知边界"）。
//
// # 为什么不用 systemd-run / bwrap
//
// systemd-run 需要 dbus 与 root，且每个 worker 会变成一条 transient scope；
// bwrap 依赖 user namespace（本机可用，但多一层依赖）。Landlock 是**进程内**的：
// 不需要额外权限、不需要守护进程，`restrict_self` 之后本进程及其所有子孙都出不去。
//
// # 用法
//
//	am-sandbox --rw /home/program/x [--rw /tmp] [--rw-file /path/f] -- cmd args...
//	am-sandbox --print-abi                     # 打印内核支持的 Landlock ABI 版本
//
// 退出码：0 = 命令正常退出；126 = 边界套不上（**fail closed**，不降级为"裸跑"）；
// 其余为命令自身的退出码。
//
// # 已知边界（如实写下来，不假装没有）
//
//   - ABI < 3 的内核上 `truncate(2)`/`ftruncate(2)` 不受 Landlock 管辖（本机 ABI=2），
//     即"把界外某个已有文件截断为 0 字节"仍可能成功。`open(O_TRUNC)` 不受影响。
//     --print-abi 就是给排查这个用的。
//   - 无法限制**读**与网络出向（要那两样得上容器/网络命名空间）。
//   - 硬链接：在同一受控目录里对界外文件建硬链（需要 REFER，通常会被挡），
//     但若界内**已经存在**一个指向界外文件的硬链接，经它写入走的是界内路径。
//   - `PR_SET_NO_NEW_PRIVS` 是 Landlock 的前提，设了之后 setuid 提权失效
//     （`sudo` 这类在沙箱内不再能提权）—— 对 agent 来说这是想要的方向。
package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"unsafe"

	"golang.org/x/sys/unix"
)

// 我们**要管**的权限：只挑"会改变文件系统状态"的那些。
//
// 刻意不管 READ_FILE / READ_DIR / EXECUTE：管了就等于"只允许白名单目录可读"，
// 而 Landlock 没有否定式规则，会把整个系统读权限一起收掉（agent 连 /usr/include
// 都读不了）。写权限才是"越界"的实害面。
const (
	fsWrite  = unix.LANDLOCK_ACCESS_FS_WRITE_FILE
	fsRemove = unix.LANDLOCK_ACCESS_FS_REMOVE_FILE | unix.LANDLOCK_ACCESS_FS_REMOVE_DIR
	fsMake   = unix.LANDLOCK_ACCESS_FS_MAKE_CHAR | unix.LANDLOCK_ACCESS_FS_MAKE_DIR |
		unix.LANDLOCK_ACCESS_FS_MAKE_REG | unix.LANDLOCK_ACCESS_FS_MAKE_SOCK |
		unix.LANDLOCK_ACCESS_FS_MAKE_FIFO | unix.LANDLOCK_ACCESS_FS_MAKE_BLOCK |
		unix.LANDLOCK_ACCESS_FS_MAKE_SYM
	fsRefer    = unix.LANDLOCK_ACCESS_FS_REFER    // ABI 2+
	fsTruncate = unix.LANDLOCK_ACCESS_FS_TRUNCATE // ABI 3+
)

// handleFor 按内核 ABI 取"这一版内核认得"的写权限集合。
//
// 关键：**不能把内核不认识的位塞进 handled_access_fs** —— 那会让 create_ruleset
// 直接 EINVAL（整条边界都套不上）。所以按 ABI 逐位裁剪。
func handleFor(abi int) uint64 {
	// 这些位 ABI 1 就有。
	h := uint64(fsWrite | fsRemove | fsMake)
	if abi >= 2 {
		h |= fsRefer
	}
	if abi >= 3 {
		h |= fsTruncate
	}
	return h
}

// fileHandleFor 是"规则目标是**单个文件**"时能用的写权限子集。
//
// ★ 这里有个坑值得记下来：把目录类权限（MAKE_*/REMOVE_*/REFER）用在文件目标上，
// `landlock_add_rule` 直接返回 **EINVAL**，而错误信息只说"invalid argument"，
// 看不出是权限位不匹配。实测为了放行 `cmd 2>/dev/null` 给 /dev/null 加规则时才撞上。
func fileHandleFor(abi int) uint64 {
	h := uint64(fsWrite) // WRITE_FILE：文件目标真正有意义的那一位
	if abi >= 3 {
		h |= fsTruncate
	}
	// 刻意**不**带 IOCTL_DEV（ABI 4 才有）：我们没把它列进 handled_access_fs
	// —— 规则里的 allowed_access 必须是 handled 的子集，否则内核直接 EINVAL。
	// 单元测试里那条"文件权限必须是 handled 的子集"就是拦这个的。
	return h
}

func abiVersion() (int, error) {
	r, _, errno := unix.Syscall(unix.SYS_LANDLOCK_CREATE_RULESET, 0, 0,
		unix.LANDLOCK_CREATE_RULESET_VERSION)
	if errno != 0 {
		return 0, fmt.Errorf("landlock_create_ruleset(VERSION) 失败: %w", errno)
	}
	return int(r), nil
}

type opt struct {
	rwDirs  []string
	rwFiles []string
	cmd     []string
}

func usage() {
	fmt.Fprint(os.Stderr, `am-sandbox —— 给命令套 Landlock 内核边界后 exec 它

用法：
  am-sandbox --rw <目录> [--rw <目录>]... [--rw-file <文件>]... -- <命令> [参数...]
  am-sandbox --print-abi

语义：--rw 列出的目录（含子树）可写，其余路径写操作一律 EACCES；读与执行不限制。
套不上边界时**不降级**（退出码 126），绝不"静默裸跑"。
`)
}

func parseArgs(args []string) (opt, error) {
	var o opt
	i := 0
	for i < len(args) {
		a := args[i]
		switch {
		case a == "--":
			o.cmd = args[i+1:]
			return o, nil
		case a == "--rw":
			if i+1 >= len(args) {
				return o, errors.New("--rw 后面要给目录")
			}
			o.rwDirs = append(o.rwDirs, args[i+1])
			i += 2
		case a == "--rw-file":
			if i+1 >= len(args) {
				return o, errors.New("--rw-file 后面要给文件路径")
			}
			o.rwFiles = append(o.rwFiles, args[i+1])
			i += 2
		case strings.HasPrefix(a, "--rw="):
			o.rwDirs = append(o.rwDirs, strings.TrimPrefix(a, "--rw="))
			i++
		case strings.HasPrefix(a, "--rw-file="):
			o.rwFiles = append(o.rwFiles, strings.TrimPrefix(a, "--rw-file="))
			i++
		case a == "-h" || a == "--help":
			usage()
			os.Exit(0)
		default:
			return o, fmt.Errorf("看不懂的参数 %q（命令前要有 --）", a)
		}
	}
	return o, errors.New("缺少 `-- <命令>`")
}

// restrictSelf 把当前进程关进 Landlock 域。返回 error 时**调用方必须放弃执行命令**。
func restrictSelf(rwDirs, rwFiles []string) (int, error) {
	abi, err := abiVersion()
	if err != nil {
		return 0, err
	}
	handled := handleFor(abi)

	attr := unix.LandlockRulesetAttr{Access_fs: handled}
	rfd, _, errno := unix.Syscall(unix.SYS_LANDLOCK_CREATE_RULESET,
		uintptr(unsafe.Pointer(&attr)), unsafe.Sizeof(attr), 0)
	if errno != 0 {
		return abi, fmt.Errorf("create_ruleset 失败: %w", errno)
	}
	defer unix.Close(int(rfd))

	// 目录：允许整个子树。文件：只允许这个文件本身（Landlock 的 path_beneath 对
	// 文件同样成立，且不含"父目录可写"）。
	addRule := func(path string, isDir bool) error {
		st, err := os.Stat(path)
		if err != nil {
			// 路径不存在就说清楚是哪一个：静默跳过 = 少一条规则而没人知道。
			return fmt.Errorf("--rw%s 指向的路径不存在: %s（%w）", map[bool]string{true: "", false: "-file"}[isDir], path, err)
		}
		if isDir && !st.IsDir() {
			return fmt.Errorf("--rw 要求目录，但 %s 不是目录", path)
		}
		if !isDir && st.IsDir() {
			return fmt.Errorf("--rw-file 要求文件，但 %s 是目录", path)
		}
		fd, err := unix.Open(path, unix.O_PATH|unix.O_CLOEXEC, 0)
		if err != nil {
			return fmt.Errorf("打开 %s 取句柄失败: %w", path, err)
		}
		defer unix.Close(fd)
		allowed := handled
		if !isDir {
			// 文件目标只能用文件类权限，否则 EINVAL（见 fileHandleFor 的注释）。
			// 由此也限定了语义：`--rw-file /dev/null` 只放行"写这个设备"，
			// **不放行**在同目录里新建/删除任何东西。
			allowed = fileHandleFor(abi)
		}
		pa := unix.LandlockPathBeneathAttr{Allowed_access: allowed, Parent_fd: int32(fd)}
		if _, _, e := unix.Syscall(unix.SYS_LANDLOCK_ADD_RULE, rfd,
			unix.LANDLOCK_RULE_PATH_BENEATH, uintptr(unsafe.Pointer(&pa))); e != 0 {
			return fmt.Errorf("add_rule(%s) 失败: %w", path, e)
		}
		return nil
	}
	for _, d := range rwDirs {
		if err := addRule(d, true); err != nil {
			return abi, err
		}
	}
	for _, f := range rwFiles {
		if err := addRule(f, false); err != nil {
			return abi, err
		}
	}

	// Landlock 要求先 no_new_privs（否则 restrict_self 直接 EPERM）。
	if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil {
		return abi, fmt.Errorf("PR_SET_NO_NEW_PRIVS 失败: %w", err)
	}
	if _, _, e := unix.Syscall(unix.SYS_LANDLOCK_RESTRICT_SELF, rfd, 0, 0); e != 0 {
		return abi, fmt.Errorf("restrict_self 失败: %w", e)
	}
	return abi, nil
}

func main() {
	args := os.Args[1:]
	if len(args) == 1 && args[0] == "--print-abi" {
		abi, err := abiVersion()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Landlock 不可用：%v\n", err)
			os.Exit(126)
		}
		fmt.Printf("landlock abi=%d handled_write_rights=0x%x\n", abi, handleFor(abi))
		return
	}

	o, err := parseArgs(args)
	if err != nil {
		fmt.Fprintf(os.Stderr, "am-sandbox: %v\n", err)
		usage()
		os.Exit(2)
	}

	// unix.Exec 是裸 execve，不查 PATH —— 先解析成绝对路径。
	// 解析放在关进边界**之前**：这一步不该受边界影响，出问题也要能给出清楚的报错。
	bin, err := exec.LookPath(o.cmd[0])
	if err != nil {
		fmt.Fprintf(os.Stderr, "am-sandbox: 找不到可执行文件 %q: %v\n", o.cmd[0], err)
		os.Exit(127)
	}
	o.cmd[0] = bin

	abi, err := restrictSelf(o.rwDirs, o.rwFiles)
	if err != nil {
		// ★ fail closed：套不上边界就**不执行命令**。
		// 悄悄裸跑会让"档位=workspace"变成一句谎话，而谎话比做不到更危险
		// （模型以为越界会被拦，于是自己不设防）。
		fmt.Fprintf(os.Stderr, "am-sandbox: 边界未建立，拒绝执行（abi=%d）：%v\n", abi, err)
		os.Exit(126)
	}

	if err := unix.Exec(o.cmd[0], o.cmd, os.Environ()); err != nil {
		fmt.Fprintf(os.Stderr, "am-sandbox: exec %q 失败: %v\n", o.cmd[0], err)
		os.Exit(127)
	}
}