Added support for new style sso-session config

tests/test_cli.py

@@ -28,6 +28,19 @@ class CLIUnitTests(TestCase):
     def setUp(self) -> None:
         self.config = tempfile.NamedTemporaryFile(delete=False)
         conf_ini = b"""
+        [sso-session petshop]
+        sso_start_url = https://petshop.awsapps.com/start
+        sso_region = ap-southeast-2
+        sso_registration_scopes = sso:account:access
+
+        [profile dev2]
+        sso_session = petshop
+        sso_account_id = 123456789123
+        sso_role_name = AdministratorAccess
+        region = ap-southeast-2
+        output = json
+        cli_pager =
+
         [default]
         sso_start_url = https://petshop.awsapps.com/start
         sso_region = ap-southeast-2
@@ -80,7 +93,7 @@ def setUp(self) -> None:
         sso_role_name = AdministratorAccess
         region = ap-southeast-2
         output = json
-        ca_bundle = dev/apps/ca-certs/cabundle-2019mar05.pem        
+        ca_bundle = dev/apps/ca-certs/cabundle-2019mar05.pem
         """
         self.config.write(conf_ini)
         self.config.seek(0)
@@ -210,6 +223,45 @@ def test_main(self):
             self.assertEqual(new_tok, 'VeryLongBase664String==')
             verify(cli.utils, times=2).invoke(...)
 
+    def test_sso_session_config(self):
+        """
+        python -m unittest tests.test_cli.CLIUnitTests.test_sso_session_config
+        """
+        with ArgvContext(program, '-p', 'dev2', '--debug'):
+            cli.main()
+            cred = cli.utils.read_config(self.credentials.name)
+            new_tok = cred['dev2']['aws_session_token']
+            self.assertNotEqual(new_tok, 'tok')
+            self.assertEqual(new_tok, 'VeryLongBase664String==')
+            verify(cli.utils, times=2).invoke(...)
+
+    def test_sso_session_config_no_section(self):
+        """
+        python -m unittest tests.test_cli.CLIUnitTests.test_sso_session_config_no_section
+        """
+        with ArgvContext(program, '-p', 'dev2', '-t'), self.assertRaises(SystemExit) as x:
+            # clean up as going to mutate this
+            self.config.close()
+            os.unlink(self.config.name)
+            # now start new test case
+            self.config = tempfile.NamedTemporaryFile(delete=False)
+            conf_ini = b"""
+            [profile dev2]
+            sso_session = petshop
+            sso_account_id = 123456789123
+            sso_role_name = AdministratorAccess
+            region = ap-southeast-2
+            output = json
+            cli_pager =
+            """
+            self.config.write(conf_ini)
+            self.config.seek(0)
+            self.config.read()
+            cli.core.aws_config_file = self.config.name
+            cli.main()
+        self.assertEqual(x.exception.code, 1)
+        verify(cli.utils, times=1).invoke(...)
+
     def test_profile_prefix(self):
         """
         python -m unittest tests.test_cli.CLIUnitTests.test_profile_prefix

yawsso/cmd.py

@@ -11,9 +11,18 @@ class Command(object):
     def __init__(self, args):
         self.args = args
         self.config = utils.read_config(core.aws_config_file)
+        self._drop_sso_session_sections_from_config()
         self.profiles_new_name = dict()
         self.export_vars = self._build_export_vars()
 
+    def _drop_sso_session_sections_from_config(self):
+        sso_sessions = []
+        for s in self.config.sections():
+            if s.startswith("sso-session"):
+                sso_sessions.append(s)
+        for entry in sso_sessions:
+            self.config.remove_section(entry)
+
     def _build_export_vars(self):
         """Make export_vars avail either side of subcommand"""
         x_vars = self.args.export_vars if hasattr(self.args, 'export_vars') and self.args.export_vars else False

yawsso/core.py

@@ -209,9 +209,25 @@ def load_profile_from_config(profile_name, config):
 
 
 def is_sso_profile(profile):
+    return is_sso_legacy_profile(profile) or is_sso_session_profile(profile)
+
+
+def is_sso_legacy_profile(profile):
     return {"sso_start_url", "sso_account_id", "sso_role_name", "sso_region"} <= profile.keys()
 
 
+def is_sso_session_profile(profile):
+    is_sso_session = {"sso_session", "sso_account_id", "sso_role_name"} <= profile.keys()
+    if is_sso_session:
+        try:
+            config = u.read_config(aws_config_file)
+            sso_session_config = dict(config.items(f"sso-session {profile['sso_session']}"))
+            profile.update(sso_session_config)  # merge profile with sso-session section
+        except NoSectionError as e:
+            u.halt(e)
+    return is_sso_session
+
+
 def is_source_profile(profile):
     return {"source_profile", "role_arn", "region"} <= profile.keys()