//go:build linux

package agentworkspace

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"sync"
	"time"

	"golang.org/x/sys/unix"

	"iop/packages/go/agenttask"
)

const (
	linuxConfinementHelperArg  = "__iop_agentworkspace_mountns_v2__"
	linuxMetadataProbeTarget   = "__iop_agentworkspace_metadata_probe_v1__"
	linuxMetadataProbeXattrKey = "user.iop_agentworkspace_probe"
)

var (
	linuxConfinementOnce     sync.Once
	linuxConfinementIdentity string
	linuxConfinementErr      error
)

func init() {
	if len(os.Args) < 4 || os.Args[1] != linuxConfinementHelperArg {
		return
	}
	if err := runLinuxConfinementHelper(os.Args[2], os.Args[3], os.Args[4:]); err != nil {
		_, _ = fmt.Fprintln(os.Stderr, err)
		os.Exit(126)
	}
	os.Exit(0)
}

func platformConfinementRevision() (string, error) {
	linuxConfinementOnce.Do(func() {
		if err := probeLinuxConfinement(); err != nil {
			linuxConfinementErr = fmt.Errorf(
				"%w: metadata-complete mount namespace policy: %v",
				ErrConfinementUnavailable,
				err,
			)
			return
		}
		linuxConfinementIdentity = "linux-user-mount-namespace-metadata-policy-v2"
	})
	return linuxConfinementIdentity, linuxConfinementErr
}

func probeLinuxConfinement() error {
	target, err := exec.LookPath("touch")
	if err != nil {
		return err
	}
	target, err = filepath.Abs(target)
	if err != nil {
		return err
	}
	probeRoot, err := os.MkdirTemp("", "iop-confinement-probe-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(probeRoot)
	probeRoot, err = filepath.EvalSymlinks(probeRoot)
	if err != nil {
		return err
	}
	writableRoot := filepath.Join(probeRoot, "writable")
	protectedRoot := filepath.Join(probeRoot, "protected")
	for _, root := range []string{writableRoot, protectedRoot} {
		if err := os.Mkdir(root, 0o700); err != nil {
			return err
		}
	}
	probe := confinementPolicy{
		SchemaVersion: confinementPolicySchemaVersion,
		Revision:      "mountns-metadata-probe",
		Backend:       "mountns",
		WritableRoots: []string{writableRoot},
	}
	allowedPath := filepath.Join(writableRoot, "allowed")
	command, err := linuxHelperCommand(
		context.Background(),
		probe,
		target,
		[]string{allowedPath},
	)
	if err != nil {
		return err
	}
	if output, err := command.CombinedOutput(); err != nil {
		return fmt.Errorf(
			"%s writable-root probe failed: %v: %s",
			"mount namespace",
			err,
			boundedText(output),
		)
	}
	if _, err := os.Stat(allowedPath); err != nil {
		return fmt.Errorf("mount namespace writable-root probe did not create output: %w", err)
	}

	protectedPath := filepath.Join(protectedRoot, "denied")
	command, err = linuxHelperCommand(
		context.Background(),
		probe,
		target,
		[]string{protectedPath},
	)
	if err != nil {
		return err
	}
	if output, err := command.CombinedOutput(); err == nil {
		return errors.New("mount namespace protected-root probe unexpectedly wrote output")
	} else if _, statErr := os.Stat(protectedPath); !os.IsNotExist(statErr) {
		return fmt.Errorf(
			"mount namespace protected-root probe changed output after %v: %s",
			err,
			boundedText(output),
		)
	}

	metadataPath := filepath.Join(protectedRoot, "metadata")
	if err := os.WriteFile(metadataPath, []byte("metadata\n"), 0o600); err != nil {
		return err
	}
	if err := unix.Setxattr(
		metadataPath,
		linuxMetadataProbeXattrKey,
		[]byte("supported"),
		0,
	); err != nil {
		return fmt.Errorf("probe filesystem xattr support: %w", err)
	}
	if err := unix.Removexattr(metadataPath, linuxMetadataProbeXattrKey); err != nil {
		return fmt.Errorf("clear filesystem xattr probe: %w", err)
	}
	before, err := captureLinuxMetadata(metadataPath)
	if err != nil {
		return err
	}
	command, err = linuxHelperCommand(
		context.Background(),
		probe,
		linuxMetadataProbeTarget,
		[]string{metadataPath},
	)
	if err != nil {
		return err
	}
	if output, err := command.CombinedOutput(); err != nil {
		return fmt.Errorf(
			"mount namespace protected metadata probe failed: %v: %s",
			err,
			boundedText(output),
		)
	}
	after, err := captureLinuxMetadata(metadataPath)
	if err != nil {
		return err
	}
	if !reflect.DeepEqual(before, after) {
		return errors.New("mount namespace protected metadata probe changed the protected file")
	}
	return nil
}

func platformConfinementCommand(
	ctx context.Context,
	_ agenttask.ConfinementBinding,
	policy confinementPolicy,
	name string,
	args []string,
) (*exec.Cmd, error) {
	if _, err := platformConfinementRevision(); err != nil {
		return nil, err
	}
	policy.Backend = "mountns"
	target, err := exec.LookPath(name)
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: resolve confined command: %w", err)
	}
	target, err = filepath.Abs(target)
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: resolve confined command path: %w", err)
	}
	return linuxHelperCommand(ctx, policy, target, args)
}

func linuxHelperCommand(
	ctx context.Context,
	policy confinementPolicy,
	target string,
	args []string,
) (*exec.Cmd, error) {
	encoded, err := json.Marshal(policy)
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: encode mount namespace policy: %w", err)
	}
	executable, err := os.Executable()
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: resolve confinement helper: %w", err)
	}
	executable, err = filepath.EvalSymlinks(executable)
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: canonicalize confinement helper: %w", err)
	}
	helperArgs := []string{
		linuxConfinementHelperArg,
		base64.RawURLEncoding.EncodeToString(encoded),
		target,
	}
	helperArgs = append(helperArgs, args...)
	unshare, err := exec.LookPath("unshare")
	if err != nil {
		return nil, fmt.Errorf(
			"%w: locate unshare helper: %v",
			ErrConfinementUnavailable,
			err,
		)
	}
	unshareArgs := []string{
		"--user",
		"--map-root-user",
		"--mount",
		executable,
	}
	unshareArgs = append(unshareArgs, helperArgs...)
	return exec.CommandContext(ctx, unshare, unshareArgs...), nil
}

func runLinuxConfinementHelper(
	encodedPolicy string,
	target string,
	args []string,
) error {
	encoded, err := base64.RawURLEncoding.DecodeString(encodedPolicy)
	if err != nil {
		return fmt.Errorf("agentworkspace: decode mount namespace policy: %w", err)
	}
	var policy confinementPolicy
	if err := json.Unmarshal(encoded, &policy); err != nil {
		return fmt.Errorf("agentworkspace: parse mount namespace policy: %w", err)
	}
	if policy.SchemaVersion != confinementPolicySchemaVersion ||
		policy.Revision == "" ||
		policy.Backend != "mountns" {
		return fmt.Errorf("agentworkspace: invalid mount namespace policy identity")
	}
	err = installMountNamespacePolicy(policy.WritableRoots)
	if err != nil {
		return fmt.Errorf(
			"agentworkspace: install %s filesystem policy: %w",
			policy.Backend,
			err,
		)
	}
	if target == linuxMetadataProbeTarget {
		return runLinuxProtectedMetadataProbe(args)
	}
	argv := append([]string{target}, args...)
	if err := unix.Exec(target, argv, os.Environ()); err != nil {
		return fmt.Errorf("agentworkspace: execute confined child: %w", err)
	}
	return nil
}

type linuxMetadataSnapshot struct {
	Mode         os.FileMode
	ModTime      time.Time
	UID          uint32
	GID          uint32
	Xattr        []byte
	XattrPresent bool
}

func captureLinuxMetadata(path string) (linuxMetadataSnapshot, error) {
	info, err := os.Stat(path)
	if err != nil {
		return linuxMetadataSnapshot{}, err
	}
	var stat unix.Stat_t
	if err := unix.Lstat(path, &stat); err != nil {
		return linuxMetadataSnapshot{}, err
	}
	xattr, present, err := readLinuxProbeXattr(path)
	if err != nil {
		return linuxMetadataSnapshot{}, err
	}
	return linuxMetadataSnapshot{
		Mode:         info.Mode(),
		ModTime:      info.ModTime(),
		UID:          stat.Uid,
		GID:          stat.Gid,
		Xattr:        xattr,
		XattrPresent: present,
	}, nil
}

func readLinuxProbeXattr(path string) ([]byte, bool, error) {
	size, err := unix.Getxattr(path, linuxMetadataProbeXattrKey, nil)
	if errors.Is(err, unix.ENODATA) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	value := make([]byte, size)
	if _, err := unix.Getxattr(path, linuxMetadataProbeXattrKey, value); err != nil {
		return nil, false, err
	}
	return value, true, nil
}

func runLinuxProtectedMetadataProbe(args []string) error {
	if len(args) != 1 {
		return fmt.Errorf("agentworkspace: metadata probe expected one protected path")
	}
	path := args[0]
	before, err := captureLinuxMetadata(path)
	if err != nil {
		return fmt.Errorf("agentworkspace: capture protected metadata: %w", err)
	}
	attempts := []struct {
		name string
		fn   func() error
	}{
		{name: "chmod", fn: func() error { return os.Chmod(path, 0o777) }},
		{name: "utime", fn: func() error {
			return os.Chtimes(path, time.Unix(123456789, 0), time.Unix(123456789, 0))
		}},
		{name: "chown", fn: func() error { return os.Chown(path, 0, 0) }},
		{name: "setxattr", fn: func() error {
			return unix.Setxattr(path, linuxMetadataProbeXattrKey, []byte("denied"), 0)
		}},
	}
	for _, attempt := range attempts {
		if err := attempt.fn(); err == nil {
			return fmt.Errorf("agentworkspace: protected %s unexpectedly succeeded", attempt.name)
		} else if !linuxMetadataMutationDenied(err) {
			return fmt.Errorf("agentworkspace: protected %s returned %w", attempt.name, err)
		}
	}
	after, err := captureLinuxMetadata(path)
	if err != nil {
		return fmt.Errorf("agentworkspace: recapture protected metadata: %w", err)
	}
	if !reflect.DeepEqual(before, after) {
		return errors.New("agentworkspace: protected metadata changed despite confinement")
	}
	return nil
}

func linuxMetadataMutationDenied(err error) bool {
	return errors.Is(err, unix.EACCES) ||
		errors.Is(err, unix.EPERM) ||
		errors.Is(err, unix.EROFS)
}

func installMountNamespacePolicy(writableRoots []string) error {
	if err := unix.Mount("", "/", "", unix.MS_REC|unix.MS_PRIVATE, ""); err != nil {
		return fmt.Errorf("make mounts private: %w", err)
	}
	for _, root := range writableRoots {
		if err := unix.Mount(root, root, "", unix.MS_BIND|unix.MS_REC, ""); err != nil {
			return fmt.Errorf("bind writable root: %w", err)
		}
	}
	if err := unix.MountSetattr(
		unix.AT_FDCWD,
		"/",
		unix.AT_RECURSIVE,
		&unix.MountAttr{Attr_set: unix.MOUNT_ATTR_RDONLY},
	); err != nil {
		return fmt.Errorf("make mount tree read-only: %w", err)
	}
	for _, root := range writableRoots {
		if err := unix.MountSetattr(
			unix.AT_FDCWD,
			root,
			unix.AT_RECURSIVE,
			&unix.MountAttr{Attr_clr: unix.MOUNT_ATTR_RDONLY},
		); err != nil {
			return fmt.Errorf("restore writable root: %w", err)
		}
	}
	return dropNamespaceMountCapabilities()
}

func dropNamespaceMountCapabilities() error {
	const (
		secureNoRoot               = 1 << 0
		secureNoRootLocked         = 1 << 1
		secureNoSetuidFixup        = 1 << 2
		secureNoSetuidFixupLocked  = 1 << 3
		secureNoAmbientRaise       = 1 << 6
		secureNoAmbientRaiseLocked = 1 << 7
	)
	secureBits := uintptr(
		secureNoRoot |
			secureNoRootLocked |
			secureNoSetuidFixup |
			secureNoSetuidFixupLocked |
			secureNoAmbientRaise |
			secureNoAmbientRaiseLocked,
	)
	if err := unix.Prctl(unix.PR_SET_SECUREBITS, secureBits, 0, 0, 0); err != nil {
		return fmt.Errorf("lock namespace root capabilities: %w", err)
	}
	if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil {
		return fmt.Errorf("set no-new-privileges: %w", err)
	}
	_ = unix.Prctl(
		unix.PR_CAP_AMBIENT,
		unix.PR_CAP_AMBIENT_CLEAR_ALL,
		0,
		0,
		0,
	)
	for capability := 0; capability <= unix.CAP_LAST_CAP; capability++ {
		if err := unix.Prctl(
			unix.PR_CAPBSET_DROP,
			uintptr(capability),
			0,
			0,
			0,
		); err != nil && err != unix.EINVAL {
			return fmt.Errorf("drop capability %d from bounding set: %w", capability, err)
		}
	}
	header := unix.CapUserHeader{
		Version: unix.LINUX_CAPABILITY_VERSION_3,
		Pid:     0,
	}
	data := [2]unix.CapUserData{}
	if err := unix.Capset(&header, &data[0]); err != nil {
		return fmt.Errorf("clear effective capabilities: %w", err)
	}
	return nil
}