From fedcc71ee42770054015fea405dc7aa7783621fb Mon Sep 17 00:00:00 2001
From: yingzhan <yingzhan@redhat.com>
Date: Tue, 28 May 2024 14:16:35 +0800
Subject: [PATCH] OCM-8208 | test: Add new function: GrantValidAccessKeys for
 backend usage

---
 pkg/aws/aws_client/client.go | 38 ++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/pkg/aws/aws_client/client.go b/pkg/aws/aws_client/client.go
index fd5f02c..ff5428e 100644
--- a/pkg/aws/aws_client/client.go
+++ b/pkg/aws/aws_client/client.go
@@ -18,6 +18,8 @@ import (
 
 	elb "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancing"
 	"github.com/aws/aws-sdk-go-v2/service/route53"
+
+	CON "github.com/openshift-online/ocm-common/pkg/aws/consts"
 )
 
 type AWSClient struct {
@@ -32,6 +34,12 @@ type AWSClient struct {
 	AccountID            string
 	KmsClient            *kms.Client
 	CloudWatchLogsClient *cloudwatchlogs.Client
+	AWSConfig            *aws.Config
+}
+
+type AccessKeyMod struct {
+	AccessKeyId     string `ini:"aws_access_key_id,omitempty"`
+	SecretAccessKey string `ini:"aws_secret_access_key,omitempty"`
 }
 
 func CreateAWSClient(profileName string, region string) (*AWSClient, error) {
@@ -79,6 +87,7 @@ func CreateAWSClient(profileName string, region string) (*AWSClient, error) {
 		IamClient:            iam.NewFromConfig(cfg),
 		ClientContext:        context.TODO(),
 		KmsClient:            kms.NewFromConfig(cfg),
+		AWSConfig:            &cfg,
 	}
 	awsClient.AccountID = awsClient.GetAWSAccountID()
 	return awsClient, nil
@@ -106,3 +115,32 @@ func (client *AWSClient) CloudFormation() *cloudformation.Client {
 func (client *AWSClient) ELB() *elb.Client {
 	return client.ElbClient
 }
+
+func GrantValidAccessKeys(userName string) (*AccessKeyMod, error) {
+	var cre aws.Credentials
+	var keysMod *AccessKeyMod
+	var err error
+	retryTimes := 3
+	for retryTimes > 0 {
+		if cre.AccessKeyID != "" {
+			break
+		}
+		client, err := CreateAWSClient(userName, CON.DefaultAWSRegion)
+		if err != nil {
+			return nil, err
+		}
+
+		cre, err = client.AWSConfig.Credentials.Retrieve(client.ClientContext)
+		if err != nil {
+			return nil, err
+		}
+		log.LogInfo(">>> Access key grant successfully")
+
+		keysMod = &AccessKeyMod{
+			AccessKeyId:     cre.AccessKeyID,
+			SecretAccessKey: cre.SecretAccessKey,
+		}
+		retryTimes--
+	}
+	return keysMod, err
+}