From 6eccbcf7426e0cf662a2c2c2a0d3c9e806b4e703 Mon Sep 17 00:00:00 2001
From: Ohad Aharoni <oaharoni@redhat.com>
Date: Thu, 23 May 2024 10:38:31 -0400
Subject: [PATCH] OCM-8256 | feat: add methods to attach policies

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

diff --git a/pkg/aws/aws_client/role.go b/pkg/aws/aws_client/role.go
index c7c00d8..329b5a2 100644
--- a/pkg/aws/aws_client/role.go
+++ b/pkg/aws/aws_client/role.go
@@ -48,6 +48,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)
+	}
+	return role, err
+}
+
 func (client *AWSClient) GetRole(roleName string) (*types.Role, error) {
 	input := &iam.GetRoleInput{
 		RoleName: &roleName,
@@ -162,7 +175,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 +204,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 +222,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 +245,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 +305,52 @@ 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) error {
+	policyAttach := iam.AttachRolePolicyInput{
+		PolicyArn: &policyArn,
+		RoleName:  &roleName,
+	}
+	_, err := client.IamClient.AttachRolePolicy(context.TODO(), &policyAttach)
+	if err != nil {
+		return err
+	}
+	timeout := 2
+	start := 0
+
+	for start < timeout {
+
+		if attached, _ := client.PolicyAttachedToRole(roleName, policyArn); attached {
+			return nil
+		}
+		time.Sleep(1 * time.Minute)
+		start++
+	}
+	return err
+}
+
+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
+}