/******************************************************************************
 * Copyright (c)  2022 PingCAP                                                *
 * Licensed under the Apache License, Version 2.0 (the "License");            *
 * you may not use this file except in compliance with the License.           *
 * You may obtain a copy of the License at                                    *
 *                                                                            *
 * http://www.apache.org/licenses/LICENSE-2.0                                 *
 *                                                                            *
 * Unless required by applicable law or agreed to in writing, software        *
 * distributed under the License is distributed on an "AS IS" BASIS,          *
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
 * See the License for the specific language governing permissions and        *
 * limitations under the License.                                             *
 ******************************************************************************/

/*******************************************************************************
 * @File: sts_test.go
 * @Description:
 * @Author: jiangxunyu@pingcap.com
 * @Version: 1.0.0
 * @Date: 2022/6/27 17:08
*******************************************************************************/

package awsauth

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/stretchr/testify/assert"
	"k8s.io/client-go/rest/fake"
)

var validateToken = "eyJBY2Nlc3NLZXlJZCI6IkFTSUFZQkVHU1VNS0o0UUc1RzRRIiwiRXhwaXJhdGlvbiI6IjIwMjItMDYtMjNUMTE6MTU6MDdaIiwiU2VjcmV0QWNjZXNzS2V5IjoiU05iS2xEQytyWGZGNUhxQ0hQTUptYmJuMG5CMkVWVk04N1RlT2dUUCIsIlNlc3Npb25Ub2tlbiI6IklRb0piM0pwWjJsdVgyVmpFRnNhQ1hWekxYZGxjM1F0TWlKR01FUUNJQUo3VVFYTHZtRVRKRUw4cW5oamdTbmVEWCs0dEs3U1NGWlNLTzZ3YTBINkFpQThyN0FDRFBZSkRiWEcvSW9FaGFQTkhqbWJNZngxZUpNV1ZxaVdIK3p2VENxd0JRaDBFQUFhRERVMU1qRTROVFV6TnpNd01DSU1JZVp3RFo5NGlYYUFQbGlvS28wRkhZYjR3WVY1ODNBSCszQVN0WTlOcmVreUJlVXpzVUdKb2dWbEtBOFBJZU9LcjdKcHkrejJibU9WN2NhMk5NR2hsZmJYdW1kVkkyY0k0YUczbVN5T3I2eWc4MG1PeWlFOE1SRTQxN0FTSmdiaFpXQUc0YmQ5aVNiS09TdUxGVWp1bzFFd3MyNUp4SENPN3drMlJnc0F5MlFSOFhVakJoSjVSOEdpRUE4cVJ4ZDVjZHBEMllBeUlOR3l2Y0NtL0Y4QWtuUm1idk9IT2FtTkZBbnpiNGJNSENKN1hsOFpuOFdjWjNBYUJ2U1d3K3puVkIwdTQ1QUtxeWErZFoxc0VHS20rOFU5bDdtcmFWMjZndmdmMnFWV2F3TW45eU5aNmpRQ0lFNnFJTlV6RzY3SUsyWWtSaVdIM2xvNDByc1RWVXYxbjJLY1k0U3dUWk5ZOVgwSHd2a3pPTW1GVEd1anVRR2MrRzJGQjJsSUo1em1TQm1laWJEcFcrem8wRUl3OTNacnFjb2lHbzFYYko5NG1FT3FGQTdmRHFQZDF5d3pKby80Sk1LeSt1LzVGY1lzbHpaeVV2YmVQSitBVWVjSHE0dit0aHJCdXlZVXFyNnhmOUwvNDh3ekVzMUp6RitsM2V4U2hBNEovcWt4ZzdkVHNtK3JaVDQ0QU1CT21FZ3FNaUtMOHlVUVR4OHhCeDV1SmpNRS9wL0x5dk1vN0UwY2FsdVFLaTV3MUpBKzhPV2RVcVNGRGVuTmJtUmUwRU1WVlRBaWZTQ0dTL3V6Q3gwNGlYL0luamptckh1YmtPaytjNEhSSDNnUXl6UUZEVXVVd2xCSFN5ZUx5bEoveGQ5U1VYQ1FwUEpGeWdmYVlkelZULzZXRjlPZXY4bVRTbEJ2dUVHTXVtSnBSS2t6RkF4V3daRGhkNHo4Q3NIeDZyVDVENnFZUzFPUlZhT2IxMnlZNWpMM0FUSkxkb0loUWxVTVBqck9mSndYUThpUmpwcU9jNW5obmk2OXR6bDhrb0hjdm5CVStmTlB2ajJyZEw3dzMxYmpQTTUvR2l5UUExbWw2a1k0TVkyT0ljTExIRlFCM3VLTmpEL1hRRXFGOHlNZXIvWWZydWJIV0tXU2h3TG44U1VocmRKNDc2Wjh4OGRYWG1NOXNLU1BDSjR3dDQvUmxRWTZuQUd0MHE2SzVWVmVRanVpQnNiWDJ6cktwbU1HYUxXWW5GT2tEeDBzY3FieG1XQlAyTUVEUER3dlVhcmxYaGk2azVuL3Joc3NCZ0JYZ205Q1BqV0FoVjhJek8zYXFNWTd6cjdjcWhBTnJKMkxiTWRoWGNoVkFDL0RORTdHNm44YjI2VEo1WG8ybUI0VjZSbk1mU0M0citIVkhGR0Vtd1dtWndsbUtrOGxDeThTUWQrVjUvR1NtMGlmUXQ0cG5NYU05VGo1YU1EUmUvaXlFWnFNcW9NPSJ9"

// mock session which is used to hit the mock server
func mockAwsConfig() {
	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion("us-west-2"),
		config.WithCredentialsProvider(credentials.StaticCredentialsProvider{Value: aws.Credentials{
			AccessKeyID:     "test_id",
			SecretAccessKey: "test_key",
			SessionToken:    "test_session_token",
		}}),
		config.WithHTTPClient(fake.CreateHTTPClient(func(request *http.Request) (*http.Response, error) {
			return &http.Response{StatusCode: 200}, nil
		})))
	if err != nil {
		panic(err)
	}
	awsCfg = &cfg
}

func Test_GetStsTokenCache(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		cache := GetStsTokenCache()
		assert.NotNil(t, cache)
		cache.Set("testKey", "testValue", 100*time.Second)
		v, b := cache.Get("testKey")
		assert.Equal(t, true, b)
		assert.Equal(t, "testValue", v)
		cache.Flush()
	})
}

func Test_GetStsToken(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "us-west-2")

		mockAwsConfig()

		token, err := GetStsToken(context.TODO())
		assert.NoError(t, err)
		assert.NotEmpty(t, token)
	})

	t.Run("get cache token", func(t *testing.T) {
		cache := GetStsTokenCache()
		assert.NotNil(t, cache)
		cache.Set(cacheStsTokenKey, "tokenString", 100*time.Second)

		token, err := GetStsToken(context.TODO())
		assert.NoError(t, err)
		assert.Equal(t, "tokenString", token)

		cache.Flush()
	})

	t.Run("env is empty", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "")

		mockAwsConfig()
		token, err := GetStsToken(context.TODO())
		assert.Error(t, err)
		assert.Empty(t, token)
	})
}

func Test_ValidateStsToken(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		t.Setenv("AWS_REGION", "")
		token := validateToken

		mockAwsConfig()
		err := ValidateStsToken(context.TODO(), token)
		assert.NoError(t, err)
		err = ValidateStsToken(context.TODO(), token)
		assert.NoError(t, err)
	})

	t.Run("validate by cache", func(t *testing.T) {
		mockAwsConfig()
		cache := GetStsTokenCache()
		cache.Set("tokenString", time.Now().Unix()+10000, 100*time.Second)
		err := ValidateStsToken(context.TODO(), "tokenString")
		assert.NoError(t, err)
		cache.Flush()
	})

	t.Run("base64 decode error", func(t *testing.T) {
		var token = "123456"

		mockAwsConfig()
		err := ValidateStsToken(context.TODO(), token)
		assert.Error(t, err)
	})

	t.Run("json unmarshal error", func(t *testing.T) {
		var token = "MTIzNDU2Cg=="

		mockAwsConfig()
		err := ValidateStsToken(context.TODO(), token)
		assert.Error(t, err)
	})

	t.Run("validate token error", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "us-west-2")

		cfg, err := config.LoadDefaultConfig(context.TODO(),
			config.WithDefaultRegion("us-west-2"),
			config.WithCredentialsProvider(credentials.StaticCredentialsProvider{Value: aws.Credentials{
				AccessKeyID:     "test_id",
				SecretAccessKey: "test_key",
				SessionToken:    "test_session_token",
			}}),
			config.WithHTTPClient(fake.CreateHTTPClient(func(request *http.Request) (*http.Response, error) {
				return &http.Response{StatusCode: 500}, nil
			})))
		assert.NoError(t, err)
		SetAwsConfig(&cfg)

		err = ValidateStsToken(context.TODO(), validateToken)
		assert.Error(t, err)
	})

	t.Run("aws config is nil", func(t *testing.T) {
		awsCfg = nil
		err := ValidateStsToken(context.TODO(), validateToken)
		assert.Error(t, err)
	})
}

func Test_genStsCredentials(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "us-west-2")

		mockAwsConfig()
		cred, err := genStsCredentials(context.TODO())
		assert.NoError(t, err)
		assert.Nil(t, cred)
	})

	t.Run("env is empty", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "")
		_, err := genStsCredentials(context.TODO())
		assert.Error(t, err)

		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "")
		_, err = genStsCredentials(context.TODO())
		assert.Error(t, err)

		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "")
		_, err = genStsCredentials(context.TODO())
		assert.Error(t, err)
	})

	t.Run("token file not exists", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "/var/run/token")
		t.Setenv("AWS_REGION", "us-west-2")

		_, err := genStsCredentials(context.TODO())
		assert.Error(t, err)
	})

	t.Run("gen sts credentials error", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "us-west-2")

		cfg, err := config.LoadDefaultConfig(context.TODO(),
			config.WithDefaultRegion("us-west-2"),
			config.WithHTTPClient(fake.CreateHTTPClient(func(request *http.Request) (*http.Response, error) {
				return &http.Response{StatusCode: 500}, nil
			})))
		assert.NoError(t, err)
		SetAwsConfig(&cfg)

		_, err = genStsCredentials(context.TODO())
		assert.Error(t, err)
	})

	t.Run("aws config is nil", func(t *testing.T) {
		t.Setenv("AWS_ROLE_ARN", "arn:aws:iam::1234567890:role/test")
		t.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./sts_test.go")
		t.Setenv("AWS_REGION", "us-west-2")

		awsCfg = nil
		_, err := genStsCredentials(context.TODO())
		assert.Error(t, err)
	})
}

func Test_roleSessionName(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		name := roleSessionName()
		assert.NotEmpty(t, name)
	})
}