Merge pull request #216 from april/master

Fixes related to maws prompt
mozilla-iam · May 13, 2020 · 929fc19 · 929fc19
2 parents c858683 + b3140cc
commit 929fc19
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 15 deletions.
diff --git a/mozilla_aws_cli/login.py b/mozilla_aws_cli/login.py
@@ -26,7 +26,7 @@
     base64_without_padding,
     exit_sigint,
     generate_challenge,
-    role_arn_to_profile_name,
+    role_arn_to_display_name,
     STSWarning
 )
 
@@ -411,10 +411,14 @@ def exchange_token_for_credentials(self):
     def print_output(self):
         # TODO: Create a global config object?
         if self.credentials is not None:
-            if self.profile_name is None:
-                self.profile_name = role_arn_to_profile_name(
-                    self.role_arn, self.role_map)
             output_map = {}
+
+            # get a role name (for the command line), and use it for the profile
+            # name if it wasn't manually overridden
+            self.role = role_arn_to_display_name(self.role_arn, self.role_map)
+            if self.profile_name is None:
+                self.profile_name = self.role
+
             if self.output == "envvar":
                 output_map.update(
                     {ENV_VARIABLE_NAME_MAP[x]: self.credentials[x]
@@ -423,7 +427,7 @@ def print_output(self):
                 output_map.update({
                     "AWS_PROFILE": None,
                     "AWS_SHARED_CREDENTIALS_FILE": None,
-                    "MAWS_PROMPT": self.profile_name})
+                    "MAWS_PROMPT": self.role})
             elif self.output == "shared":
                 # Write the credentials
                 path = write_aws_shared_credentials(
@@ -433,21 +437,19 @@ def print_output(self):
                     output_map.update({
                         "AWS_PROFILE": self.profile_name,
                         "AWS_SHARED_CREDENTIALS_FILE": path,
-                        "MAWS_PROMPT": self.profile_name})
+                        "MAWS_PROMPT": self.role})
                     output_map.update({
                         x: None for x in ENV_VARIABLE_NAME_MAP.values()})
             elif self.output == "awscli":
                 # Call into aws a bunch of times
                 if write_aws_cli_credentials(self.profile_name,
                                              self.credentials):
-                    if self.profile_name != "default":
-                        output_map.update({
-                            "AWS_PROFILE": self.profile_name,
-                            "AWS_SHARED_CREDENTIALS_FILE": None,
-                            "MAWS_PROMPT": self.profile_name
-                        })
-                        output_map.update({
-                            x: None for x in ENV_VARIABLE_NAME_MAP.values()})
+                    output_map.update({
+                        "AWS_PROFILE": self.profile_name,
+                        "AWS_SHARED_CREDENTIALS_FILE": None,
+                        "MAWS_PROMPT": self.role})
+                    output_map.update({
+                        x: None for x in ENV_VARIABLE_NAME_MAP.values()})
                 else:
                     logger.error("Unable to write credentials with aws-cli.")
             elif self.output == "boto":

diff --git a/mozilla_aws_cli/utils.py b/mozilla_aws_cli/utils.py
@@ -30,7 +30,7 @@ def generate_challenge(code_verifier):
         hashlib.sha256(code_verifier.encode()).digest())
 
 
-def role_arn_to_profile_name(role_arn, role_map):
+def role_arn_to_display_name(role_arn, role_map):
     if not role_map:
         role_map = {}