package mail

import (
	"context"
	"errors"
	"reflect"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	iamtypes "github.com/aws/aws-sdk-go-v2/service/iam/types"
)

// stubIAM is an iamAPI double that records every call by name, in order,
// and lets a test script CreateUser and GetUser. Every other method
// answers success.
type stubIAM struct {
	calls []string

	createUser   func(*iam.CreateUserInput) (*iam.CreateUserOutput, error)
	getUser      func(*iam.GetUserInput) (*iam.GetUserOutput, error)
	listUsers    func(*iam.ListUsersInput) (*iam.ListUsersOutput, error)
	putBoundary  func(*iam.PutUserPermissionsBoundaryInput) error
	boundariesOn []string

	createIn   *iam.CreateUserInput
	boundaryIn *iam.PutUserPermissionsBoundaryInput
	policyIn   *iam.PutUserPolicyInput
}

func (s *stubIAM) CreateUser(_ context.Context, in *iam.CreateUserInput, _ ...func(*iam.Options)) (*iam.CreateUserOutput, error) {
	s.calls = append(s.calls, "CreateUser")
	s.createIn = in
	if s.createUser != nil {
		return s.createUser(in)
	}
	return &iam.CreateUserOutput{User: &iamtypes.User{Arn: aws.String("arn:aws:iam::123456789012:user/carlos-mail/" + aws.ToString(in.UserName))}}, nil
}

func (s *stubIAM) GetUser(_ context.Context, in *iam.GetUserInput, _ ...func(*iam.Options)) (*iam.GetUserOutput, error) {
	s.calls = append(s.calls, "GetUser")
	if s.getUser != nil {
		return s.getUser(in)
	}
	return &iam.GetUserOutput{User: &iamtypes.User{Arn: aws.String("arn:aws:iam::123456789012:user/carlos-mail/" + aws.ToString(in.UserName))}}, nil
}

func (s *stubIAM) PutUserPolicy(_ context.Context, in *iam.PutUserPolicyInput, _ ...func(*iam.Options)) (*iam.PutUserPolicyOutput, error) {
	s.calls = append(s.calls, "PutUserPolicy")
	s.policyIn = in
	return &iam.PutUserPolicyOutput{}, nil
}

func (s *stubIAM) PutUserPermissionsBoundary(_ context.Context, in *iam.PutUserPermissionsBoundaryInput, _ ...func(*iam.Options)) (*iam.PutUserPermissionsBoundaryOutput, error) {
	s.calls = append(s.calls, "PutUserPermissionsBoundary")
	s.boundaryIn = in
	if s.putBoundary != nil {
		if err := s.putBoundary(in); err != nil {
			return nil, err
		}
	}
	s.boundariesOn = append(s.boundariesOn, aws.ToString(in.UserName))
	return &iam.PutUserPermissionsBoundaryOutput{}, nil
}

func (s *stubIAM) DeleteUserPolicy(context.Context, *iam.DeleteUserPolicyInput, ...func(*iam.Options)) (*iam.DeleteUserPolicyOutput, error) {
	s.calls = append(s.calls, "DeleteUserPolicy")
	return &iam.DeleteUserPolicyOutput{}, nil
}

func (s *stubIAM) CreateAccessKey(context.Context, *iam.CreateAccessKeyInput, ...func(*iam.Options)) (*iam.CreateAccessKeyOutput, error) {
	s.calls = append(s.calls, "CreateAccessKey")
	return &iam.CreateAccessKeyOutput{AccessKey: &iamtypes.AccessKey{}}, nil
}

func (s *stubIAM) UpdateAccessKey(context.Context, *iam.UpdateAccessKeyInput, ...func(*iam.Options)) (*iam.UpdateAccessKeyOutput, error) {
	s.calls = append(s.calls, "UpdateAccessKey")
	return &iam.UpdateAccessKeyOutput{}, nil
}

func (s *stubIAM) DeleteAccessKey(context.Context, *iam.DeleteAccessKeyInput, ...func(*iam.Options)) (*iam.DeleteAccessKeyOutput, error) {
	s.calls = append(s.calls, "DeleteAccessKey")
	return &iam.DeleteAccessKeyOutput{}, nil
}

func (s *stubIAM) ListAccessKeys(context.Context, *iam.ListAccessKeysInput, ...func(*iam.Options)) (*iam.ListAccessKeysOutput, error) {
	s.calls = append(s.calls, "ListAccessKeys")
	return &iam.ListAccessKeysOutput{}, nil
}

func (s *stubIAM) ListUsers(_ context.Context, in *iam.ListUsersInput, _ ...func(*iam.Options)) (*iam.ListUsersOutput, error) {
	s.calls = append(s.calls, "ListUsers")
	if s.listUsers != nil {
		return s.listUsers(in)
	}
	return &iam.ListUsersOutput{}, nil
}

func (s *stubIAM) DeleteUser(context.Context, *iam.DeleteUserInput, ...func(*iam.Options)) (*iam.DeleteUserOutput, error) {
	s.calls = append(s.calls, "DeleteUser")
	return &iam.DeleteUserOutput{}, nil
}

const (
	testUserBoundary  = "arn:aws:iam::123456789012:policy/carlos-mail-boundary"
	otherUserBoundary = "arn:aws:iam::123456789012:policy/something-else"
)

// ensureTestUser runs the real awsProvisioner.EnsureUser against stub with
// a grant SendPolicyJSON accepts.
func ensureTestUser(t *testing.T, stub *stubIAM, boundary string) (string, error) {
	t.Helper()
	p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: boundary}
	return p.EnsureUser(context.Background(), "carlos-mail-bag-hm", UserPath,
		[]string{"arn:aws:ses:eu-west-1:123456789012:identity/example.com"}, nil, []string{"example.com"})
}

// existingUser makes CreateUser report the user already exists, and GetUser
// return it carrying pb (nil for a user with no boundary).
func existingUser(stub *stubIAM, pb *iamtypes.AttachedPermissionsBoundary) {
	stub.createUser = func(*iam.CreateUserInput) (*iam.CreateUserOutput, error) {
		return nil, &iamtypes.EntityAlreadyExistsException{Message: aws.String("exists")}
	}
	stub.getUser = func(in *iam.GetUserInput) (*iam.GetUserOutput, error) {
		return &iam.GetUserOutput{User: &iamtypes.User{
			UserName:            in.UserName,
			Arn:                 aws.String("arn:aws:iam::123456789012:user/carlos-mail/" + aws.ToString(in.UserName)),
			PermissionsBoundary: pb,
		}}, nil
	}
}

// (a) A new user is created under the configured boundary. Without it the
// console's role can make a /carlos-mail/ user with an admin policy.
func TestEnsureUserCreatesUnderTheBoundary(t *testing.T) {
	stub := &stubIAM{}
	if _, err := ensureTestUser(t, stub, testUserBoundary); err != nil {
		t.Fatalf("EnsureUser: %v", err)
	}
	if stub.createIn == nil || stub.createIn.PermissionsBoundary == nil {
		t.Fatal("CreateUser was called with NO permissions boundary")
	}
	if got := aws.ToString(stub.createIn.PermissionsBoundary); got != testUserBoundary {
		t.Fatalf("CreateUser PermissionsBoundary = %q, want %q", got, testUserBoundary)
	}
	if got := aws.ToString(stub.createIn.Path); got != UserPath {
		t.Fatalf("CreateUser Path = %q, want %q", got, UserPath)
	}
	if want := []string{"CreateUser", "PutUserPolicy"}; !reflect.DeepEqual(stub.calls, want) {
		t.Fatalf("calls = %v, want %v", stub.calls, want)
	}
}

// (b) An existing user with no boundary, or with a different one, gets the
// configured boundary set, and set BEFORE its policy is written: once the
// operator conditions PutUserPolicy on iam:PermissionsBoundary, a policy
// write on an unbounded user is refused.
func TestEnsureUserBackfillsTheBoundary(t *testing.T) {
	for _, tc := range []struct {
		name string
		pb   *iamtypes.AttachedPermissionsBoundary
	}{
		{"no boundary", nil},
		{"empty boundary struct", &iamtypes.AttachedPermissionsBoundary{}},
		{"another boundary", &iamtypes.AttachedPermissionsBoundary{
			PermissionsBoundaryArn:  aws.String(otherUserBoundary),
			PermissionsBoundaryType: iamtypes.PermissionsBoundaryAttachmentTypePolicy,
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			stub := &stubIAM{}
			existingUser(stub, tc.pb)
			arn, err := ensureTestUser(t, stub, testUserBoundary)
			if err != nil {
				t.Fatalf("EnsureUser: %v", err)
			}
			if arn != "arn:aws:iam::123456789012:user/carlos-mail/carlos-mail-bag-hm" {
				t.Fatalf("arn = %q", arn)
			}
			if want := []string{"CreateUser", "GetUser", "PutUserPermissionsBoundary", "PutUserPolicy"}; !reflect.DeepEqual(stub.calls, want) {
				t.Fatalf("calls = %v, want %v", stub.calls, want)
			}
			if got := aws.ToString(stub.boundaryIn.PermissionsBoundary); got != testUserBoundary {
				t.Fatalf("PutUserPermissionsBoundary PermissionsBoundary = %q, want %q", got, testUserBoundary)
			}
			if got := aws.ToString(stub.boundaryIn.UserName); got != "carlos-mail-bag-hm" {
				t.Fatalf("PutUserPermissionsBoundary UserName = %q", got)
			}
		})
	}
}

// (c) An existing user that already carries the boundary costs no extra
// call.
func TestEnsureUserLeavesACorrectBoundaryAlone(t *testing.T) {
	stub := &stubIAM{}
	existingUser(stub, &iamtypes.AttachedPermissionsBoundary{
		PermissionsBoundaryArn:  aws.String(testUserBoundary),
		PermissionsBoundaryType: iamtypes.PermissionsBoundaryAttachmentTypePolicy,
	})
	if _, err := ensureTestUser(t, stub, testUserBoundary); err != nil {
		t.Fatalf("EnsureUser: %v", err)
	}
	if want := []string{"CreateUser", "GetUser", "PutUserPolicy"}; !reflect.DeepEqual(stub.calls, want) {
		t.Fatalf("calls = %v, want %v", stub.calls, want)
	}
}

// A failed backfill stops the walk before the policy write: the user must
// never get a fresh policy while it still has no ceiling.
func TestEnsureUserBackfillFailureWritesNoPolicy(t *testing.T) {
	stub := &stubIAM{}
	existingUser(stub, nil)
	p := &awsProvisioner{iam: &failingBoundaryIAM{stubIAM: stub}, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
	_, err := p.EnsureUser(context.Background(), "carlos-mail-bag-hm", UserPath,
		[]string{"arn:aws:ses:eu-west-1:123456789012:identity/example.com"}, nil, []string{"example.com"})
	if err == nil {
		t.Fatal("EnsureUser succeeded although PutUserPermissionsBoundary failed")
	}
	for _, c := range stub.calls {
		if c == "PutUserPolicy" {
			t.Fatalf("PutUserPolicy ran after a failed backfill: calls = %v", stub.calls)
		}
	}
}

type failingBoundaryIAM struct{ *stubIAM }

func (f *failingBoundaryIAM) PutUserPermissionsBoundary(context.Context, *iam.PutUserPermissionsBoundaryInput, ...func(*iam.Options)) (*iam.PutUserPermissionsBoundaryOutput, error) {
	f.calls = append(f.calls, "PutUserPermissionsBoundary")
	return nil, errors.New("AccessDenied")
}

// Every path that creates or refreshes a sending user goes through
// EnsureUser, so the boundary has to hold there even for a provisioner
// built without the constructor: no IAM call at all.
func TestEnsureUserRefusesWithoutABoundary(t *testing.T) {
	stub := &stubIAM{}
	_, err := ensureTestUser(t, stub, "")
	if !errors.Is(err, ErrNoUserBoundary) {
		t.Fatalf("err = %v, want ErrNoUserBoundary", err)
	}
	if len(stub.calls) != 0 {
		t.Fatalf("IAM was called without a boundary: %v", stub.calls)
	}
}

// (d) The constructor refuses a missing or malformed boundary before it
// loads AWS config or calls STS, so this test makes no network call.
func TestNewAWSProvisionerRefusesABadBoundary(t *testing.T) {
	for _, b := range []string{
		"",
		"carlos-mail-boundary",
		"arn:aws:iam::123456789012:role/carlos-mail-boundary",
		"arn:aws:iam::123456789012:policy/",
		"arn:aws:iam::123456789012:policy/*",
		"arn:aws:s3:::policy/x",
		"arn:aws:iam:eu-west-1:123456789012:policy/x",
		"arn:aws:iam::not-an-account:policy/example",
		"arn:aws:iam::12345678901:policy/example",
		"arn:aws:iam::1234567890123:policy/example",
		"arn:aws:iam::aws:policy/AdministratorAccess",
		"arn:aws:iam::123456789012:policy/exa\tmple",
		"arn:aws:iam::123456789012:policy/example\n",
		"arn:aws:iam::123456789012:policy/exa mple",
		"arn:aws:iam::123456789012:policy/exa?mple",
		"arn:aws:iam::123456789012:policy//example",
		"arn:aws:iam::123456789012:policy/carlos/",
		"arn:AWS:iam::123456789012:policy/example",
		"\x00arn:aws:iam::123456789012:policy/example",
	} {
		p, err := NewAWSProvisioner(context.Background(), "eu-west-1", b)
		if !errors.Is(err, ErrNoUserBoundary) {
			t.Errorf("NewAWSProvisioner(boundary %q) err = %v, want ErrNoUserBoundary", b, err)
		}
		if p != nil {
			t.Errorf("NewAWSProvisioner(boundary %q) returned a provisioner", b)
		}
	}
}

func TestCheckUserBoundaryAcceptsPolicyARNs(t *testing.T) {
	for _, b := range []string{
		testUserBoundary,
		"arn:aws:iam::123456789012:policy/carlos/mail-boundary",
		"arn:aws-us-gov:iam::123456789012:policy/carlos-mail-boundary",
		"arn:aws-cn:iam::123456789012:policy/a/b/c+=,.@_-name",
	} {
		if err := CheckUserBoundary(b); err != nil {
			t.Errorf("CheckUserBoundary(%q) = %v, want nil", b, err)
		}
	}
}

// Rotation reaches MintKey without EnsureUser, so MintKey itself backfills:
// no key is minted for a user without the boundary.
func TestMintKeyBackfillsTheBoundary(t *testing.T) {
	for _, tc := range []struct {
		name string
		pb   *iamtypes.AttachedPermissionsBoundary
		want []string
	}{
		{"unbounded", nil, []string{"GetUser", "PutUserPermissionsBoundary", "CreateAccessKey"}},
		{"another boundary", &iamtypes.AttachedPermissionsBoundary{PermissionsBoundaryArn: aws.String(otherUserBoundary)},
			[]string{"GetUser", "PutUserPermissionsBoundary", "CreateAccessKey"}},
		{"bounded", &iamtypes.AttachedPermissionsBoundary{PermissionsBoundaryArn: aws.String(testUserBoundary)},
			[]string{"GetUser", "CreateAccessKey"}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			stub := &stubIAM{}
			existingUser(stub, tc.pb)
			p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
			if _, _, err := p.MintKey(context.Background(), "carlos-mail-bag-hm"); err != nil {
				t.Fatalf("MintKey: %v", err)
			}
			if !reflect.DeepEqual(stub.calls, tc.want) {
				t.Fatalf("calls = %v, want %v", stub.calls, tc.want)
			}
			if tc.pb == nil && aws.ToString(stub.boundaryIn.PermissionsBoundary) != testUserBoundary {
				t.Fatalf("boundary set = %q, want %q", aws.ToString(stub.boundaryIn.PermissionsBoundary), testUserBoundary)
			}
		})
	}
}

func TestMintKeyRefusesWhenTheBackfillFails(t *testing.T) {
	stub := &stubIAM{putBoundary: func(*iam.PutUserPermissionsBoundaryInput) error { return errors.New("AccessDenied") }}
	existingUser(stub, nil)
	p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
	if _, _, err := p.MintKey(context.Background(), "carlos-mail-bag-hm"); err == nil {
		t.Fatal("MintKey succeeded although the boundary could not be set")
	}
	for _, c := range stub.calls {
		if c == "CreateAccessKey" {
			t.Fatalf("a key was minted for an unbounded user: calls = %v", stub.calls)
		}
	}
}

// sweepStub serves users over two ListUsers pages and answers GetUser from
// boundaries (absent = no boundary).
func sweepStub(t *testing.T, pages [][]string, boundaries map[string]string) *stubIAM {
	t.Helper()
	stub := &stubIAM{}
	stub.listUsers = func(in *iam.ListUsersInput) (*iam.ListUsersOutput, error) {
		if aws.ToString(in.PathPrefix) != UserPath {
			t.Fatalf("ListUsers PathPrefix = %q, want %q", aws.ToString(in.PathPrefix), UserPath)
		}
		page := 0
		if in.Marker != nil {
			page = 1
		}
		out := &iam.ListUsersOutput{}
		for _, n := range pages[page] {
			out.Users = append(out.Users, iamtypes.User{UserName: aws.String(n)})
		}
		if page+1 < len(pages) {
			out.IsTruncated, out.Marker = true, aws.String("next")
		}
		return out, nil
	}
	stub.getUser = func(in *iam.GetUserInput) (*iam.GetUserOutput, error) {
		u := &iamtypes.User{UserName: in.UserName, Arn: aws.String("arn:aws:iam::123456789012:user/carlos-mail/" + aws.ToString(in.UserName))}
		if b, ok := boundaries[aws.ToString(in.UserName)]; ok {
			u.PermissionsBoundary = &iamtypes.AttachedPermissionsBoundary{PermissionsBoundaryArn: aws.String(b)}
		}
		return &iam.GetUserOutput{User: u}, nil
	}
	return stub
}

// The sweep bounds every unbounded or wrongly bounded user under the path,
// across pages, and leaves bounded ones alone.
func TestSweepUserBoundariesBoundsEveryUser(t *testing.T) {
	stub := sweepStub(t,
		[][]string{{"carlos-mail-bag-a", "carlos-mail-bag-b"}, {"carlos-mail-bag-c", "carlos-mail-bag-d+s+laptop"}},
		map[string]string{"carlos-mail-bag-b": testUserBoundary, "carlos-mail-bag-c": otherUserBoundary})
	p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
	res, err := p.SweepUserBoundaries(context.Background())
	if err != nil {
		t.Fatalf("SweepUserBoundaries: %v", err)
	}
	if res.Total != 4 || res.Bounded != 1 || len(res.Failed) != 0 {
		t.Fatalf("result = %+v, want Total 4, Bounded 1, no failures", res)
	}
	want := []string{"carlos-mail-bag-a", "carlos-mail-bag-c", "carlos-mail-bag-d+s+laptop"}
	if !reflect.DeepEqual(res.Fixed, want) || !reflect.DeepEqual(stub.boundariesOn, want) {
		t.Fatalf("fixed = %v, boundaries set on %v, want %v", res.Fixed, stub.boundariesOn, want)
	}
	if aws.ToString(stub.boundaryIn.PermissionsBoundary) != testUserBoundary {
		t.Fatalf("boundary set = %q", aws.ToString(stub.boundaryIn.PermissionsBoundary))
	}
}

// A user the sweep cannot fix or read is reported, not swallowed, and the
// sweep goes on to the rest.
func TestSweepUserBoundariesReportsFailures(t *testing.T) {
	stub := sweepStub(t, [][]string{{"carlos-mail-bag-a", "carlos-mail-bag-b", "carlos-mail-bag-c"}}, nil)
	stub.putBoundary = func(in *iam.PutUserPermissionsBoundaryInput) error {
		if aws.ToString(in.UserName) == "carlos-mail-bag-a" {
			return errors.New("AccessDenied")
		}
		return nil
	}
	read := stub.getUser
	stub.getUser = func(in *iam.GetUserInput) (*iam.GetUserOutput, error) {
		if aws.ToString(in.UserName) == "carlos-mail-bag-b" {
			return nil, errors.New("Throttling")
		}
		return read(in)
	}
	p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
	res, err := p.SweepUserBoundaries(context.Background())
	if err != nil {
		t.Fatalf("SweepUserBoundaries: %v", err)
	}
	if res.Total != 3 || !reflect.DeepEqual(res.Fixed, []string{"carlos-mail-bag-c"}) || len(res.Failed) != 2 {
		t.Fatalf("result = %+v, want 3 total, c fixed, a and b failed", res)
	}
	if res.Failed[0].User != "carlos-mail-bag-a" || res.Failed[1].User != "carlos-mail-bag-b" || res.Failed[0].Err == nil || res.Failed[1].Err == nil {
		t.Fatalf("failed = %+v", res.Failed)
	}
}

// A failed listing is an error: the counts cannot claim every user.
func TestSweepUserBoundariesListFailureIsAnError(t *testing.T) {
	stub := &stubIAM{listUsers: func(*iam.ListUsersInput) (*iam.ListUsersOutput, error) {
		return nil, errors.New("AccessDenied: iam:ListUsers")
	}}
	p := &awsProvisioner{iam: stub, account: "123456789012", region: "eu-west-1", boundary: testUserBoundary}
	if _, err := p.SweepUserBoundaries(context.Background()); err == nil {
		t.Fatal("SweepUserBoundaries succeeded although ListUsers failed")
	}
}