Add an option to set the profile name, both for shared/awscli output

mozilla_aws_cli/cache.py

@@ -124,19 +124,18 @@ def disable_caching(*args, **kwargs):
 
 
 @_requires_safe_cache_dir
-def write_aws_cli_credentials(credentials, role_arn, role_map):
-    # such as infosec-somerole
-    profile = role_arn_to_profile_name(role_arn, role_map)
-
+def write_aws_cli_credentials(profile, credentials, role_arn, role_map):
     # We call aws a bunch of times, getting all the return values
     retval = 0
 
     # Update the values
     for cred_key, aws_key in viewitems(CREDENTIALS_TO_AWS_MAP):
         if cred_key in credentials:
             process = ["aws", "configure", "set",
-                       aws_key, credentials[cred_key],
-                       "--profile", profile]
+                       aws_key, credentials[cred_key]]
+
+            if profile != "default":
+                process += ["--profile", profile]
 
             _retval = subprocess.call(process)
             retval = retval | _retval
@@ -182,15 +181,12 @@ def read_aws_shared_credentials():
 
 
 @_requires_safe_cache_dir
-def write_aws_shared_credentials(credentials, role_arn, role_map=None):
+def write_aws_shared_credentials(profile, credentials, role_arn, role_map=None):
     path = os.path.join(DOT_DIR, "credentials")
 
     # Try to read in the existing credentials
     config = read_aws_shared_credentials()
 
-    # such as infosec-somerole
-    profile = role_arn_to_profile_name(role_arn, role_map)
-
     # Add all the new credentials to the config object
     if not config.has_section(profile):
         config.add_section(profile)

mozilla_aws_cli/cli.py

@@ -152,6 +152,9 @@ def validate_disable_caching(ctx, param, disabled):
     help="How to output the AWS API keys",
     callback=validate_output
 )
+@click.option("--profile",
+              metavar="<profile>",
+              help="Override profile name used with `awscli` or `shared` output")
 @click.option(
     "-r",
     "--role-arn",
@@ -160,12 +163,14 @@ def validate_disable_caching(ctx, param, disabled):
     callback=validate_arn)
 @click.option("-v", "--verbose", is_flag=True, help="Print debugging messages")
 @click.option("-w", "--web-console", is_flag=True, help="Open AWS web console")
-def main(batch, config, no_cache, output, role_arn, verbose, web_console):
+def main(batch, config, no_cache, output,
+         profile, role_arn, verbose, web_console):
     """Fetch AWS API Keys using SSO web login"""
     if verbose:
         logger.setLevel(logging.DEBUG)
 
     # Order of precedence : output, config["output"], "envvar"
+    profile = config.get("profile") if profile is None else profile
     config["output"] = output if output is not None else config.get("output", "envvar")
     config["openid-configuration"] = requests.get(config["well_known_url"]).json()
     config["jwks"] = requests.get(config["openid-configuration"]["jwks_uri"]).json()
@@ -182,6 +187,7 @@ def main(batch, config, no_cache, output, role_arn, verbose, web_console):
         jwks=config["jwks"],
         openid_configuration=config["openid-configuration"],
         config=config,
+        profile_name=profile,
         role_arn=role_arn,
         scope=config.get("scope"),
         token_endpoint=config["openid-configuration"]["token_endpoint"],

mozilla_aws_cli/login.py

@@ -58,6 +58,7 @@ def __init__(
         jwks=None,
         openid_configuration=None,
         config=None,
+        profile_name=None,
         role_arn=None,
         scope="openid",
         token_endpoint="https://auth.mozilla.auth0.com/oauth/token",
@@ -76,10 +77,10 @@ def __init__(
         self.client_id = client_id
         self.code_verifier = base64_without_padding(os.urandom(32))
         self.code_challenge = generate_challenge(self.code_verifier)
+        self.config = {} if config is None else config
         self.idtoken_for_roles_url = idtoken_for_roles_url
         self.jwks = jwks
         self.openid_configuration = openid_configuration
-        self.config = {} if config is None else config
         self.output = self.config.get("output", "envvar")
         self.print_role_arn = self.config.get("print_role_arn", True)
         self.role = None
@@ -91,9 +92,8 @@ def __init__(
         self.oidc_state = self.id + "-" + base64_without_padding(os.urandom(32))
 
         # URL of the OIDC token endpoint obtained from the discovery document
-        self.token_endpoint = token_endpoint
         self.batch = batch
-        self.web_console = web_console
+        self.token_endpoint = token_endpoint
         self.issuer_domain = issuer_domain
 
         # Whether or not we have opened a browser tab
@@ -110,6 +110,12 @@ def __init__(
         self.last_state_check = None
         self.max_sleep_no_state_check = 2  # seconds
 
+        # Whether we should open the AWS web console or not
+        self.web_console = web_console
+
+        # If we're using the AWS CLI output, what profile should we use
+        self.profile_name = profile_name
+
         # This used by the web application to poll the login state
         self.state = "pending"
         self.web_state = {
@@ -377,8 +383,9 @@ def exchange_token_for_credentials(self):
     def print_output(self):
         # TODO: Create a global config object?
         if self.credentials is not None:
-            profile_name = role_arn_to_profile_name(
-                self.role_arn, self.role_map)
+            if self.profile_name is None:
+                self.profile_name = role_arn_to_profile_name(
+                    self.role_arn, self.role_map)
             output_map = {}
             if self.output == "envvar":
                 output_map.update({ENV_VARIABLE_NAME_MAP[x]: self.credentials[x]
@@ -387,24 +394,30 @@ def print_output(self):
             elif self.output == "shared":
                 # Write the credentials
                 path = write_aws_shared_credentials(
+                    self.profile_name,
                     self.credentials,
                     self.role_arn,
                     self.role_map)
                 if path:
                     output_map.update({
-                        'AWS_PROFILE': profile_name,
+                        'AWS_PROFILE': self.profile_name,
                         'AWS_SHARED_CREDENTIALS_FILE': path})
             elif self.output == "awscli":
                 # Call into aws a bunch of times
-                if write_aws_cli_credentials(self.credentials,
+                if write_aws_cli_credentials(self.profile_name,
+                                             self.credentials,
                                              self.role_arn,
                                              self.role_map):
-                    output_map.update({'AWS_PROFILE': profile_name})
+                    if self.profile_name != "default":
+                        output_map.update({'AWS_PROFILE': self.profile_name})
                 else:
                     logger.error('Unable to write credentials with aws-cli.')
             else:
                 raise ValueError('Output setting unknown : {}'.format(self.output))
-            print(output_set_env_vars(output_map))
+
+            if output_map:
+                print(output_set_env_vars(output_map))
+
             if self.print_role_arn:
                 print("Environment variables set for role {}".format(
                     self.role_arn), file=sys.stderr)