From 3e91655315ba63d9c40d1fcb16e92b5516c8c9a8 Mon Sep 17 00:00:00 2001
From: Ohad Aharoni <oaharoni@redhat.com>
Date: Mon, 3 Jun 2024 11:43:15 -0400
Subject: [PATCH] OCM-8256 | feat: add methods to attach policies

---
 pkg/aws/aws_client/role.go | 91 +++++++++++++++++++++++++++++++++++---
 1 file changed, 85 insertions(+), 6 deletions(-)

diff --git a/pkg/aws/aws_client/role.go b/pkg/aws/aws_client/role.go
index c7c00d8a..95c0ea4f 100644
--- a/pkg/aws/aws_client/role.go
+++ b/pkg/aws/aws_client/role.go
@@ -11,6 +11,11 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/iam/types"
 )
 
+var (
+	DEFAULT_RETRIES        = 3
+	DEFAULT_RETRY_DURATION = 60 * time.Second
+)
+
 func (client *AWSClient) CreateRole(roleName string,
 	assumeRolePolicyDocument string,
 	permissionBoundry string,
@@ -48,6 +53,19 @@ func (client *AWSClient) CreateRole(roleName string,
 	return *resp.Role, err
 }
 
+func (client *AWSClient) CreateRoleAndAttachPolicy(roleName string,
+	assumeRolePolicyDocument string,
+	permissionBoundry string,
+	tags map[string]string,
+	path string,
+	policyArn string) (types.Role, error) {
+	role, err := client.CreateRole(roleName, assumeRolePolicyDocument, permissionBoundry, tags, path)
+	if err == nil {
+		err = client.AttachPolicy(*role.RoleName, policyArn, DEFAULT_RETRIES, DEFAULT_RETRY_DURATION)
+	}
+	return role, err
+}
+
 func (client *AWSClient) GetRole(roleName string) (*types.Role, error) {
 	input := &iam.GetRoleInput{
 		RoleName: &roleName,
@@ -162,7 +180,7 @@ func (client *AWSClient) DeleteRoleInstanceProfiles(roleName string) error {
 	return nil
 }
 
-func (client *AWSClient) CreateIAMRole(roleName string, ProdENVTrustedRole string, StageENVTrustedRole string, StageIssuerTrustedRole string,
+func (client *AWSClient) CreateIAMRole(roleName string, ProdENVTrustedRole string, StageENVTrustedRole string, StageIssuerTrustedRole string, policyArn string,
 	externalID ...string) (types.Role, error) {
 	statement := map[string]interface{}{
 		"Effect": "Allow",
@@ -191,10 +209,10 @@ func (client *AWSClient) CreateIAMRole(roleName string, ProdENVTrustedRole strin
 		return types.Role{}, err
 	}
 
-	return client.CreateRole(roleName, string(assumeRolePolicyDocument), "", make(map[string]string), "/")
+	return client.CreateRoleAndAttachPolicy(roleName, string(assumeRolePolicyDocument), "", make(map[string]string), "/", policyArn)
 }
 
-func (client *AWSClient) CreateRegularRole(roleName string) (types.Role, error) {
+func (client *AWSClient) CreateRegularRole(roleName string, policyArn string) (types.Role, error) {
 
 	statement := map[string]interface{}{
 		"Effect": "Allow",
@@ -209,10 +227,10 @@ func (client *AWSClient) CreateRegularRole(roleName string) (types.Role, error)
 		fmt.Println("Failed to convert Role Policy Document into JSON: ", err)
 		return types.Role{}, err
 	}
-	return client.CreateRole(roleName, assumeRolePolicyDocument, "", make(map[string]string), "/")
+	return client.CreateRoleAndAttachPolicy(roleName, assumeRolePolicyDocument, "", make(map[string]string), "/", policyArn)
 }
 
-func (client *AWSClient) CreateRoleForAuditLogForward(roleName, awsAccountID string, oidcEndpointURL string) (types.Role, error) {
+func (client *AWSClient) CreateRoleForAuditLogForward(roleName, awsAccountID string, oidcEndpointURL string, policyArn string) (types.Role, error) {
 	statement := map[string]interface{}{
 		"Effect": "Allow",
 		"Principal": map[string]interface{}{
@@ -232,7 +250,7 @@ func (client *AWSClient) CreateRoleForAuditLogForward(roleName, awsAccountID str
 		return types.Role{}, err
 	}
 
-	return client.CreateRole(roleName, string(assumeRolePolicyDocument), "", make(map[string]string), "/")
+	return client.CreateRoleAndAttachPolicy(roleName, string(assumeRolePolicyDocument), "", make(map[string]string), "/", policyArn)
 }
 
 func (client *AWSClient) CreatePolicy(policyName string, statements ...map[string]interface{}) (string, error) {
@@ -292,3 +310,64 @@ func completeRolePolicyDocument(statement map[string]interface{}) (string, error
 	assumeRolePolicyDocument, err := json.Marshal(rolePolicyDocument)
 	return string(assumeRolePolicyDocument), err
 }
+
+func (client *AWSClient) AttachPolicy(roleName string, policyArn string, retries int, retryIntervalInSeconds time.Duration) error {
+	policyAttach := iam.AttachRolePolicyInput{
+		PolicyArn: &policyArn,
+		RoleName:  &roleName,
+	}
+	_, err := client.IamClient.AttachRolePolicy(context.TODO(), &policyAttach)
+	if err != nil {
+		return err
+	}
+
+	attached, err := client.PolicyAttachedToRole(roleName, policyArn)
+
+	if err != nil {
+		return err
+	}
+	if attached {
+		return nil
+	}
+
+	start := 0
+
+	for start < retries {
+		attached, err := client.PolicyAttachedToRole(roleName, policyArn)
+		if err != nil && start == retries {
+			return err
+		}
+		if attached {
+			return nil
+		}
+		time.Sleep(retryIntervalInSeconds)
+		start++
+	}
+	return fmt.Errorf("failed to attach policy to role but no errors were thrown, please investigate")
+}
+
+func (client *AWSClient) PolicyAttachedToRole(roleName string, policyArn string) (bool, error) {
+	policies, err := client.ListRoleAttachedPolicies(roleName)
+	if err != nil {
+		return false, err
+	}
+	for _, policy := range policies {
+		if *policy.PolicyArn == policyArn {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+func (client *AWSClient) ListRoleAttachedPolicies(roleName string) ([]types.AttachedPolicy, error) {
+	policies := []types.AttachedPolicy{}
+	policyLister := iam.ListAttachedRolePoliciesInput{
+		RoleName: &roleName,
+	}
+	policyOut, err := client.IamClient.ListAttachedRolePolicies(context.TODO(), &policyLister)
+	if err != nil {
+		return policies, err
+	}
+	policies = policyOut.AttachedPolicies
+	return policies, nil
+}