[WIP] Add option for Role session name to EKS subcommands

awscli/customizations/eks/get_token.py

@@ -95,6 +95,14 @@ class GetTokenCommand(BasicCommand):
             ),
             'required': False,
         },
+        {
+            'name': 'session-name',
+            'help_text': (
+                "Use this parameter with --role-arn to specify a role session name. "
+                "When omitted, the role session name defaults to 'EKSGetTokenAuth'."
+            ),
+            'required': False,
+        },
         {
             'name': 'cluster-id',
             # When EKS in-region cluster supports cluster-id, we will need to update this help text
@@ -114,7 +122,7 @@ def get_expiration_time(self):
     def _run_main(self, parsed_args, parsed_globals):
         client_factory = STSClientFactory(self._session)
         sts_client = client_factory.get_sts_client(
-            region_name=parsed_globals.region, role_arn=parsed_args.role_arn
+            region_name=parsed_globals.region, role_arn=parsed_args.role_arn, role_session_name=parsed_args.session_name
         )
 
         validate_mutually_exclusive(parsed_args, ['cluster_name'], ['cluster_id'])
@@ -240,21 +248,23 @@ class STSClientFactory(object):
     def __init__(self, session):
         self._session = session
 
-    def get_sts_client(self, region_name=None, role_arn=None):
+    def get_sts_client(self, region_name=None, role_arn=None, role_session_name=None):
         client_kwargs = {'region_name': region_name}
         if role_arn is not None:
-            creds = self._get_role_credentials(region_name, role_arn)
+            creds = self._get_role_credentials(region_name, role_arn, role_session_name)
             client_kwargs['aws_access_key_id'] = creds['AccessKeyId']
             client_kwargs['aws_secret_access_key'] = creds['SecretAccessKey']
             client_kwargs['aws_session_token'] = creds['SessionToken']
         sts = self._session.create_client('sts', **client_kwargs)
         self._register_k8s_aws_id_handlers(sts)
         return sts
 
-    def _get_role_credentials(self, region_name, role_arn):
+    def _get_role_credentials(self, region_name, role_arn, role_session_name):
         sts = self._session.create_client('sts', region_name)
+        if role_session_name is None:
+            role_session_name = 'EKSGetTokenAuth'
         return sts.assume_role(
-            RoleArn=role_arn, RoleSessionName='EKSGetTokenAuth'
+            RoleArn=role_arn, RoleSessionName=role_session_name
         )['Credentials']
 
     def _register_k8s_aws_id_handlers(self, sts_client):

awscli/customizations/eks/update_kubeconfig.py

@@ -78,6 +78,12 @@ class UpdateKubeconfigCommand(BasicCommand):
                           "connect to the cluster the first time."),
             'required': False
         },
+        {
+            'name': 'session-name',
+            'help_text': ("The name of the role session to be passed down "
+                          "to further commands."),
+            'required': False
+        },
         {
             'name': 'dry-run',
             'action': 'store_true',
@@ -334,6 +340,12 @@ def get_user_entry(self, user_alias=None):
                 self._parsed_args.role_arn
             ])
 
+        if self._parsed_args.session_name is not None:
+            generated_user["user"]["exec"]["args"].extend([
+                "--session-name",
+                self._parsed_args.session_name
+            ])
+
         if self._session.profile:
             generated_user["user"]["exec"]["env"] = [OrderedDict([
                 ("name", "AWS_PROFILE"),

tests/functional/eks/test_get_token.py

@@ -25,6 +25,7 @@ def setUp(self):
         super(TestGetTokenCommand, self).setUp()
         self.cluster_name = 'MyCluster'
         self.role_arn = 'arn:aws:iam::012345678910:role/RoleArn'
+        self.session_name = 'CustomSessionName123'
         self.access_key = 'ABCDEFGHIJKLMNOPQRST'
         self.secret_key = 'TSRQPONMLKJUHGFEDCBA'
         self.session_token = 'TOKENTOKENTOKENTOKEN'
@@ -174,6 +175,28 @@ def test_url_with_arn(self):
         )
         self.assert_url_correct(response, has_session_token=True)
 
+    def test_url_with_arn_and_session_name(self):
+        cmd = 'eks get-token --cluster-name %s' % self.cluster_name
+        cmd += ' --role-arn %s' % self.role_arn
+        cmd += ' --session-name %s' % self.session_name
+        self.parsed_responses = [
+            {
+                "Credentials": {
+                    "AccessKeyId": self.access_key,
+                    "SecretAccessKey": self.secret_key,
+                    "SessionToken": self.session_token,
+                },
+            }
+        ]
+        response = self.run_get_token(cmd)
+        assume_role_call = self.operations_called[0]
+        self.assertEqual(assume_role_call[0].name, 'AssumeRole')
+        self.assertEqual(
+            assume_role_call[1],
+            {'RoleArn': self.role_arn, 'RoleSessionName': self.session_name},
+        )
+        self.assert_url_correct(response, has_session_token=True)
+
     def test_token_has_no_padding(self):
         cmd = 'eks get-token --cluster-name %s' % self.cluster_name
         num_rounds = 100