Merge pull request #35 from nathan-v/Feature_region_arg

Add region parameter, screen only output option, command wrapper, bug fixes, v0.7.5

Author: nathan-v

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
 

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,14 +299,20 @@ 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)
 
     @staticmethod
     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

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

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! 👍')

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'

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):

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': ''}]