package rootaws

import (
	"bytes"
	"context"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
)

func TestScrubEnvironment(t *testing.T) {
	for _, tc := range []struct {
		name string
		euid int
	}{
		{"root", 0},
		{"non-root", 1001},
	} {
		t.Run(tc.name, func(t *testing.T) {
			for _, key := range rootScrubbedEnvironment {
				t.Setenv(key, "untrusted")
			}
			t.Setenv("AWS_ENDPOINT_URL_S3", "untrusted")
			t.Setenv("AWS_ENDPOINT_URL_STS", "untrusted")
			t.Setenv("AWS_ACCESS_KEY_ID", "KEEP")
			t.Setenv("AWS_SECRET_ACCESS_KEY", "KEEP_SECRET")
			t.Setenv("AWS_SESSION_TOKEN", "KEEP_TOKEN")
			var logs bytes.Buffer
			old := log.Writer()
			log.SetOutput(&logs)
			defer log.SetOutput(old)
			scrubEnvironment(tc.euid)
			for _, key := range rootScrubbedEnvironment {
				got, set := os.LookupEnv(key)
				want := "untrusted"
				if tc.euid == 0 {
					want = ""
					if key == "AWS_CONFIG_FILE" || key == "AWS_SHARED_CREDENTIALS_FILE" {
						want = disabledSharedFile
					}
				}
				if got != want || set != (want != "") {
					t.Errorf("%s = %q, set=%v; want %q", key, got, set, want)
				}
			}
			for _, key := range []string{"AWS_ENDPOINT_URL_S3", "AWS_ENDPOINT_URL_STS"} {
				if _, set := os.LookupEnv(key); set != (tc.euid != 0) {
					t.Errorf("%s set=%v for euid %d", key, set, tc.euid)
				}
			}
			wantLogs := 0
			if tc.euid == 0 {
				wantLogs = 1
			}
			if strings.Count(logs.String(), "ALERT:") != wantLogs {
				t.Fatalf("alert count: %q", logs.String())
			}
			if os.Getenv("AWS_ACCESS_KEY_ID") != "KEEP" || os.Getenv("AWS_SECRET_ACCESS_KEY") != "KEEP_SECRET" || os.Getenv("AWS_SESSION_TOKEN") != "KEEP_TOKEN" {
				t.Fatal("environment credentials changed")
			}
		})
	}
}

func TestScrubEnvironmentDisablesSharedFilesWithNothingSet(t *testing.T) {
	for _, key := range rootScrubbedEnvironment {
		t.Setenv(key, "")
		os.Unsetenv(key)
	}
	var logs bytes.Buffer
	old := log.Writer()
	log.SetOutput(&logs)
	defer log.SetOutput(old)
	scrubEnvironment(0)
	if os.Getenv("AWS_CONFIG_FILE") != disabledSharedFile || os.Getenv("AWS_SHARED_CREDENTIALS_FILE") != disabledSharedFile {
		t.Fatal("shared files not disabled for root")
	}
	if logs.Len() != 0 {
		t.Fatalf("alert without overrides: %q", logs.String())
	}
}

func TestScrubEnvironmentBlocksDependencySharedConfig(t *testing.T) {
	isolatedAWSEnv(t)
	dir := t.TempDir()
	marker := filepath.Join(dir, "executed")
	// The planted config sits under HOME, where the SDK looks by default, and
	// is also named explicitly.
	if err := os.Mkdir(filepath.Join(dir, ".aws"), 0o700); err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(dir, ".aws", "config")
	if err := os.WriteFile(path, []byte("[default]\ncredential_process = touch "+marker+"\n"), 0600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, ".aws", "credentials"), []byte("[default]\ncredential_process = touch "+marker+"\n"), 0600); err != nil {
		t.Fatal(err)
	}
	t.Setenv("HOME", dir)
	t.Setenv("AWS_CONFIG_FILE", path)
	t.Setenv("AWS_SHARED_CREDENTIALS_FILE", path)
	t.Setenv("AWS_PROFILE", "default")
	t.Setenv("AWS_DEFAULT_PROFILE", "default")
	t.Setenv("AWS_SDK_LOAD_CONFIG", "1")
	scrubEnvironment(0)
	// This deliberately bypasses rootaws, as a third-party SDK call would.
	cfg, err := awsconfig.LoadDefaultConfig(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := cfg.Credentials.Retrieve(context.Background()); err == nil {
		t.Fatal("unexpected credentials")
	}
	if _, err := os.Stat(marker); !os.IsNotExist(err) {
		t.Fatal("credential_process executed")
	}
	if _, err := os.Stat(disabledSharedFile); err == nil {
		t.Fatal("disabled path exists")
	}
}

func TestScrubEnvironmentKeepsProxyButBypassesIMDS(t *testing.T) {
	for _, tc := range []struct{ in, want string }{
		{"", "169.254.169.254,fd00:ec2::254"},
		{"example.com", "example.com,169.254.169.254,fd00:ec2::254"},
		{"169.254.169.254,example.com", "169.254.169.254,example.com,fd00:ec2::254"},
	} {
		t.Setenv("HTTP_PROXY", "http://proxy.example.com:3128")
		t.Setenv("NO_PROXY", tc.in)
		t.Setenv("no_proxy", tc.in)
		scrubEnvironment(0)
		if os.Getenv("HTTP_PROXY") != "http://proxy.example.com:3128" {
			t.Fatal("proxy removed")
		}
		if got := os.Getenv("NO_PROXY"); got != tc.want {
			t.Errorf("NO_PROXY(%q) = %q, want %q", tc.in, got, tc.want)
		}
		if got := os.Getenv("no_proxy"); got != tc.want {
			t.Errorf("no_proxy(%q) = %q, want %q", tc.in, got, tc.want)
		}
	}
	t.Setenv("NO_PROXY", "")
	scrubEnvironment(1001)
	if os.Getenv("NO_PROXY") != "" {
		t.Fatal("non-root NO_PROXY changed")
	}
}

// A root child of a root process inherits the parent's own disabled shared
// files. That is not an override: it must stay silent, or a guarded updater
// that reads a candidate's `version` with its output combined refuses the
// candidate (an armed box could then take no router update at all).
func TestScrubEnvironmentInheritedSentinelIsNotAnOverride(t *testing.T) {
	for _, key := range rootScrubbedEnvironment {
		t.Setenv(key, "")
		os.Unsetenv(key)
	}
	t.Setenv("AWS_CONFIG_FILE", disabledSharedFile)
	t.Setenv("AWS_SHARED_CREDENTIALS_FILE", disabledSharedFile)
	var logs bytes.Buffer
	old := log.Writer()
	log.SetOutput(&logs)
	defer log.SetOutput(old)
	scrubEnvironment(0)
	if logs.Len() != 0 {
		t.Fatalf("alert for the inherited sentinel: %q", logs.String())
	}
	if os.Getenv("AWS_CONFIG_FILE") != disabledSharedFile || os.Getenv("AWS_SHARED_CREDENTIALS_FILE") != disabledSharedFile {
		t.Fatal("shared files not disabled for root")
	}
	// Any other value for either key is still an override.
	for _, key := range []string{"AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"} {
		logs.Reset()
		t.Setenv("AWS_CONFIG_FILE", disabledSharedFile)
		t.Setenv("AWS_SHARED_CREDENTIALS_FILE", disabledSharedFile)
		t.Setenv(key, disabledSharedFile+"x")
		scrubEnvironment(0)
		if strings.Count(logs.String(), "ALERT:") != 1 {
			t.Fatalf("%s set to another file did not alert: %q", key, logs.String())
		}
		if os.Getenv(key) != disabledSharedFile {
			t.Fatalf("%s not reset to the sentinel", key)
		}
	}
}