def test_assume_role(self, mock_write): mock_write.return_value = None assertion = mock.Mock() assertion.roles.return_value = [{'role': '', 'principle': ''}] session = aws.Session('BogusAssertion') 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 ret = session.assume_role() self.assertEquals(None, ret) self.assertEquals('AKI', session.creds['AccessKeyId']) self.assertEquals('squirrel', session.creds['SecretAccessKey']) self.assertEquals('token', session.creds['SessionToken']) self.assertEquals('never', session.creds['Expiration']) # Verify _write is called correctly mock_write.assert_has_calls([ mock.call() ])
def test_set_role(self): assertion = mock.Mock() assertion.roles.return_value = [{'role': '', 'principle': ''}] session = aws.Session('BogusAssertion') session.assertion = assertion session.set_role('0') self.assertEquals('', session.role['role'])
def test_generate_aws_console_url(self, requests_mock): session = aws.Session("BogusAssertion") session.duration = 3600 session.creds = { "AccessKeyId": "AKI", "SecretAccessKey": "squirrel", "SessionToken": "token", "Expiration": "never", } resp_mock = mock.MagicMock() resp_mock.json.return_value = {"SigninToken": "baz"} requests_mock.get.return_value = resp_mock issuer = "https://ex.okta.com/foo/bar" ret = session.generate_aws_console_url(issuer) expected = ( "https://signin.aws.amazon.com/federation?Action=login&Issuer=" "https://ex.okta.com/foo/bar&Destination=" "https%3A//console.aws.amazon.com/&SigninToken=baz") self.assertEqual(ret, expected) # mock.ANY required for the session due to Python 3.5 behavior requests_mock.assert_has_calls([ mock.call.get( "https://signin.aws.amazon.com/federation", params={ "Action": "getSigninToken", "SessionDuration": 3600, "Session": mock.ANY, }, ), mock.call.get().json(), ], )
def test_assume_role_multiple(self, mock_write): mock_write.return_value = None assertion = mock.Mock() roles = [ { "arn": "1", "principle": "" }, { "arn": "2", "principle": "" }, ] assertion.roles.return_value = roles session = aws.Session("BogusAssertion") 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 with self.assertRaises(aws.MultipleRoles): session.assume_role()
def test_available_roles(self): roles = [{ 'role': '::::1:role/role', 'principle': '' }, { 'role': '::::1:role/role', 'principle': '' }] session = aws.Session('BogusAssertion') session.assertion = mock.MagicMock() session.assertion.roles.return_value = roles expected = [{ 'account': '1', 'role_name': 'role', 'principle': '', 'arn': '::::1:role/role' }, { 'account': '1', 'role_name': 'role', 'principle': '', 'arn': '::::1:role/role' }] result = session.available_roles() print(result) self.assertEqual(expected, result)
def test_generate_aws_console_url(self, requests_mock): session = aws.Session('BogusAssertion') session.duration = 3600 session.creds = { 'AccessKeyId': 'AKI', 'SecretAccessKey': 'squirrel', 'SessionToken': 'token', 'Expiration': 'never' } resp_mock = mock.MagicMock() resp_mock.json.return_value = {'SigninToken': 'baz'} requests_mock.get.return_value = resp_mock issuer = 'https://ex.okta.com/foo/bar' ret = session.generate_aws_console_url(issuer) expected = ( "https://signin.aws.amazon.com/federation?Action=login&Issuer=" "https://ex.okta.com/foo/bar&Destination=" "https%3A//console.aws.amazon.com/&SigninToken=baz") self.assertEqual(ret, expected) # mock.ANY required for the session due to Python 3.5 behavior requests_mock.assert_has_calls([ mock.call.get('https://signin.aws.amazon.com/federation', params={ 'Action': 'getSigninToken', 'SessionDuration': 3600, 'Session': mock.ANY }), mock.call.get().json() ])
def test_account_ids_to_names_map(self): session = aws.Session("BogusAssertion") session.get_account_name_map = mock.MagicMock() account_map = {"1": "One", "2": "Two"} session.get_account_name_map.return_value = account_map roles = [ { "role": "role", "account": "1" }, { "role": "role", "account": "2" }, ] expected = [ { "account": "One", "role": "role" }, { "account": "Two", "role": "role" }, ] ret = session.account_ids_to_names(roles) self.assertEqual(ret, expected)
def test_init_folder_exists(self, exists_mock, _makedirs_mock, expuser_mock): exists_mock.return_value = True expuser_mock.return_value = '/home/fakeuser' aws.Session('BogusAssertion') exists_mock.assert_has_calls([mock.call('/home/fakeuser')])
def test_account_ids_to_names_call_failed(self): session = aws.Session('BogusAssertion') session.get_account_name_map = mock.MagicMock() session.get_account_name_map.side_effect = Exception() roles = [{'role': '::::1:role'}, {'role': '::::2:role'}] ret = session.account_ids_to_names(roles) self.assertEqual(ret, [{'role': '::::1:role'}, {'role': '::::2:role'}])
def test_assume_role_preset(self, mock_write): mock_write.return_value = None assertion = mock.Mock() roles = [ { "role": "::::1:role/role1", "principle": "", "arn": "1" }, { "role": "::::1:role/role2", "principle": "", "arn": "2" }, { "role": "::::1:role/role3", "principle": "", "arn": "3" }, ] assertion.roles.return_value = roles session = aws.Session("BogusAssertion") session.role = 1 session.roles = roles 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 ret = session.assume_role() self.assertEqual(None, ret) self.assertEqual("AKI", session.creds["AccessKeyId"]) self.assertEqual("squirrel", session.creds["SecretAccessKey"]) self.assertEqual("token", session.creds["SessionToken"]) self.assertEqual("never", session.creds["Expiration"]) # Verify _write is called correctly mock_write.assert_has_calls([ mock.call(), ], ) session.sts.assert_has_calls([ mock.call.assume_role_with_saml( RoleArn="2", PrincipalArn="", SAMLAssertion=mock.ANY, DurationSeconds=3600, ), ], )
def test_available_roles(self): assertion = mock.Mock() roles = [{'role': '1', 'principle': ''}, {'role': '2', 'principle': ''}] assertion.roles.return_value = roles session = aws.Session('BogusAssertion') session.assertion = assertion result = session.available_roles() self.assertEquals(roles, result)
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 start_session(self): """Initialize AWS session object.""" try: assertion = self.okta_client.get_assertion(appid=self.config.appid, apptype='amazon_aws') except okta.UnknownError: sys.exit(1) return aws.Session(assertion, profile=self.config.name)
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_is_valid_false_missing_expiration(self): session = aws.Session('BogusAssertion') # Set expiration to None like we failed to set the value session.expiration = None # Should return False - the time comparison isn't possible if # expiration hasn't been set yet. ret = session.is_valid self.assertEquals(False, ret)
def test_get_account_name_map_error(self): def post(*args, **kwargs): return MockResponse('text', True) session = aws.Session('BogusAssertion') session.assertion = mock.MagicMock() session.assertion.encode.response_value = '' with mock.patch('aws_okta_keyman.aws.requests.post', side_effect=post): with self.assertRaises(Exception): session.get_account_name_map()
def test_init_folder_exists( self, exists_mock, _makedirs_mock, expuser_mock, ): exists_mock.return_value = True expuser_mock.return_value = "/home/fakeuser" aws.Session("BogusAssertion") exists_mock.assert_has_calls([mock.call("/home/fakeuser")])
def test_init_with_duration( self, exists_mock, _makedirs_mock, expuser_mock, ): exists_mock.return_value = True expuser_mock.return_value = "/home/fakeuser" aws.Session("BogusAssertion", session_duration=6000) exists_mock.assert_has_calls([mock.call("/home/fakeuser")])
def test_write(self, mock_add_profile): session = aws.Session('BogusAssertion') ret = session._write() self.assertEquals(None, ret) # Verify add_profile is called with the correct args creds = {'AccessKeyId': None, 'SecretAccessKey': None, 'SessionToken': None, 'Expiration': None} mock_add_profile.assert_has_calls([ mock.call(creds=creds, name='default', region='us-east-1') ])
def test_account_ids_to_names_map(self): session = aws.Session('BogusAssertion') session.get_account_name_map = mock.MagicMock() account_map = {'1': 'One', '2': 'Two'} session.get_account_name_map.return_value = account_map roles = [{'role': 'role', 'account': '1'}, {'role': 'role', 'account': '2'}] expected = [{'account': 'One', 'role': 'role'}, {'account': 'Two', 'role': 'role'}] ret = session.account_ids_to_names(roles) self.assertEqual(ret, expected)
def test_available_roles_multiple_accounts(self): roles = [ { "role": "::::1:role/role", "principle": "" }, { "role": "::::2:role/role", "principle": "" }, ] roles_full = [ { "account": "1", "role_name": "role", "arn": "::::1:role/role", "principle": "", }, { "account": "2", "role_name": "role", "arn": "::::2:role/role", "principle": "", }, ] session = aws.Session("BogusAssertion") session.assertion = mock.MagicMock() session.assertion.roles.return_value = roles session.account_ids_to_names = mock.MagicMock() session.account_ids_to_names.return_value = roles_full expected = [ { "account": "1", "role_name": "role", "principle": "", "arn": "::::1:role/role", "roleIdx": 0, }, { "account": "2", "role_name": "role", "principle": "", "arn": "::::2:role/role", "roleIdx": 1, }, ] result = session.available_roles() print(result) self.assertEqual(expected, result)
def test_get_account_name_map(self): def post(*args, **kwargs): return MockResponse('html') session = aws.Session('BogusAssertion') session.assertion = mock.MagicMock() session.assertion.encode.response_value = '' session.account_names_from_html = mock.MagicMock() session.account_names_from_html.return_value = {} with mock.patch('aws_okta_keyman.aws.requests.post', side_effect=post): ret = session.get_account_name_map() self.assertEqual(ret, {}) session.account_names_from_html.assert_has_calls([mock.call('html')])
def test_is_valid_true(self): session = aws.Session('BogusAssertion') # Mock out the expiration time to 4:10PM UTC expir_mock = datetime.datetime(2017, 7, 25, 16, 10, 00, 000000) # Now set our current time to 3:55PM UTC mock_now = datetime.datetime(2017, 7, 25, 15, 55, 00, 000000) # Should return True - more than 600 seconds to expiration with mock.patch('datetime.datetime') as dt_mock: dt_mock.utcnow.return_value = mock_now dt_mock.strptime.return_value = expir_mock ret = session.is_valid self.assertEquals(True, ret)
def test_is_valid_false(self): session = aws.Session('BogusAssertion') # Mock out the expiration time to 4:10PM UTC expir_mock = datetime.datetime(2017, 7, 25, 16, 10, 00, 000000) # Now set our current time to 4:05PM UTC mock_now = datetime.datetime(2017, 7, 25, 16, 4, 00, 000000) # Should return False - less than 600 seconds away from expiration with mock.patch('datetime.datetime') as dt_mock: dt_mock.utcnow.return_value = mock_now dt_mock.strptime.return_value = expir_mock ret = session.is_valid self.assertEquals(False, ret)
def test_available_roles(self): roles = [ { "role": "::::1:role/role1", "principle": "" }, { "role": "::::1:role/role3", "principle": "" }, { "role": "::::1:role/role2", "principle": "" }, ] session = aws.Session("BogusAssertion") session.assertion = mock.MagicMock() session.assertion.roles.return_value = roles result = session.available_roles() print(result) expected = [ { "account": "1", "role_name": "role1", "principle": "", "arn": "::::1:role/role1", "roleIdx": 0, }, { "account": "1", "role_name": "role2", "principle": "", "arn": "::::1:role/role2", "roleIdx": 1, }, { "account": "1", "role_name": "role3", "principle": "", "arn": "::::1:role/role3", "roleIdx": 2, }, ] self.assertEqual(expected, result)
def test_write(self, mock_add_profile): session = aws.Session("BogusAssertion") ret = session._write() self.assertEqual(None, ret) # Verify add_profile is called with the correct args creds = { "AccessKeyId": None, "SecretAccessKey": None, "SessionToken": None, "Expiration": None, } mock_add_profile.assert_has_calls([ mock.call(creds=creds, name="default", region="us-east-1"), ], )
def start_session(self): """Initialize AWS session object.""" self.log.info('Getting SAML Assertion from {org}'.format( org=self.config.org)) assertion = self.okta_client.get_assertion( appid=self.config.appid) try: session = aws.Session(assertion, profile=self.config.name, role=self.role) except xml.etree.ElementTree.ParseError: self.log.error('Could not find any Role in the SAML assertion') self.log.error(assertion.__dict__) raise aws.InvalidSaml() return session
def test_assume_role_multiple(self, mock_write): mock_write.return_value = None assertion = mock.Mock() roles = [{'role': '1', 'principle': ''}, {'role': '2', 'principle': ''}] assertion.roles.return_value = roles session = aws.Session('BogusAssertion') 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 with self.assertRaises(aws.MultipleRoles): session.assume_role()
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
def test_assume_role_duration_rejected(self, mock_write): mock_write.return_value = None assertion = mock.Mock() assertion.roles.return_value = [{"arn": "", "principle": ""}] session = aws.Session("BogusAssertion") session.duration = 1000000 session.roles = [{"arn": "", "principle": ""}] session.assertion = assertion sts = { "Credentials": { "AccessKeyId": "AKI", "SecretAccessKey": "squirrel", "SessionToken": "token", "Expiration": "never", }, } session.sts = mock.Mock() err_mock = mock.MagicMock() err = botocore.exceptions.ClientError(err_mock, err_mock) session.sts.assume_role_with_saml.side_effect = [err, sts] session.assume_role() self.assertEqual("AKI", session.creds["AccessKeyId"]) self.assertEqual("squirrel", session.creds["SecretAccessKey"]) self.assertEqual("token", session.creds["SessionToken"]) self.assertEqual("never", session.creds["Expiration"]) session.sts.assert_has_calls([ mock.call.assume_role_with_saml( RoleArn="", PrincipalArn="", SAMLAssertion=mock.ANY, DurationSeconds=1000000, ), mock.call.assume_role_with_saml( RoleArn="", PrincipalArn="", SAMLAssertion=mock.ANY, DurationSeconds=3600, ), ], )