package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"reflect"
	"sync/atomic"
	"testing"

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

// testFileSeq generates unique numeric suffixes so concurrent/sequential
// subtests can each own a distinct /private/tmp/hazmat-<n>.sb path.
var testFileSeq atomic.Int64

// writePolicyFile creates /private/tmp/hazmat-<seq>.sb with the given
// content and mode.  The path satisfies policyFilePattern.
func writePolicyFile(t *testing.T, content string, mode os.FileMode) string {
	t.Helper()
	seq := testFileSeq.Add(1)
	path := fmt.Sprintf("/private/tmp/hazmat-%d.sb", seq)
	if err := os.WriteFile(path, []byte(content), mode); err != nil {
		t.Fatalf("write policy file: %v", err)
	}
	// Explicit chmod so the mode is exactly what we asked for regardless of
	// the process umask — mirrors what runAgentSeatbeltScript does.
	if err := os.Chmod(path, mode); err != nil {
		t.Fatalf("chmod policy file: %v", err)
	}
	t.Cleanup(func() { os.Remove(path) })
	return path
}

// ── Path pattern ──────────────────────────────────────────────────────────────

func TestValidatePolicyFile_PathPattern(t *testing.T) {
	tests := []struct {
		path    string
		wantErr bool
	}{
		{"/private/tmp/hazmat-1234.sb", false},
		{"/private/tmp/hazmat-99999.sb", false},
		{"/private/tmp/hazmat-abc.sb", true},          // non-numeric
		{"/tmp/hazmat-1234.sb", true},                 // wrong dir (symlink target)
		{"/private/tmp/evil/../hazmat-1234.sb", true}, // traversal
		{"/private/tmp/hazmat-.sb", true},             // empty PID component
		{"/private/tmp/xhazmat-1234.sb", true},        // wrong prefix
	}
	for _, tc := range tests {
		got := policyFilePattern.MatchString(tc.path)
		if tc.wantErr && got {
			t.Errorf("path %q: expected pattern to NOT match, but it did", tc.path)
		}
		if !tc.wantErr && !got {
			t.Errorf("path %q: expected pattern to match, but it did not", tc.path)
		}
	}
}

// ── Content check ─────────────────────────────────────────────────────────────

const (
	goodPolicy  = ";; Generated by hazmat\n(version 1)\n(deny default)\n(allow process-fork)\n"
	badPolicy   = "(version 1)\n(allow default)\n"
	emptyPolicy = ""
)

func TestValidatePolicyFile_ContentCheck(t *testing.T) {
	// Set SUDO_UID so the ownership check sees us as the invoking user.
	t.Setenv("SUDO_UID", fmt.Sprint(os.Getuid()))

	t.Run("deny-default passes", func(t *testing.T) {
		path := writePolicyFile(t, goodPolicy, 0o644)
		if err := validatePolicyFile(path); err != nil {
			t.Errorf("unexpected error for good policy: %v", err)
		}
	})

	t.Run("allow-default rejected", func(t *testing.T) {
		path := writePolicyFile(t, badPolicy, 0o644)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for allow-default policy, got nil")
		}
	})

	t.Run("empty policy rejected", func(t *testing.T) {
		path := writePolicyFile(t, emptyPolicy, 0o644)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for empty policy, got nil")
		}
	})

	t.Run("wrong mode rejected", func(t *testing.T) {
		path := writePolicyFile(t, goodPolicy, 0o600)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for 0600 mode, got nil")
		}
	})
}

// ── Ownership check ───────────────────────────────────────────────────────────

func TestValidatePolicyFile_OwnershipCheck(t *testing.T) {
	t.Run("matching SUDO_UID passes", func(t *testing.T) {
		t.Setenv("SUDO_UID", fmt.Sprint(os.Getuid()))
		path := writePolicyFile(t, goodPolicy, 0o644)
		if err := validatePolicyFile(path); err != nil {
			t.Errorf("unexpected error: %v", err)
		}
	})

	t.Run("wrong SUDO_UID rejected", func(t *testing.T) {
		if os.Getuid() == 9999 {
			t.Skip("running as uid 9999")
		}
		t.Setenv("SUDO_UID", "9999")
		path := writePolicyFile(t, goodPolicy, 0o644)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for mismatched SUDO_UID, got nil")
		}
	})

	t.Run("missing SUDO_UID rejected", func(t *testing.T) {
		t.Setenv("SUDO_UID", "")
		path := writePolicyFile(t, goodPolicy, 0o644)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for missing SUDO_UID, got nil")
		}
	})

	t.Run("root SUDO_UID rejected", func(t *testing.T) {
		t.Setenv("SUDO_UID", "0")
		path := writePolicyFile(t, goodPolicy, 0o644)
		if err := validatePolicyFile(path); err == nil {
			t.Error("expected error for SUDO_UID=0, got nil")
		}
	})
}

func TestCloseInheritedFDs_ClosesExtraFiles(t *testing.T) {
	const childEnv = "HAZMAT_LAUNCH_TEST_CHILD"

	if os.Getenv(childEnv) == "1" {
		if err := closeInheritedFDs(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(2)
		}
		openFDs, err := listOpenFDs(6)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(3)
		}
		if err := json.NewEncoder(os.Stdout).Encode(openFDs); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(4)
		}
		os.Exit(0)
	}

	sentinel, err := os.CreateTemp(t.TempDir(), "hazmat-launch-fd-*")
	if err != nil {
		t.Fatalf("create sentinel file: %v", err)
	}
	defer sentinel.Close()

	if _, err := unix.FcntlInt(sentinel.Fd(), unix.F_SETFD, 0); err != nil {
		t.Fatalf("clear FD_CLOEXEC on sentinel: %v", err)
	}

	cmd := exec.Command(os.Args[0], "-test.run=^TestCloseInheritedFDs_ClosesExtraFiles$")
	cmd.Env = append(os.Environ(), childEnv+"=1")
	cmd.ExtraFiles = []*os.File{sentinel}

	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("run helper child: %v\n%s", err, out)
	}

	var openFDs []int
	if err := json.Unmarshal(out, &openFDs); err != nil {
		t.Fatalf("decode open fds %q: %v", out, err)
	}

	if want := []int{0, 1, 2}; !reflect.DeepEqual(openFDs, want) {
		t.Fatalf("open fds after cleanup = %v, want %v", openFDs, want)
	}
}

func TestExecModeBypassesPolicyValidation(t *testing.T) {
	const childEnv = "HAZMAT_LAUNCH_EXEC_CHILD"

	if os.Getenv(childEnv) == "1" {
		os.Args = []string{"hazmat-launch", "exec", "/usr/bin/true"}
		main()
		os.Exit(99)
	}

	cmd := exec.Command(os.Args[0], "-test.run=^TestExecModeBypassesPolicyValidation$")
	cmd.Env = append(os.Environ(), childEnv+"=1")

	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("run exec-mode child: %v\n%s", err, out)
	}
	if got := string(out); got != "" {
		t.Fatalf("exec-mode child output = %q, want empty output", got)
	}
}

func TestParseLaunchModeArgsMetadata(t *testing.T) {
	metadata, cmdArgs, err := parseLaunchModeArgs([]string{
		"--hazmat-metadata-json",
		`{"kind":"hazmat.session"}`,
		"/usr/bin/env",
		"-i",
	})
	if err != nil {
		t.Fatalf("parseLaunchModeArgs: %v", err)
	}
	if metadata != `{"kind":"hazmat.session"}` {
		t.Fatalf("metadata = %q", metadata)
	}
	if !reflect.DeepEqual(cmdArgs, []string{"/usr/bin/env", "-i"}) {
		t.Fatalf("cmdArgs = %v", cmdArgs)
	}

	if _, _, err := parseLaunchModeArgs([]string{"--hazmat-metadata-json", `{}`}); err == nil {
		t.Fatal("expected missing command after metadata to be rejected")
	}
	if _, _, err := parseLaunchModeArgs(nil); err == nil {
		t.Fatal("expected missing command to be rejected")
	}
}

func listOpenFDs(maxFD int) ([]int, error) {
	openFDs := make([]int, 0, maxFD+1)
	for fd := 0; fd <= maxFD; fd++ {
		if _, err := unix.FcntlInt(uintptr(fd), unix.F_GETFD, 0); err != nil {
			if errors.Is(err, unix.EBADF) {
				continue
			}
			return nil, fmt.Errorf("inspect fd %d: %w", fd, err)
		}
		openFDs = append(openFDs, fd)
	}
	return openFDs, nil
}
