From 6ed4fd0e42ce73614431a10e88f680363dbc99a7 Mon Sep 17 00:00:00 2001
From: Mohamed Habib <moe.habib9@gmail.com>
Date: Thu, 21 Mar 2024 15:06:53 +0000
Subject: [PATCH] Correctly pass project roles from orchestrator to cli (#1297)

* Correctly pass project roles from orchestrator to cli
---
 backend/utils/graphs.go          |  4 +--
 cli/pkg/digger/digger.go         |  6 +++++
 libs/orchestrator/aws.go         |  3 +++
 libs/orchestrator/json_models.go | 45 +++++++++++++++++++++-----------
 4 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/backend/utils/graphs.go b/backend/utils/graphs.go
index a3dec48d1..d250a3705 100644
--- a/backend/utils/graphs.go
+++ b/backend/utils/graphs.go
@@ -19,8 +19,8 @@ func ConvertJobsToDiggerJobs(jobsMap map[string]orchestrator.Job, projectMap map
 
 	log.Printf("Number of Jobs: %v\n", len(jobsMap))
 	marshalledJobsMap := map[string][]byte{}
-	for _, job := range jobsMap {
-		marshalled, err := json.Marshal(orchestrator.JobToJson(job))
+	for projectName, job := range jobsMap {
+		marshalled, err := json.Marshal(orchestrator.JobToJson(job, projectMap[projectName]))
 		if err != nil {
 			return nil, nil, err
 		}
diff --git a/cli/pkg/digger/digger.go b/cli/pkg/digger/digger.go
index 905411c94..e44cc0a3d 100644
--- a/cli/pkg/digger/digger.go
+++ b/cli/pkg/digger/digger.go
@@ -545,6 +545,12 @@ func RunJob(
 			}
 			return errors.New(msg)
 		}
+
+		err = job.PopulateAwsCredentialsEnvVarsForJob()
+		if err != nil {
+			log.Fatalf("failed to fetch AWS keys, %v", err)
+		}
+
 		var terraformExecutor terraform.TerraformExecutor
 		projectPath := path.Join(workingDir, job.ProjectDir)
 		if job.Terragrunt {
diff --git a/libs/orchestrator/aws.go b/libs/orchestrator/aws.go
index 45cfeedad..837f3f9fe 100644
--- a/libs/orchestrator/aws.go
+++ b/libs/orchestrator/aws.go
@@ -105,6 +105,9 @@ func (fetcher GithubAwsTokenFetcher) FetchToken(context awssdkcreds.Context) ([]
 }
 
 func GetProviderFromRole(role string) *stscreds.WebIdentityRoleProvider {
+	if role == "" {
+		return nil
+	}
 	mySession := session.Must(session.NewSession())
 	stsSTS := sts.New(mySession, &awssdk.Config{Region: awssdk.String("us-east-1")})
 	x := stscreds.NewWebIdentityRoleProviderWithOptions(stsSTS, role, "diggerSess", GithubAwsTokenFetcher{})
diff --git a/libs/orchestrator/json_models.go b/libs/orchestrator/json_models.go
index cfad685ac..c780ff8c3 100644
--- a/libs/orchestrator/json_models.go
+++ b/libs/orchestrator/json_models.go
@@ -1,6 +1,9 @@
 package orchestrator
 
-import "slices"
+import (
+	"github.com/diggerhq/digger/libs/digger_config"
+	"slices"
+)
 
 type StepJson struct {
 	Action    string   `json:"action"`
@@ -27,6 +30,8 @@ type JobJson struct {
 	Namespace         string            `json:"namespace"`
 	StateEnvVars      map[string]string `json:"stateEnvVars"`
 	CommandEnvVars    map[string]string `json:"commandEnvVars"`
+	StateRoleName     string            `json:"state_role_name"`
+	CommandRoleName   string            `json:"command_role_name"`
 }
 
 func (j *JobJson) IsPlan() bool {
@@ -37,7 +42,13 @@ func (j *JobJson) IsApply() bool {
 	return slices.Contains(j.Commands, "digger apply")
 }
 
-func JobToJson(job Job) JobJson {
+func JobToJson(job Job, project digger_config.Project) JobJson {
+	stateRole, commandRole := "", ""
+	if project.AwsRoleToAssume != nil {
+		stateRole = project.AwsRoleToAssume.State
+		commandRole = project.AwsRoleToAssume.Command
+
+	}
 	return JobJson{
 		ProjectName:       job.ProjectName,
 		ProjectDir:        job.ProjectDir,
@@ -52,24 +63,28 @@ func JobToJson(job Job) JobJson {
 		Namespace:         job.Namespace,
 		StateEnvVars:      job.StateEnvVars,
 		CommandEnvVars:    job.CommandEnvVars,
+		StateRoleName:     stateRole,
+		CommandRoleName:   commandRole,
 	}
 }
 
 func JsonToJob(jobJson JobJson) Job {
 	return Job{
-		ProjectName:       jobJson.ProjectName,
-		ProjectDir:        jobJson.ProjectDir,
-		ProjectWorkspace:  jobJson.ProjectWorkspace,
-		Terragrunt:        jobJson.Terragrunt,
-		Commands:          jobJson.Commands,
-		ApplyStage:        jsonToStage(jobJson.ApplyStage),
-		PlanStage:         jsonToStage(jobJson.PlanStage),
-		PullRequestNumber: jobJson.PullRequestNumber,
-		EventName:         jobJson.EventName,
-		RequestedBy:       jobJson.RequestedBy,
-		Namespace:         jobJson.Namespace,
-		StateEnvVars:      jobJson.StateEnvVars,
-		CommandEnvVars:    jobJson.CommandEnvVars,
+		ProjectName:        jobJson.ProjectName,
+		ProjectDir:         jobJson.ProjectDir,
+		ProjectWorkspace:   jobJson.ProjectWorkspace,
+		Terragrunt:         jobJson.Terragrunt,
+		Commands:           jobJson.Commands,
+		ApplyStage:         jsonToStage(jobJson.ApplyStage),
+		PlanStage:          jsonToStage(jobJson.PlanStage),
+		PullRequestNumber:  jobJson.PullRequestNumber,
+		EventName:          jobJson.EventName,
+		RequestedBy:        jobJson.RequestedBy,
+		Namespace:          jobJson.Namespace,
+		StateEnvVars:       jobJson.StateEnvVars,
+		CommandEnvVars:     jobJson.CommandEnvVars,
+		StateEnvProvider:   GetProviderFromRole(jobJson.StateRoleName),
+		CommandEnvProvider: GetProviderFromRole(jobJson.CommandRoleName),
 	}
 }