def authenticate(self):
cache = JSONFileCache()
saml_fetcher = SAMLFetcher(self, cache=cache)
credentials = saml_fetcher.fetch_credentials()
return credentials
def get_accounts_and_roles(self):
cache = JSONFileCache()
saml_fetcher = SAMLFetcher(self, cache=cache)
app_and_role = saml_fetcher.get_app_roles()
result_accounts = []
results = {
"application_url": app_and_role["Application"],
"accounts": result_accounts,
"user": app_and_role["User"],
"organization": app_and_role["Organization"],
}
accounts = app_and_role["Accounts"]
for name_raw in accounts:
account_parts = re.match(
r"(Account:) ([a-zA-Z0-9-_]+) \(([0-9]+)\)", name_raw)
account = account_parts[2]
account_id = account_parts[3]
roles = accounts[name_raw]
result_roles = []
result_account = {
"name": account,
"id": account_id,
"name_raw": name_raw,
"roles": result_roles
}
result_accounts.append(result_account)
for role in roles:
role_suffix = role.split(
os.environ.get("AWS_OKTA_ROLE_SUFFIX_DELIMITER", "-"))[-1]
result_roles.append({"name": role, "suffix": role_suffix})
return results
def test_fetcher(self, mock_okta, mock_print_tty, mock_client):
self.OPTIONS["--role"] = "arn:aws:iam::2:role/Role-One"
mock_okta().get_saml_response.return_value = SAML_RESPONSE
mock_cache = MagicMock()
authenticate = Authenticate(self.OPTIONS)
fetcher = SAMLFetcher(authenticate, cache=mock_cache)
fetcher.fetch_credentials()
def test_get_app_roles(self, mock_get_app_roles):
mock_get_app_roles.return_value = ("accounts", None, "app-url", "jdoe", 'test-org')
authenticate = Authenticate(self.OPTIONS)
fetcher = SAMLFetcher(authenticate, cache={})
actual = fetcher.get_app_roles()
self.assertEqual({
'Accounts': 'accounts',
'Application': 'app-url',
'Organization': 'test-org',
'User': '******'
}, actual)
def test_fetcher_should_prompt_all_accounts(
self,
mock_okta,
mock_prompt,
mock_prompt_print_tty,
mock_print_tty,
mock_client
):
def assume_role_side_effect(*args, **kwargs):
if kwargs['RoleArn'] == 'arn:aws:iam::1:role/Role-One':
return {
'Credentials': {
'AccessKeyId': 'test-key1',
'SecretAccessKey': 'test-secret1',
'SessionToken': 'test-token1',
'Expiration': datetime(2020, 4, 17, 12, 0, 0, 0)
}
}
raise RuntimeError('invalid RoleArn')
self.OPTIONS["--pass"] = '******'
mock_c = mock.Mock()
mock_c.assume_role_with_saml.side_effect = assume_role_side_effect
mock_okta().get_saml_response.return_value = SAML_RESPONSE
mock_client.return_value = mock_c
authenticate = Authenticate(self.OPTIONS)
fetcher = SAMLFetcher(authenticate, cache={})
creds = fetcher.fetch_credentials()
self.assertDictEqual({
'AccessKeyId': 'test-key1',
'Expiration': '2020-04-17T12:00:00',
'SecretAccessKey': 'test-secret1',
'SessionToken': 'test-token1'
}, creds)
self.assertEqual(7, mock_prompt_print_tty.call_count)
MagicMock.assert_has_calls(mock_prompt_print_tty, [
call('Select AWS Role:'),
call('Account: 1', indents=0),
call('[ 1 ] Role-One', indents=1),
call('[ 2 ] Role-Two', indents=1),
call('Account: 2', indents=0),
call('[ 3 ] Role-One', indents=1),
call('Selection: ', newline=False)
])
def run(self):
cache = JSONFileCache()
saml_fetcher = SAMLFetcher(self, cache=cache)
credentials = saml_fetcher.fetch_credentials()
if self.configuration["AWS_OKTA_ENVIRONMENT"]:
if os.name == 'nt':
print(
NT_EXPORT_STRING.format(credentials["AccessKeyId"],
credentials["SecretAccessKey"],
credentials["SessionToken"]))
else:
print(
UNIX_EXPORT_STRING.format(credentials["AccessKeyId"],
credentials["SecretAccessKey"],
credentials["SessionToken"]))
else:
credentials["Version"] = 1
print(json.dumps(credentials))