Fix get API key for user org

Author: thijslemmens

scalewayclient.go

@@ -24,17 +24,13 @@ const (
 
 type ScalewayClient struct {
 	httpClient            http.Client
-	organizationToProfile map[string]OrganizationProfile
+	organizationToJwt     map[string]string
+	organizationToUserId  map[string]string
 	defaultOrganizationId string
 	jti                   string
 	issuer                string
 }
 
-type OrganizationProfile struct {
-	jwt    string
-	userId string
-}
-
 type ApiKey struct {
 	AccessKey string `json:"access_key"`
 	SecretKey string `json:"secret_key"`
@@ -48,8 +44,9 @@ type PostBody struct {
 
 func NewScalewayClient(email, password, token string) (*ScalewayClient, error) {
 	client := ScalewayClient{
-		httpClient:            http.Client{},
-		organizationToProfile: make(map[string]OrganizationProfile),
+		httpClient:           http.Client{},
+		organizationToJwt:    make(map[string]string),
+		organizationToUserId: make(map[string]string),
 	}
 
 	postBody := PostBody{
@@ -70,21 +67,27 @@ func NewScalewayClient(email, password, token string) (*ScalewayClient, error) {
 	client.jti = jsonResponse["jwt"].(map[string]interface{})["jti"].(string)
 	client.issuer = jsonResponse["jwt"].(map[string]interface{})["issuer_id"].(string)
 
-	client.organizationToProfile[client.defaultOrganizationId] = OrganizationProfile{
-		jwt: jsonResponse["token"].(string),
+	jwt := jsonResponse["token"].(string)
+
+	jsonResponse, err = client.sendRequest(IAM_USERS_API+client.issuer, nil, jwt, "GET")
+	if err != nil {
+		return nil, err
 	}
+	client.defaultOrganizationId = jsonResponse["organization_id"].(string)
+	client.organizationToJwt[client.defaultOrganizationId] = jwt
+	client.organizationToUserId[client.defaultOrganizationId] = client.issuer
 
 	return &client, nil
 }
 
 func (client *ScalewayClient) ListOrganizations() (map[string]string, error) {
-	jsonResponse, err := client.sendRequest(IAM_USERS_API+client.issuer, nil, client.getOrCreateOrganizationProfile(client.defaultOrganizationId).jwt, "GET")
+	jsonResponse, err := client.sendRequest(IAM_USERS_API+client.issuer, nil, client.getOrCreateOrganizationJwt(client.defaultOrganizationId), "GET")
 	if err != nil {
 		return nil, err
 	}
 	accountRootUserId := jsonResponse["account_root_user_id"].(string)
 
-	jsonResponse, err = client.sendRequest(ACCOUNT_USERS_API+accountRootUserId, nil, client.getOrCreateOrganizationProfile(client.defaultOrganizationId).jwt, "GET")
+	jsonResponse, err = client.sendRequest(ACCOUNT_USERS_API+accountRootUserId, nil, client.getOrCreateOrganizationJwt(client.defaultOrganizationId), "GET")
 	if err != nil {
 		return nil, err
 	}
@@ -102,11 +105,11 @@ func (sc *ScalewayClient) CreateAPIKey(organizationId string, duration time.Dura
 	jsonBody, _ := json.Marshal(map[string]string{
 		"default_project_id": organizationId,
 		"description":        "generated by scw-2fa-init",
-		"user_id":            sc.getOrCreateOrganizationProfile(organizationId).userId,
+		"user_id":            sc.getOrCreateOrganizationUserId(organizationId),
 		"expiresAt":          time.Now().Add(duration).UTC().Format(time.RFC3339),
 	})
 
-	jsonResponse, err := sc.sendRequest(API_KEYS_API, jsonBody, sc.getOrCreateOrganizationProfile(organizationId).jwt, http.MethodPost)
+	jsonResponse, err := sc.sendRequest(API_KEYS_API, jsonBody, sc.getOrCreateOrganizationJwt(organizationId), http.MethodPost)
 	if err != nil {
 		return nil, err
 	}
@@ -117,21 +120,27 @@ func (sc *ScalewayClient) CreateAPIKey(organizationId string, duration time.Dura
 	return &ApiKey{AccessKey: accessKey, SecretKey: secretKey}, nil
 }
 
-func (client *ScalewayClient) getOrCreateOrganizationProfile(organizationId string) OrganizationProfile {
-	organizationProfile, ok := client.organizationToProfile[organizationId]
+func (client *ScalewayClient) getOrCreateOrganizationUserId(organizationId string) string {
+	userId, ok := client.organizationToUserId[organizationId]
 	if !ok {
-		jsonBody, _ := json.Marshal(map[string]string{
-			"organization_id": organizationId,
-		})
-
-		jsonResponse, _ := client.sendRequest(fmt.Sprintf(SWITCH_ORGANIZATION_API, client.jti), jsonBody, client.getOrCreateOrganizationProfile(client.defaultOrganizationId).jwt, "POST")
-		organizationProfile = OrganizationProfile{
-			jwt:    jsonResponse["token"].(string),
-			userId: jsonResponse["user_id"].(string),
-		}
-		client.organizationToProfile[organizationId] = organizationProfile
-	}
-	return organizationProfile
+		client.completeOrganizationData(organizationId)
+	}
+	userId, _ = client.organizationToUserId[organizationId]
+	return userId
+}
+
+func (client *ScalewayClient) completeOrganizationData(organizationId string) {
+	jsonBody, _ := json.Marshal(map[string]string{
+		"organization_id": organizationId,
+	})
+
+	jsonResponse, err := client.sendRequest(fmt.Sprintf(SWITCH_ORGANIZATION_API, client.jti), jsonBody, client.getOrCreateOrganizationJwt(client.defaultOrganizationId), "POST")
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	client.organizationToUserId[organizationId] = jsonResponse["user_id"].(string)
+	client.organizationToJwt[organizationId] = jsonResponse["token"].(string)
 }
 
 func (client *ScalewayClient) sendRequest(url string, body []byte, jwt string, method string) (map[string]interface{}, error) {
@@ -150,13 +159,23 @@ func (client *ScalewayClient) sendRequest(url string, body []byte, jwt string, m
 		return nil, err
 	}
 
+	responseBody, _ := ioutil.ReadAll(response.Body)
+
 	if response.StatusCode >= 400 {
-		return nil, errors.New("Request failed with status code: " + strconv.Itoa(response.StatusCode))
+		return nil, errors.New("Request failed with status code: " + strconv.Itoa(response.StatusCode) + "\n" + string(responseBody))
 	}
 
-	responseBody, _ := ioutil.ReadAll(response.Body)
 	var jsonResponse map[string]interface{}
 	json.Unmarshal(responseBody, &jsonResponse)
 
 	return jsonResponse, nil
 }
+
+func (client *ScalewayClient) getOrCreateOrganizationJwt(organizationId string) string {
+	jwt, ok := client.organizationToJwt[organizationId]
+	if !ok {
+		client.completeOrganizationData(organizationId)
+	}
+	jwt, _ = client.organizationToJwt[organizationId]
+	return jwt
+}