diff --git a/aws_okta_keyman/aws.py b/aws_okta_keyman/aws.py index 8a0fcb2..1bd2b26 100644 --- a/aws_okta_keyman/aws.py +++ b/aws_okta_keyman/aws.py @@ -127,11 +127,9 @@ def __init__(self, dir=cred_dir)) os.makedirs(cred_dir) - self.sts = boto3.client('sts') - self.profile = profile self.region = region - + self.sts = boto3.client('sts', region_name=self.region) self.assertion = SamlAssertion(assertion) self.writer = Credentials(cred_file) @@ -197,7 +195,7 @@ def available_roles(self): key=lambda k: (k['account'], k['role_name'])) return self.roles - def assume_role(self): + def assume_role(self, print_only=False): """Use the SAML Assertion to actually get the credentials. Uses the supplied (one time use!) SAML Assertion to go out to Amazon @@ -216,7 +214,11 @@ def assume_role(self): PrincipalArn=self.roles[self.role]['principle'], SAMLAssertion=self.assertion.encode()) self.creds = session['Credentials'] - self._write() + + if print_only: + self._print_creds() + else: + self._write() def _write(self): """Write out our secrets to the Credentials object.""" @@ -229,6 +231,30 @@ def _write(self): LOG.info('Session expires at {time} ⏳'.format( time=self.creds['Expiration'])) + def _print_creds(self): + """ Print out the retrieved credentials to the screen + """ + cred_str = "AWS_ACCESS_KEY_ID = {}\n".format(self.creds['AccessKeyId']) + cred_str = "{}AWS_SECRET_ACCESS_KEY = {}\n".format( + cred_str, self.creds['SecretAccessKey']) + cred_str = "{}AWS_SESSION_TOKEN = {}".format( + cred_str, self.creds['SessionToken']) + LOG.info("AWS Credentials: \n\n\n{}\n\n".format(cred_str)) + + def export_creds_to_var_string(self): + """ Export the current credentials as environment vaiables + """ + var_string = ( + "export AWS_ACCESS_KEY_ID={}; " + "export AWS_SECRET_ACCESS_KEY={}; " + "export AWS_SESSION_TOKEN={};" + ).format( + self.creds['AccessKeyId'], + self.creds['SecretAccessKey'], + self.creds['SessionToken'] + ) + return var_string + def account_ids_to_names(self, roles): """Turn account IDs into user-friendly names diff --git a/aws_okta_keyman/config.py b/aws_okta_keyman/config.py index 2a5873e..fb2be85 100644 --- a/aws_okta_keyman/config.py +++ b/aws_okta_keyman/config.py @@ -48,6 +48,9 @@ def __init__(self, argv): self.oktapreview = None self.password_cache = None self.password_reset = None + self.command = None + self.screen = None + self.region = None if len(argv) > 1: if argv[1] == 'config': @@ -65,6 +68,9 @@ def validate(self): "or as an argument") raise ValueError(err) + if self.region is None: + self.region = 'us-east-1' + if self.username is None: user = getpass.getuser() LOG.info( @@ -166,13 +172,14 @@ def main_args(arg_group, required=False): @staticmethod def optional_args(optional_args): """Define the always-optional arguments.""" - optional_args.add_argument('-u', '--username', type=str, help=( - 'Okta Login Name - either ' - 'bob@foobar.com, or just bob works too,' - ' depending on your organization ' - 'settings. Will use the current user if not' - 'specified.' - )), + optional_args.add_argument('-u', '--username', type=str, + help=( + 'Okta Login Name - either ' + 'bob@foobar.com, or just bob works too,' + ' depending on your organization ' + 'settings. Will use the current user if ' + 'not specified.' + )) optional_args.add_argument('-a', '--appid', type=str, help=( 'The "redirect link" Application ID - ' 'this can be found by mousing over the ' @@ -228,6 +235,24 @@ def optional_args(optional_args): ' if it has changed or is incorrect.' ), default=False) + optional_args.add_argument('-C', '--command', type=str, + help=( + 'Command to run with the requested ' + 'AWS keys provided as environment ' + 'variables.' + )) + optional_args.add_argument('-s', '--screen', action='store_true', + help=( + 'Print the retrieved key ' + 'only and do not write to the AWS ' + 'credentials file.' + ), + default=False) + optional_args.add_argument('-re', '--region', type=str, + help=( + 'AWS region to use for calls. ' + 'Required for GovCloud.' + )) @staticmethod def read_yaml(filename, raise_on_error=False): @@ -274,6 +299,12 @@ def write_config(self): LOG.debug("YAML being saved: {}".format(config_out)) + file_folder = os.path.dirname(os.path.abspath(file_path)) + if not os.path.exists(file_folder): + LOG.debug("Creating missin config file folder : {}".format( + file_folder)) + os.makedirs(file_folder) + with open(file_path, 'w') as outfile: yaml.safe_dump(config_out, outfile, default_flow_style=False) @@ -281,7 +312,7 @@ def write_config(self): def clean_config_for_write(config): """Remove args we don't want to save to a config file.""" ignore = ['name', 'appid', 'argv', 'writepath', 'config', 'debug', - 'oktapreview', 'password_reset'] + 'oktapreview', 'password_reset', 'command'] for var in ignore: del config[var] @@ -293,7 +324,7 @@ def clean_config_for_write(config): @staticmethod def user_input(text): """Wrap input() making testing support of py2 and py3 easier.""" - return input(text) + return input(text).strip() def interactive_config(self): """ Runs an interactive configuration to make it simpler to create diff --git a/aws_okta_keyman/duo.py b/aws_okta_keyman/duo.py index edfb3d7..98cad08 100644 --- a/aws_okta_keyman/duo.py +++ b/aws_okta_keyman/duo.py @@ -273,3 +273,4 @@ def do_redirect(self, url, sid): if 'cookie' in result['response']: return result['response']['cookie'] + return None diff --git a/aws_okta_keyman/keyman.py b/aws_okta_keyman/keyman.py index 07a6cda..ecd0e62 100644 --- a/aws_okta_keyman/keyman.py +++ b/aws_okta_keyman/keyman.py @@ -20,6 +20,7 @@ import getpass import logging +import os import sys import time import traceback @@ -108,7 +109,7 @@ def setup_logging(): @staticmethod def user_input(text): """Wrap input() making testing support of py2 and py3 easier.""" - return input(text) + return input(text).strip() def user_password(self): """Wrap getpass to simplify testing.""" @@ -351,8 +352,10 @@ def start_session(self): appid=self.config.appid) try: + self.log.info("Starting AWS session for {}".format( + self.config.region)) session = aws.Session(assertion, profile=self.config.name, - role=self.role) + role=self.role, region=self.config.region) except xml.etree.ElementTree.ParseError: self.log.error('Could not find any Role in the SAML assertion') @@ -377,11 +380,11 @@ def aws_auth_loop(self): try: session = self.start_session() - session.assume_role() + session.assume_role(self.config.screen) except aws.MultipleRoles: self.handle_multiple_roles(session) - session.assume_role() + session.assume_role(self.config.screen) except requests.exceptions.ConnectionError: self.log.warning('Connection error... will retry') time.sleep(5) @@ -404,10 +407,23 @@ def aws_auth_loop(self): self.auth_okta(state_token=err.state_token) continue - # If we're not running in re-up mode, once we have the assertion - # and creds, go ahead and quit. if not self.config.reup: - self.log.info('All done! 👍') - return None + return self.wrap_up(session) self.log.info('Reup enabled, sleeping... 💤') + + def wrap_up(self, session): + """ Execute any final steps when we're not in reup mode + + Args: + session: aws.session object + """ + if self.config.command: + command_string = "{} {}".format( + session.export_creds_to_var_string(), + self.config.command + ) + self.log.info("Running requested command...\n\n") + os.system(command_string) + else: + self.log.info('All done! 👍') diff --git a/aws_okta_keyman/metadata.py b/aws_okta_keyman/metadata.py index 909ea27..2951f1f 100644 --- a/aws_okta_keyman/metadata.py +++ b/aws_okta_keyman/metadata.py @@ -14,5 +14,5 @@ # Copyright 2018 Nathan V """Package metadata.""" -__version__ = '0.7.4' +__version__ = '0.7.5' __desc__ = 'AWS Okta Keyman' diff --git a/aws_okta_keyman/okta.py b/aws_okta_keyman/okta.py index 3a2a8b3..1e5454b 100644 --- a/aws_okta_keyman/okta.py +++ b/aws_okta_keyman/okta.py @@ -181,12 +181,13 @@ def validate_mfa(self, fid, state_token, passcode): """ if len(passcode) > 6 or len(passcode) < 5: LOG.error('Passcodes must be 5 or 6 digits') - return + return None valid = self.send_user_response(fid, state_token, passcode, 'passCode') if valid: self.set_token(valid) return True + return None def validate_answer(self, fid, state_token, answer): """Validate an Okta user with Question-based MFA. @@ -206,12 +207,13 @@ def validate_answer(self, fid, state_token, answer): """ if not answer: LOG.error('Answer cannot be blank') - return + return None valid = self.send_user_response(fid, state_token, answer, 'answer') if valid: self.set_token(valid) return True + return None def send_user_response(self, fid, state_token, user_response, resp_type): """Call Okta with a factor response and verify it. @@ -450,6 +452,7 @@ def handle_mfa_response(self, ret): return True self.handle_response_factors(response_factors, ret['stateToken']) + return None def handle_push_factors(self, factors, state_token): """Handle any push-type factors. @@ -547,9 +550,9 @@ def get_aws_apps(self): if i['appName'] == 'amazon_aws'} accounts = [] - for k, v in aws_list.items(): - appid = v.split("/", 5)[5] - accounts.append({'name': k, 'appid': appid}) + for key, val in aws_list.items(): + appid = val.split("/", 5)[5] + accounts.append({'name': key, 'appid': appid}) return accounts def mfa_callback(self, auth, verification, state_token): diff --git a/aws_okta_keyman/test/aws_test.py b/aws_okta_keyman/test/aws_test.py index b213d11..968b051 100644 --- a/aws_okta_keyman/test/aws_test.py +++ b/aws_okta_keyman/test/aws_test.py @@ -123,6 +123,8 @@ def setUp(self): self.mock_saml = self.patcher.start() self.fake_assertion = mock.MagicMock(name='FakeAssertion') self.mock_saml.return_value = self.fake_assertion + self.botopatch = mock.patch('aws_okta_keyman.aws.boto3') + self.mock_boto = self.botopatch.start() @mock.patch('os.path.expanduser') @mock.patch('os.makedirs') @@ -286,6 +288,57 @@ def test_assume_role_preset(self, mock_write): mock.call() ]) + @mock.patch('aws_okta_keyman.aws.Session._print_creds') + @mock.patch('aws_okta_keyman.aws.Session._write') + def test_assume_role_print(self, mock_write, mock_print): + assertion = mock.Mock() + assertion.roles.return_value = [{'arn': '', 'principle': ''}] + session = aws.Session('BogusAssertion') + session.role = 0 + session.roles = [{'arn': '', 'principle': ''}] + session.assertion = assertion + sts = {'Credentials': + {'AccessKeyId': 'AKI', + 'SecretAccessKey': 'squirrel', + 'SessionToken': 'token', + 'Expiration': 'never' + }} + session.sts = mock.Mock() + session.sts.assume_role_with_saml.return_value = sts + + session.assume_role(print_only=True) + + assert not mock_write.called + assert mock_print.called + + @mock.patch('aws_okta_keyman.aws.LOG') + def test_print_creds(self, log_mock): + session = aws.Session('BogusAssertion') + expected = ( + 'AWS Credentials: \n\n\n' + 'AWS_ACCESS_KEY_ID = None\n' + 'AWS_SECRET_ACCESS_KEY = None\n' + 'AWS_SESSION_TOKEN = None\n\n' + ) + + session._print_creds() + + log_mock.assert_has_calls([ + mock.call.info(expected) + ]) + + def test_export_creds_to_var_string(self): + session = aws.Session('BogusAssertion') + expected = ( + 'export AWS_ACCESS_KEY_ID=None; ' + 'export AWS_SECRET_ACCESS_KEY=None; ' + 'export AWS_SESSION_TOKEN=None;' + ) + + ret = session.export_creds_to_var_string() + + self.assertEqual(ret, expected) + def test_available_roles(self): roles = [{'role': '::::1:role/role', 'principle': ''}, {'role': '::::1:role/role', 'principle': ''}] diff --git a/aws_okta_keyman/test/config_test.py b/aws_okta_keyman/test/config_test.py index 96f7a30..26f764e 100644 --- a/aws_okta_keyman/test/config_test.py +++ b/aws_okta_keyman/test/config_test.py @@ -497,6 +497,25 @@ def test_write_config_path_expansion(self): m.assert_has_calls([mock.call(expected_path, 'w')]) + @mock.patch('aws_okta_keyman.config.os') + def test_write_config_path_create_when_missing(self, os_mock): + config = Config(['aws_okta_keyman.py']) + config.clean_config_for_write = mock.MagicMock() + config.clean_config_for_write.return_value = {} + config.read_yaml = mock.MagicMock() + config.read_yaml.return_value = {} + folderpath = '/home/user/.config/' + os_mock.path.dirname.return_value = folderpath + os_mock.path.exists.return_value = False + + m = mock.mock_open() + with mock.patch('aws_okta_keyman.config.open', m): + config.write_config() + + os_mock.assert_has_calls([ + mock.call.makedirs(folderpath) + ]) + def test_clean_config_for_write(self): config_in = { 'name': 'foo', @@ -508,7 +527,8 @@ def test_clean_config_for_write(self): 'oktapreview': 'foo', 'accounts': None, 'shouldstillbehere': 'woohoo', - 'password_reset': True + 'password_reset': True, + 'command': None } config_out = { 'shouldstillbehere': 'woohoo' @@ -531,7 +551,8 @@ def test_clean_config_for_write_with_accounts(self): 'oktapreview': 'foo', 'accounts': accounts, 'shouldstillbehere': 'woohoo', - 'password_reset': True + 'password_reset': True, + 'command': None } config_out = { 'accounts': accounts, @@ -542,7 +563,7 @@ def test_clean_config_for_write_with_accounts(self): @mock.patch('aws_okta_keyman.config.input') def test_user_input(self, input_mock): - input_mock.return_value = 'test' + input_mock.return_value = ' test ' self.assertEqual('test', Config.user_input('input test')) @mock.patch('aws_okta_keyman.config.getpass') diff --git a/aws_okta_keyman/test/duo_test.py b/aws_okta_keyman/test/duo_test.py index 669c26c..b97243c 100644 --- a/aws_okta_keyman/test/duo_test.py +++ b/aws_okta_keyman/test/duo_test.py @@ -302,6 +302,17 @@ def test_do_redirect_failure(self): with self.assertRaises(Exception): duo_test.do_redirect('url', 'sid') + def test_do_redirect_missing_cookie(self): + duo_test = duo.Duo(DETAILS, 'token') + duo_test.session = mock.MagicMock() + json_ok = {'response': {'crumbs': 'yum'}, 'stat': 'OK'} + headers = {'Location': 'https://someurl/foo?sid=somesid'} + duo_test.session.post.return_value = MockResponse( + headers, 200, json_ok) + ret = duo_test.do_redirect('url', 'sid') + + self.assertEqual(ret, None) + class TestQuietHandler(unittest.TestCase): # noinspection PyArgumentList diff --git a/aws_okta_keyman/test/keyman_test.py b/aws_okta_keyman/test/keyman_test.py index 39fff73..8592f93 100644 --- a/aws_okta_keyman/test/keyman_test.py +++ b/aws_okta_keyman/test/keyman_test.py @@ -125,7 +125,7 @@ def test_main_aws_auth_error(self, _config_mock): @mock.patch('aws_okta_keyman.keyman.input') def test_user_input(self, input_mock): - input_mock.return_value = 'test' + input_mock.return_value = ' test ' self.assertEqual('test', Keyman.user_input('input test')) @@ -600,7 +600,8 @@ def test_start_session(self, aws_mock, _config_mock): mock.call.get_assertion(appid=mock.ANY) ]) aws_mock.assert_has_calls([ - mock.call.Session('assertion', profile=mock.ANY, role=None) + mock.call.Session('assertion', profile=mock.ANY, role=None, + region=mock.ANY) ]) @mock.patch('aws_okta_keyman.keyman.Config') @@ -628,13 +629,15 @@ def test_aws_auth_loop(self, config_mock): config_mock().reup = False keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz']) keyman.start_session = mock.MagicMock() + keyman.wrap_up = mock.MagicMock() keyman.aws_auth_loop() keyman.start_session.assert_has_calls([ mock.call(), - mock.call().assume_role() + mock.call().assume_role(mock.ANY) ]) + assert keyman.wrap_up.called @mock.patch('time.sleep', return_value=None) @mock.patch('aws_okta_keyman.keyman.Config') @@ -671,6 +674,7 @@ def test_aws_auth_loop_continue(self, config_mock, sleep_mock): @mock.patch('aws_okta_keyman.keyman.Config') def test_aws_auth_loop_multirole(self, config_mock): config_mock().reup = False + config_mock().screen = False keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz']) session_instance = mock.MagicMock() session_instance.assume_role.side_effect = aws.MultipleRoles @@ -682,8 +686,8 @@ def test_aws_auth_loop_multirole(self, config_mock): keyman.aws_auth_loop() session_instance.assume_role.assert_has_calls([ - mock.call(), - mock.call(), + mock.call(False), + mock.call(False), ]) keyman.handle_multiple_roles.assert_has_calls([ mock.call(mock.ANY) @@ -735,3 +739,32 @@ def test_aws_auth_loop_exception(self, config_mock, _sleep_mock): keyman.aws_auth_loop() assert keyman.start_session.called + + @mock.patch('aws_okta_keyman.keyman.Config') + def test_wrap_up_no_command(self, config_mock): + config_mock().command = None + keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz']) + keyman.log = mock.MagicMock() + + keyman.wrap_up(None) + + keyman.log.assert_has_calls([ + mock.call.info('All done! 👍') + ]) + + @mock.patch('aws_okta_keyman.keyman.os') + @mock.patch('aws_okta_keyman.keyman.Config') + def test_wrap_up_with_command(self, config_mock, os_mock): + config_mock().command = 'echo w00t' + keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz']) + fake_session = mock.MagicMock() + fake_session.export_creds_to_var_string.return_value = 'foo' + + keyman.wrap_up(fake_session) + + fake_session.assert_has_calls([ + mock.call.export_creds_to_var_string() + ]) + os_mock.assert_has_calls([ + mock.call.system('foo echo w00t') + ]) diff --git a/aws_okta_keyman/test/okta_test.py b/aws_okta_keyman/test/okta_test.py index 3cf38ac..dc99754 100644 --- a/aws_okta_keyman/test/okta_test.py +++ b/aws_okta_keyman/test/okta_test.py @@ -649,6 +649,16 @@ def test_handle_mfa_response_trigger_sms_otp(self): 'token') ]) + def test_handle_mfa_response_returns_none(self): + client = okta.Okta('organization', 'username', 'password') + client.handle_push_factors = mock.MagicMock() + client.handle_push_factors.return_value = False + client.handle_response_factors = mock.MagicMock() + + ret = client.handle_mfa_response(MFA_CHALLENGE_SMS_OTP) + + self.assertEqual(ret, None) + def test_handle_mfa_response_unsupported(self): client = okta.Okta('organization', 'username', 'password') client.handle_push_factors = mock.MagicMock()