From 17b1dd9d697e20319243a150d413b5ab4797515d Mon Sep 17 00:00:00 2001 From: April King Date: Thu, 5 Dec 2019 15:59:37 -0600 Subject: [PATCH] Add an option to set the profile name, both for shared and awscli output. Fixes #132. --- mozilla_aws_cli/cache.py | 16 ++++++---------- mozilla_aws_cli/cli.py | 8 +++++++- mozilla_aws_cli/login.py | 31 ++++++++++++++++++++++--------- 3 files changed, 35 insertions(+), 20 deletions(-) diff --git a/mozilla_aws_cli/cache.py b/mozilla_aws_cli/cache.py index b02ffb2..3878be3 100644 --- a/mozilla_aws_cli/cache.py +++ b/mozilla_aws_cli/cache.py @@ -124,10 +124,7 @@ 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 @@ -135,8 +132,10 @@ def write_aws_cli_credentials(credentials, role_arn, role_map): 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) diff --git a/mozilla_aws_cli/cli.py b/mozilla_aws_cli/cli.py index 633797a..283878f 100644 --- a/mozilla_aws_cli/cli.py +++ b/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="", + 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"], diff --git a/mozilla_aws_cli/login.py b/mozilla_aws_cli/login.py index 47cb340..0aaec44 100644 --- a/mozilla_aws_cli/login.py +++ b/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)