| 240 | |
| 241 | |
| 242 | class STSClientFactory: |
| 243 | def __init__(self, session): |
| 244 | self._session = session |
| 245 | |
| 246 | def get_sts_client(self, region_name=None, role_arn=None): |
| 247 | client_kwargs = {'region_name': region_name} |
| 248 | if role_arn is not None: |
| 249 | creds = self._get_role_credentials(region_name, role_arn) |
| 250 | client_kwargs['aws_access_key_id'] = creds['AccessKeyId'] |
| 251 | client_kwargs['aws_secret_access_key'] = creds['SecretAccessKey'] |
| 252 | client_kwargs['aws_session_token'] = creds['SessionToken'] |
| 253 | sts = self._session.create_client('sts', **client_kwargs) |
| 254 | self._register_k8s_aws_id_handlers(sts) |
| 255 | return sts |
| 256 | |
| 257 | def _get_role_credentials(self, region_name, role_arn): |
| 258 | sts = self._session.create_client('sts', region_name) |
| 259 | return sts.assume_role( |
| 260 | RoleArn=role_arn, RoleSessionName='EKSGetTokenAuth' |
| 261 | )['Credentials'] |
| 262 | |
| 263 | def _register_k8s_aws_id_handlers(self, sts_client): |
| 264 | sts_client.meta.events.register( |
| 265 | 'provide-client-params.sts.GetCallerIdentity', |
| 266 | self._retrieve_k8s_aws_id, |
| 267 | ) |
| 268 | sts_client.meta.events.register( |
| 269 | 'before-sign.sts.GetCallerIdentity', |
| 270 | self._inject_k8s_aws_id_header, |
| 271 | ) |
| 272 | |
| 273 | def _retrieve_k8s_aws_id(self, params, context, **kwargs): |
| 274 | if K8S_AWS_ID_HEADER in params: |
| 275 | context[K8S_AWS_ID_HEADER] = params.pop(K8S_AWS_ID_HEADER) |
| 276 | |
| 277 | def _inject_k8s_aws_id_header(self, request, **kwargs): |
| 278 | if K8S_AWS_ID_HEADER in request.context: |
| 279 | request.headers[K8S_AWS_ID_HEADER] = request.context[ |
| 280 | K8S_AWS_ID_HEADER |
| 281 | ] |