From b78d8d28ed169bef4b9f97c1d62a30a949b8e9d1 Mon Sep 17 00:00:00 2001
From: Daniel Hougaard <62331820+DanielHougaard@users.noreply.github.com>
Date: Tue, 12 Mar 2024 19:56:37 +0100
Subject: [PATCH] Feat: CLI support for scoped JWT tokens

---
 cli/packages/api/api.go   | 22 ++++++++++++++++++++
 cli/packages/api/model.go |  8 ++++++++
 cli/packages/cmd/login.go | 42 ++++++++++++++++++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/cli/packages/api/api.go b/cli/packages/api/api.go
index 1b69b8f23..dfb0cf7bc 100644
--- a/cli/packages/api/api.go
+++ b/cli/packages/api/api.go
@@ -164,6 +164,28 @@ func CallGetAllOrganizations(httpClient *resty.Client) (GetOrganizationsResponse
 	return orgResponse, nil
 }
 
+func CallSelectOrganization(httpClient *resty.Client, request SelectOrganizationRequest) (SelectOrganizationResponse, error) {
+	var selectOrgResponse SelectOrganizationResponse
+
+	response, err := httpClient.
+		R().
+		SetBody(request).
+		SetResult(&selectOrgResponse).
+		SetHeader("User-Agent", USER_AGENT).
+		Post(fmt.Sprintf("%v/v3/auth/select-organization", config.INFISICAL_URL))
+
+	if err != nil {
+		return SelectOrganizationResponse{}, err
+	}
+
+	if response.IsError() {
+		return SelectOrganizationResponse{}, fmt.Errorf("CallSelectOrganization: Unsuccessful response: [response=%v]", response)
+	}
+
+	return selectOrgResponse, nil
+
+}
+
 func CallGetAllWorkSpacesUserBelongsTo(httpClient *resty.Client) (GetWorkSpacesResponse, error) {
 	var workSpacesResponse GetWorkSpacesResponse
 	response, err := httpClient.
diff --git a/cli/packages/api/model.go b/cli/packages/api/model.go
index 310976078..f80c17a92 100644
--- a/cli/packages/api/model.go
+++ b/cli/packages/api/model.go
@@ -135,6 +135,14 @@ type GetOrganizationsResponse struct {
 	} `json:"organizations"`
 }
 
+type SelectOrganizationResponse struct {
+	Token string `json:"token"`
+}
+
+type SelectOrganizationRequest struct {
+	OrganizationId string `json:"organizationId"`
+}
+
 type Secret struct {
 	SecretKeyCiphertext     string `json:"secretKeyCiphertext,omitempty"`
 	SecretKeyIV             string `json:"secretKeyIV,omitempty"`
diff --git a/cli/packages/cmd/login.go b/cli/packages/cmd/login.go
index 5cff7770f..75e8f12d8 100644
--- a/cli/packages/cmd/login.go
+++ b/cli/packages/cmd/login.go
@@ -302,10 +302,12 @@ func cliDefaultLogin(userCredentialsToBeStored *models.UserCredentials) {
 		util.PrintErrorMessageAndExit("We were unable to fetch required details to complete your login. Run with -d to see more info")
 	}
 
+	newJwtToken := GetJwtTokenWithOrganizationId(loginTwoResponse.Token)
+
 	//updating usercredentials
 	userCredentialsToBeStored.Email = email
 	userCredentialsToBeStored.PrivateKey = string(decryptedPrivateKey)
-	userCredentialsToBeStored.JTWToken = loginTwoResponse.Token
+	userCredentialsToBeStored.JTWToken = newJwtToken
 }
 
 func init() {
@@ -480,6 +482,44 @@ func getFreshUserCredentials(email string, password string) (*api.GetLoginOneV2R
 	return &loginOneResponseResult, &loginTwoResponseResult, nil
 }
 
+func GetJwtTokenWithOrganizationId(oldJwtToken string) string {
+	log.Debug().Msg(fmt.Sprint("GetJwtTokenWithOrganizationId: ", "oldJwtToken", oldJwtToken))
+
+	httpClient := resty.New()
+	httpClient.SetAuthToken(oldJwtToken)
+
+	organizationResponse, err := api.CallGetAllOrganizations(httpClient)
+
+	if err != nil {
+		util.HandleError(err, "Unable to pull organizations that belong to you")
+	}
+
+	organizations := organizationResponse.Organizations
+
+	organizationNames := util.GetOrganizationsNameList(organizationResponse)
+
+	prompt := promptui.Select{
+		Label: "Which Infisical organization would you like to log into?",
+		Items: organizationNames,
+	}
+
+	index, _, err := prompt.Run()
+	if err != nil {
+		util.HandleError(err)
+	}
+
+	selectedOrganization := organizations[index]
+
+	selectedOrgRes, err := api.CallSelectOrganization(httpClient, api.SelectOrganizationRequest{OrganizationId: selectedOrganization.ID})
+
+	if err != nil {
+		util.HandleError(err)
+	}
+
+	return selectedOrgRes.Token
+
+}
+
 func userLoginMenu(currentLoggedInUserEmail string) (bool, error) {
 	label := fmt.Sprintf("Current logged in user email: %s on domain: %s", currentLoggedInUserEmail, config.INFISICAL_URL)