//go:build unit || all
// +build unit all

package k8s

import (
	"context"
	"math/rand"
	"os"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts/stsiface"
	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/aws-iam-authenticator/pkg/token"
	"sigs.k8s.io/controller-runtime/pkg/client"
	_ "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	_ "sigs.k8s.io/controller-runtime/pkg/client/fake"
)

const (
	defaultSecretName     = "cred-123456789012-xaccounts3access-us-west-2-eks1"
	AwsAccessKeyStr       = "AWS_ACCESS_KEY_ID"
	AwsSecretAccessKeyStr = "AWS_SECRET_ACCESS_KEY"
)

func defaultTarget() EKSTarget {
	return EKSTarget{
		AssumeRole: "arn:aws:iam::123456789012:role/xaccounts3access",
		Region:     "us-west-2",
		Cluster:    "eks1",
		Namespace:  "default",
	}
}

func tokenOpt(f func(opt *EKSBearerTokenOpt)) EKSBearerTokenOpt {
	res := EKSBearerTokenOpt{
		EKSTarget: defaultTarget(),
	}
	if f != nil {
		f(&res)
	}
	return res
}

func randomString(n int) string {
	var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")

	s := make([]rune, n)
	for i := range s {
		s[i] = letters[rand.Intn(len(letters))]
	}
	return string(s)
}

func TestParseAssumeRoleKey(t *testing.T) {
	cases := map[string]struct {
		opt       EKSBearerTokenOpt
		expectErr bool
		expectKey string
	}{
		"valid-role-arn": {
			opt:       tokenOpt(nil),
			expectErr: false,
			expectKey: defaultSecretName,
		},
		"invalid-role-arn": {
			opt: tokenOpt(func(opt *EKSBearerTokenOpt) {
				opt.AssumeRole = "invalid-arn"
			}),
			expectErr: true,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			key, err := retrieveRoleCredSecretKey(tc.opt)
			if !tc.expectErr {
				g.Expect(err).To(Succeed())
			} else {
				g.Expect(err).NotTo(Succeed())
				return
			}
			g.Expect(key.Name).To(Equal(tc.expectKey))
		})
	}
}

func TestUpdateTokenSecret(t *testing.T) {
	cases := map[string]struct {
		opt       EKSBearerTokenOpt
		secret    *corev1.Secret
		expectErr bool
	}{
		"token-not-exist": {
			opt:       tokenOpt(nil),
			expectErr: false,
		},
		"token-exist": {
			opt: tokenOpt(nil),
			secret: &corev1.Secret{
				ObjectMeta: metav1.ObjectMeta{
					Name:      defaultSecretName,
					Namespace: tokenOpt(nil).Namespace,
				},
				StringData: map[string]string{
					tokenExpirationKey: "old-expire",
					bearerTokenKey:     "old-token",
				},
			},
			expectErr: false,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(
				func() []client.Object {
					if tc.secret == nil {
						return nil
					} else {
						return []client.Object{tc.secret}
					}
				}()...,
			).Build()
			e := &eksTokenMngr{cli: cli}

			randToken, randExpire := randomString(20), randomString(20)
			err := e.updateSecret(context.TODO(), randToken, randExpire, tc.secret, tc.opt)
			if !tc.expectErr {
				g.Expect(err).To(Succeed())
			} else {
				g.Expect(err).NotTo(Succeed())
				return
			}
			key, err := retrieveRoleCredSecretKey(tc.opt)
			g.Expect(err).To(Succeed())
			newSecret := &corev1.Secret{}
			g.Expect(cli.Get(context.TODO(), key, newSecret)).To(Succeed())
			retrieveToken, err := retrieveValueFromSecret(newSecret, bearerTokenKey)
			g.Expect(err).To(Succeed())
			g.Expect(retrieveToken).To(Equal(randToken))
			retrieveExpire, err := retrieveValueFromSecret(newSecret, tokenExpirationKey)
			g.Expect(err).To(Succeed())
			g.Expect(retrieveExpire).To(Equal(randExpire))
		})
	}
}

func TestCheckSecretExpired(t *testing.T) {
	cases := map[string]struct {
		opt      EKSBearerTokenOpt
		expireTS string
		now      func() time.Time
		expired  bool
	}{
		"expired": {
			opt:      tokenOpt(nil),
			expireTS: "2021-10-28T19:43:33Z",
			now: func() time.Time {
				// year 2020, before 2021 in expire ts
				return time.Date(2020, 1, 2, 0, 0, 0, 0, time.UTC)
			},
			expired: false,
		},
		"not-expired": {
			opt:      tokenOpt(nil),
			expireTS: "2021-10-28T19:43:33Z",
			now: func() time.Time {
				// year 2030, after 2021 in expire ts
				return time.Date(2030, 1, 2, 0, 0, 0, 0, time.UTC)
			},
			expired: true,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects().Build()
			e := &eksTokenMngr{cli: cli}

			randToken := randomString(20)
			g.Expect(e.updateSecret(context.TODO(), randToken, tc.expireTS, nil, tc.opt)).To(Succeed())
			key, err := retrieveRoleCredSecretKey(tc.opt)
			g.Expect(err).To(Succeed())
			newSecret := &corev1.Secret{}
			g.Expect(cli.Get(context.TODO(), key, newSecret)).To(Succeed())

			expired, err := isSecretExpired(newSecret, tc.now)
			g.Expect(err).To(Succeed())
			g.Expect(expired).To(Equal(tc.expired))
		})
	}
}

func restoreEnv(key string) func() {
	val, exist := os.LookupEnv(key)
	return func() {
		if exist {
			os.Setenv(key, val)
		} else {
			os.Unsetenv(key)
		}
	}
}

// fakeGenerator is used to mock token.Generator
type fakeGenerator struct {
	fakeToken token.Token
	fakeErr   error
	fakeJson  string
}

func (f *fakeGenerator) Get(string) (token.Token, error) {
	return f.fakeToken, f.fakeErr
}

func (f *fakeGenerator) GetWithRole(clusterID, roleARN string) (token.Token, error) {
	return f.fakeToken, f.fakeErr
}

func (f *fakeGenerator) GetWithRoleForSession(clusterID string, roleARN string, sess *session.Session) (token.Token, error) {
	return f.fakeToken, f.fakeErr
}

func (f *fakeGenerator) GetWithOptions(options *token.GetTokenOptions) (token.Token, error) {
	return f.fakeToken, f.fakeErr
}

func (f *fakeGenerator) GetWithSTS(clusterID string, stsAPI stsiface.STSAPI) (token.Token, error) {
	return f.fakeToken, f.fakeErr
}

func (f *fakeGenerator) FormatJSON(token.Token) string {
	return f.fakeJson
}

func TestRetrieveToken(t *testing.T) {
	defer restoreEnv(AwsAccessKeyStr)()
	defer restoreEnv(AwsSecretAccessKeyStr)()

	os.Setenv(AwsAccessKeyStr, "fake-key")
	os.Setenv(AwsSecretAccessKeyStr, "fake-secret-key")

	cases := map[string]struct {
		opt       EKSBearerTokenOpt
		secret    *corev1.Secret
		oldExpire string // expire ts in existing secret
		oldTokStr string // token string in existing secret
		expectErr bool
		expectTok string
		fakeToken token.Token
		fakeErr   error
		fakeJson  string
		fakeNow   func() time.Time // overwrites the 'now' oracle in option
	}{
		"no-existing-secret": {
			opt:       tokenOpt(nil),
			expectErr: false,
			expectTok: "fake-token",
			fakeToken: token.Token{
				Token:      "fake-token",
				Expiration: time.Date(2030, 1, 2, 0, 0, 0, 0, time.UTC),
			},
		},
		"has-non-expired-secret": {
			opt:       tokenOpt(nil),
			oldExpire: "2021-10-28T19:43:33Z",
			oldTokStr: "old-non-expired-token",
			expectErr: false,
			expectTok: "old-non-expired-token",
			fakeToken: token.Token{
				Token:      "new-token",
				Expiration: time.Date(2030, 1, 2, 0, 0, 0, 0, time.UTC),
			},
			fakeNow: func() time.Time { return time.Date(2020, 1, 2, 0, 0, 0, 0, time.UTC) },
		},
		"has-expired-secret": {
			opt:       tokenOpt(nil),
			oldExpire: "2021-10-28T19:43:33Z",
			oldTokStr: "old-non-expired-token",
			expectErr: false,
			expectTok: "new-token",
			fakeToken: token.Token{
				Token:      "new-token",
				Expiration: time.Date(2030, 1, 2, 0, 0, 0, 0, time.UTC),
			},
			fakeNow: func() time.Time { return time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC) },
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects().Build()
			tokenGen := &fakeGenerator{
				fakeToken: tc.fakeToken,
				fakeErr:   tc.fakeErr,
				fakeJson:  tc.fakeJson,
			}
			opt := tc.opt
			opt.now = tc.fakeNow
			if tc.oldExpire != "" {
				e := eksTokenMngr{cli: cli}
				g.Expect(e.updateSecret(context.TODO(), tc.oldTokStr, tc.oldExpire, nil, opt)).To(Succeed())
			}

			e := NewEKSTokenManager(cli, tokenGen)
			tokenStr, err := e.RetrieveToken(context.TODO(), opt)
			if tc.expectErr {
				g.Expect(err).NotTo(Succeed())
				return
			}
			g.Expect(err).To(Succeed())
			g.Expect(tokenStr).To(Equal(tc.expectTok))
		})
	}
}