// agent-sandbox runs Claude Code in a rootless podman sandbox whose only way
// out is a per-session router that resolves and forwards allowlisted hosts.
package main

import (
	"context"
	_ "embed"
	"fmt"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/containers/podman/v5/libpod/define"
	"github.com/containers/podman/v5/pkg/specgen"
	spec "github.com/opencontainers/runtime-spec/specs-go"
	"go.podman.io/common/libnetwork/types"
	"golang.org/x/term"
)

const (
	registry  = "artifactapi.k8s.syd1.au.unkin.net/docker-internal"
	egressNet = "agent-sandbox-egress"
)

var version = "dev"

//go:embed allowed-hosts
var defaultAllowlist []byte

func main() {
	// A closed stdout must end the session through cleanup, not kill us.
	signal.Ignore(syscall.SIGPIPE)
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
	// Every argument belongs to claude or --exec, as in the bash launcher.
	code := launch(ctx, os.Args[1:])
	stop()
	os.Exit(code)
}

func imageNames() (string, string) {
	img, router := registry+"/agent-sandbox:"+version, registry+"/agent-sandbox-router:"+version
	if version == "dev" {
		img, router = "localhost/agent-sandbox:dev", "localhost/agent-sandbox-router:dev"
	}
	return envOr("AGENT_SANDBOX_IMAGE", img), envOr("AGENT_SANDBOX_ROUTER_IMAGE", router)
}

func envOr(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func configDir() string {
	return filepath.Join(envOr("XDG_CONFIG_HOME", filepath.Join(os.Getenv("HOME"), ".config")), "agent-sandbox")
}

func allowlistPath() string {
	return envOr("AGENT_SANDBOX_ALLOWLIST", filepath.Join(configDir(), "allowed-hosts"))
}

// runtimeDir holds the session dirs, including the private API socket's.
func runtimeDir() string {
	return envOr("XDG_RUNTIME_DIR", "/tmp")
}

func warn(format string, a ...any) {
	fmt.Fprintf(os.Stderr, "agent-sandbox: "+format+"\n", a...)
}

// launch runs one session and returns the process exit code.
func launch(ctx context.Context, args []string) int {
	cmd := append([]string{"claude"}, args...)
	if len(args) > 0 && args[0] == "--exec" {
		cmd = args[1:]
	}
	pwd, err := os.Getwd()
	if err != nil {
		warn("%v", err)
		return 1
	}
	allowlist := allowlistPath()
	src, err := checkLaunchDir(pwd, allowlist)
	if err != nil {
		warn("%v", err)
		return 1
	}
	if strings.Contains(pwd+src, ":") {
		warn("cannot mount a directory containing ':'")
		return 1
	}
	if err := ensureAllowlist(allowlist); err != nil {
		warn("%v", err)
		return 1
	}
	entries, err := parseAllowlist(allowlist)
	if err != nil {
		warn("%v", err)
		return 1
	}

	session := fmt.Sprintf("agent-sandbox-%d-%d", time.Now().Unix(), os.Getpid())
	runtimeDir := runtimeDir()
	tmp, err := os.MkdirTemp(runtimeDir, session+".")
	if err != nil {
		warn("%v", err)
		return 1
	}
	defer func() { _ = os.RemoveAll(tmp) }()
	// The API socket lives apart from tmp, which is mounted into the router.
	apiDir, err := os.MkdirTemp(runtimeDir, session+"-api.")
	if err != nil {
		warn("%v", err)
		return 1
	}
	defer func() { _ = os.RemoveAll(apiDir) }()

	conn, eng, stopService, err := startService(apiDir)
	if err != nil {
		warn("%v", err)
		return 125
	}
	defer stopService()
	defer func() {
		_ = eng.PodRemove(conn, session)
		_ = eng.ContainerRemove(conn, session+"-router")
		_ = eng.NetworkRemove(conn, session)
	}()

	opCtx, cancel := context.WithCancel(conn)
	defer cancel()
	defer context.AfterFunc(ctx, cancel)()
	code := runSession(opCtx, eng, session, tmp, src, pwd, entries, cmd)
	if ctx.Err() != nil {
		return 130
	}
	return code
}

// fail reports a podman API error with podman's own exit code for it.
func fail(err error) int {
	warn("%v", err)
	return 125
}

// runSession mounts the checked src at pwd, the path the user launched from.
func runSession(ctx context.Context, eng engine, session, tmp, src, pwd string, entries []entry, cmd []string) int {
	image, routerImage := imageNames()
	homeVolume := envOr("AGENT_SANDBOX_HOME_VOLUME", "agent-sandbox-home")
	router := session + "-router"
	yes := true

	if _, err := eng.NetworkCreate(ctx, &types.Network{Name: egressNet, Driver: "bridge", DNSEnabled: true}, true); err != nil {
		return fail(err)
	}
	if err := eng.VolumeCreate(ctx, homeVolume); err != nil {
		return fail(err)
	}
	n, err := eng.NetworkCreate(ctx, &types.Network{Name: session, Driver: "bridge", Internal: true}, false)
	if err != nil {
		return fail(err)
	}
	subnet := ""
	for _, s := range n.Subnets {
		subnet += s.Subnet.String()
	}
	if !strings.HasSuffix(subnet, ".0/24") {
		warn("unexpected subnet %s", subnet)
		return 1
	}
	rip := strings.TrimSuffix(subnet, ".0/24") + ".2"

	dnsmasq, nginx, forwards := routerConfig(entries, rip)
	for name, body := range map[string]string{"dnsmasq.conf": dnsmasq, "nginx.conf": nginx} {
		f := filepath.Join(tmp, name)
		err := os.WriteFile(f, []byte(body), 0o644)
		if err == nil {
			err = os.Chmod(f, 0o644)
		}
		if err != nil {
			warn("%v", err)
			return 1
		}
	}
	if err := os.Chmod(tmp, 0o755); err != nil {
		warn("%v", err)
		return 1
	}

	r := newSpec(routerImage, router)
	r.Remove = &yes
	r.NetNS = specgen.Namespace{NSMode: specgen.Bridge}
	r.Networks = map[string]types.PerNetworkOptions{
		session:   {StaticIPs: []net.IP{net.ParseIP(rip)}},
		egressNet: {},
	}
	r.Mounts = []spec.Mount{{Type: "bind", Source: tmp, Destination: "/etc/router", Options: []string{"ro", "Z"}}}
	if err := eng.ContainerCreate(ctx, r); err != nil {
		return fail(err)
	}
	if err := eng.ContainerStart(ctx, router); err != nil {
		return fail(err)
	}
	if err := waitPort(ctx, eng, router, rip, "443"); err != nil {
		warn("%v", err)
		return 1
	}

	// The pod owns the netns, DNS and hosts file; claude only joins it.
	userns, _ := specgen.ParseUserNamespace("keep-id:uid=1000,gid=1000")
	p := &specgen.PodSpecGenerator{}
	p.Name = session
	p.Hostname = "agent-sandbox"
	p.NetNS = specgen.Namespace{NSMode: specgen.Bridge}
	p.Networks = map[string]types.PerNetworkOptions{session: {}}
	p.DNSServer = []net.IP{net.ParseIP(rip)}
	p.NoManageHosts = true
	p.Userns = userns
	if err := eng.PodCreate(ctx, p); err != nil {
		return fail(err)
	}

	// Loopback forwarders live in their own container: claude can't see, signal
	// or replace them. Claude runs unprivileged; podman's default bounding set
	// (no NET_ADMIN/NET_RAW/SYS_ADMIN) stays only so nested containers work.
	f := newSpec(routerImage, session+"-forwarder")
	f.Pod = session
	f.User = "0:0"
	f.ReadOnlyFilesystem = &yes
	f.ReadWriteTmpfs = &yes
	f.CapDrop = []string{"all"}
	f.CapAdd = []string{"NET_BIND_SERVICE"}
	f.Entrypoint = []string{"/usr/local/bin/forwarder"}
	f.Env = map[string]string{"FORWARDS": forwards}
	if err := eng.ContainerCreate(ctx, f); err != nil {
		return fail(err)
	}
	if err := eng.ContainerStart(ctx, f.Name); err != nil {
		return fail(err)
	}
	if err := waitPort(ctx, eng, f.Name, "127.0.0.1", "443"); err != nil {
		warn("%v", err)
		return 1
	}

	c := newSpec(image, session+"-claude")
	c.Remove = &yes
	c.Stdin = &yes
	c.Terminal = ptr(term.IsTerminal(int(os.Stdin.Fd())) && term.IsTerminal(int(os.Stdout.Fd())))
	c.Init = &yes
	c.Pod = session
	c.CapDrop = []string{"NET_ADMIN", "NET_RAW", "SYS_ADMIN"}
	c.SelinuxOpts = []string{"disable"}
	c.Annotations = map[string]string{define.InspectAnnotationLabel: "disable"}
	c.Devices = []spec.LinuxDevice{{Path: "/dev/fuse"}, {Path: "/dev/net/tun"}}
	c.Volumes = []*specgen.NamedVolume{{Name: homeVolume, Dest: "/home/claude"}}
	c.Mounts = []spec.Mount{{Type: "bind", Source: src, Destination: pwd}}
	c.WorkDir = pwd
	c.Env = map[string]string{"DISABLE_AUTOUPDATER": "1", "CLAUDE_CODE_TMPDIR": "/var/tmp"}
	c.Command = cmd
	if err := eng.ContainerCreate(ctx, c); err != nil {
		return fail(err)
	}
	code, err := eng.ContainerRun(ctx, c.Name)
	if err != nil {
		warn("%v", err)
	}
	return code
}

func ptr[T any](v T) *T { return &v }

// newSpec mirrors `podman run` defaults the API leaves to the client.
func newSpec(image, name string) *specgen.SpecGenerator {
	s := specgen.NewSpecGenerator(image, false)
	s.RawImageName = image
	s.Name = name
	s.HTTPProxy = ptr(true)
	s.SdNotifyMode = define.SdNotifyModeContainer
	return s
}

func waitPort(ctx context.Context, eng engine, container, ip, port string) error {
	for range 50 {
		if code, err := eng.ContainerExec(ctx, container, []string{"nc", "-z", ip, port}); err == nil && code == 0 {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
	return fmt.Errorf("%s did not start: %s", container, eng.ContainerLogs(context.WithoutCancel(ctx), container, 5))
}

// ensureAllowlist seeds a missing allowlist from the packaged default.
func ensureAllowlist(path string) error {
	if fi, err := os.Stat(path); err == nil && fi.Mode().IsRegular() {
		return nil
	}
	src := defaultAllowlist
	if exe, err := os.Executable(); err == nil {
		for _, p := range []string{filepath.Join(filepath.Dir(exe), "allowed-hosts"), "/usr/share/agent-sandbox/allowed-hosts"} {
			if b, err := os.ReadFile(p); err == nil {
				src = b
				break
			}
		}
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(path, src, 0o644); err != nil {
		return err
	}
	warn("created %s from defaults", path)
	return nil
}

var forbiddenRoots = []string{"/run", "/tmp", "/var/tmp", "/proc", "/sys", "/dev"}

func allowedPathsFile() string {
	return filepath.Join(configDir(), "allowed-paths")
}

// sensitiveDirs are host dirs an allowed-paths root may neither contain nor sit under.
func sensitiveDirs(home string) []string {
	dirs := []string{home + "/.ssh", home + "/.gnupg", home + "/.config", filepath.Dir(configDir()),
		home + "/.local/share/containers", runtimeDir()}
	if v := os.Getenv("XDG_DATA_HOME"); v != "" {
		dirs = append(dirs, v+"/containers")
	}
	return dirs
}

// overlapsSensitive returns the sensitive dir that p is, sits under or contains.
func overlapsSensitive(p, home string) string {
	for _, s := range sensitiveDirs(home) {
		if s = realpath(s); inside(p, s) || inside(s, p) {
			return s
		}
	}
	return ""
}

// allowedRoots reads the launch-dir roots; without the file the only root is
// $HOME and def is true.
func allowedRoots(home string) (roots []string, def bool, err error) {
	b, err := os.ReadFile(allowedPathsFile())
	if os.IsNotExist(err) {
		return []string{home}, true, nil
	}
	if err != nil {
		return nil, false, err
	}
	for _, l := range strings.Split(string(b), "\n") {
		l = strings.TrimSpace(l)
		if l == "" || strings.HasPrefix(l, "#") {
			continue
		}
		r := l
		if l == "~" || strings.HasPrefix(l, "~/") {
			r = home + l[1:]
		}
		if !filepath.IsAbs(r) {
			return nil, false, fmt.Errorf("%s: %q is not an absolute path", allowedPathsFile(), l)
		}
		r = realpath(r)
		if inside(home, r) {
			return nil, false, fmt.Errorf("%s: %q is or contains $HOME", allowedPathsFile(), l)
		}
		if filepath.Dir(r) == home && strings.HasPrefix(filepath.Base(r), ".") {
			return nil, false, fmt.Errorf("%s: %q is a dot-directory of $HOME", allowedPathsFile(), l)
		}
		for _, bad := range forbiddenRoots {
			if bad = realpath(bad); inside(r, bad) || inside(bad, r) {
				return nil, false, fmt.Errorf("%s: %q overlaps %s", allowedPathsFile(), l, bad)
			}
		}
		for _, d := range pathDirs() {
			if inside(d, r) {
				return nil, false, fmt.Errorf("%s: %q is or contains $PATH entry %s", allowedPathsFile(), l, d)
			}
		}
		if s := overlapsSensitive(r, home); s != "" {
			return nil, false, fmt.Errorf("%s: %q overlaps %s", allowedPathsFile(), l, s)
		}
		roots = append(roots, r)
	}
	return roots, false, nil
}

// pathDirs resolves $PATH; a sandbox-written file there would run on the host.
func pathDirs() []string {
	var dirs []string
	for _, d := range filepath.SplitList(os.Getenv("PATH")) {
		dirs = append(dirs, realpath(d))
	}
	return dirs
}

// checkLaunchDir returns pwd resolved if it is strictly under an allowed root,
// never $HOME itself, anything holding the sandbox's config, the API socket or
// a $PATH entry, nor anything overlapping a forbidden system or sensitive dir. Under the
// default $HOME root, dot-dirs of $HOME are refused too.
func checkLaunchDir(pwd, allowlist string) (string, error) {
	pwd = realpath(pwd)
	home := realpath(os.Getenv("HOME"))
	roots, def, err := allowedRoots(home)
	if err != nil {
		return "", err
	}
	if inside(home, pwd) {
		return "", fmt.Errorf("refusing to mount %s: it is or contains $HOME", pwd)
	}
	if rel, ok := strings.CutPrefix(pwd, home+"/"); def && ok && strings.HasPrefix(rel, ".") {
		return "", fmt.Errorf("refusing to mount %s: it is under a dot-directory of $HOME", pwd)
	}
	paths := allowedPathsFile()
	for _, dir := range append([]string{configDir(), filepath.Dir(allowlist), filepath.Dir(realpath(allowlist)),
		filepath.Dir(realpath(paths)), runtimeDir()}, pathDirs()...) {
		if dir = realpath(dir); inside(dir, pwd) {
			return "", fmt.Errorf("refusing to mount %s: it contains %s", pwd, dir)
		}
	}
	for _, bad := range forbiddenRoots {
		if bad = realpath(bad); inside(pwd, bad) || inside(bad, pwd) {
			return "", fmt.Errorf("refusing to mount %s: it overlaps %s", pwd, bad)
		}
	}
	if s := overlapsSensitive(pwd, home); s != "" {
		return "", fmt.Errorf("refusing to mount %s: it overlaps %s", pwd, s)
	}
	for _, r := range roots {
		if pwd != r && inside(pwd, r) {
			return pwd, nil
		}
	}
	return "", fmt.Errorf("refusing to mount %s: it is not under an allowed path (%s)", pwd, paths)
}

// inside reports whether path p is dir or lies under it.
func inside(p, dir string) bool {
	return p == dir || strings.HasPrefix(p, strings.TrimSuffix(dir, "/")+"/")
}

// realpath resolves every symlink in p like `realpath -m`: missing components
// are allowed, and ".." applies to the resolved path.
func realpath(p string) string {
	if !filepath.IsAbs(p) {
		wd, _ := os.Getwd()
		p = filepath.Join(wd, p)
	}
	rest := strings.Split(p, "/")
	resolved := "/"
	for hops := 0; len(rest) > 0; {
		part := rest[0]
		rest = rest[1:]
		switch part {
		case "", ".":
			continue
		case "..":
			resolved = filepath.Dir(resolved)
			continue
		}
		next := filepath.Join(resolved, part)
		fi, err := os.Lstat(next)
		if err != nil || fi.Mode()&os.ModeSymlink == 0 || hops > 255 {
			resolved = next
			continue
		}
		target, err := os.Readlink(next)
		if err != nil {
			resolved = next
			continue
		}
		hops++
		if filepath.IsAbs(target) {
			resolved = "/"
		}
		rest = append(strings.Split(target, "/"), rest...)
	}
	return resolved
}