baolsen commented on a change in pull request #7619: [AIRFLOW-6975] Base
AWSHook AssumeRoleWithSAML
URL: https://github.com/apache/airflow/pull/7619#discussion_r388081436
##########
File path: airflow/providers/amazon/aws/hooks/base_aws.py
##########
@@ -206,6 +212,104 @@ def _get_credentials(self, region_name):
endpoint_url,
)
+ def _assume_role(
+ self,
+ sts_client,
+ extra_config,
+ role_arn,
+ assume_role_kwargs):
+ if "external_id" in extra_config: # Backwards compatibility
+ assume_role_kwargs["ExternalId"] = extra_config.get(
+ "external_id"
+ )
+ role_session_name = "Airflow_" + self.aws_conn_id
+ self.log.info(
+ "Doing sts_client.assume_role to role_arn=%s
(role_session_name=%s)",
+ role_arn,
+ role_session_name,
+ )
+ return sts_client.assume_role(
+ RoleArn=role_arn,
+ RoleSessionName=role_session_name,
+ **assume_role_kwargs
+ )
+
+ def _assume_role_with_saml(
+ self,
+ sts_client,
+ extra_config,
+ role_arn,
+ assume_role_kwargs):
+
+ saml_config = extra_config['assume_role_with_saml']
+ principal_arn = saml_config['principal_arn']
+
+ idp_url = saml_config["idp_url"]
+ self.log.info("idp_url= %s", idp_url)
+
+ idp_request_kwargs = saml_config["idp_request_kwargs"]
+
+ idp_auth_method = saml_config['idp_auth_method']
+ if idp_auth_method == 'http_spegno_auth':
+ # requests_gssapi will need paramiko > 2.6 since you'll need
+ # 'gssapi' not 'python-gssapi' from PyPi.
+ # https://github.com/paramiko/paramiko/pull/1311
+ import requests_gssapi
+ auth = requests_gssapi.HTTPSPNEGOAuth()
+ if 'mutual_authentication' in saml_config:
+ mutual_auth = saml_config['mutual_authentication']
+ if mutual_auth == 'REQUIRED':
+ auth = requests_gssapi.HTTPSPNEGOAuth(
+ requests_gssapi.REQUIRED)
+ elif mutual_auth == 'OPTIONAL':
+ auth = requests_gssapi.HTTPSPNEGOAuth(
+ requests_gssapi.OPTIONAL)
+ elif mutual_auth == 'DISABLED':
+ auth = requests_gssapi.HTTPSPNEGOAuth(
+ requests_gssapi.DISABLED)
+ else:
+ raise NotImplementedError(
+ 'mutual_authentication=%s' % mutual_auth)
+ # Query the IDP
+ import requests
+ idp_reponse = requests.get(
+ idp_url, auth=auth, **idp_request_kwargs)
+ idp_reponse.raise_for_status()
+
+ # Assist with debugging. Note: contains sensitive info!
+ xpath = saml_config['saml_response_xpath']
+ log_idp_response = 'log_idp_response' in saml_config and
saml_config[
+ 'log_idp_response']
+ if log_idp_response:
+ self.log.warning(
+ 'The IDP response contains sensitive information,'
+ ' but log_idp_response is ON (%s).', log_idp_response)
+ self.log.info('idp_reponse.content= %s', idp_reponse.content)
+ self.log.info('xpath= %s', xpath)
+
+ # Extract SAML Assertion from the returned HTML / XML
+ from lxml import etree
+ xml = etree.fromstring(idp_reponse.content)
+ saml_assertion = xml.xpath(xpath)
+ if isinstance(saml_assertion, list):
+ if len(saml_assertion) == 1:
+ saml_assertion = saml_assertion[0]
+ if not saml_assertion:
+ raise ValueError('Invalid SAML Assertion')
+ else:
+ raise NotImplementedError('idp_auth_method=%s' % idp_auth_method)
Review comment:
Also fixed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
