Repository: incubator-airflow Updated Branches: refs/heads/master a5fa1bf88 -> 449a7fd1b
[AIRFLOW-2807] Support STS Assume Role External ID Currently the role assumption method works only if the granting account does not specify an External ID. The external ID is used to solved the confused deputy problem. When using the AWS hook to export data to multiple customers, it's good security practice to use the external ID. There is no backwards compatibility break, the ID will be `None` in existing cases. Moto doesn't provide any convenient way to verify the value was passed in the credential response in tests, so existing test cases are kept. Documentation: https://docs.aws.amazon.com/IAM/lat est/UserGuide/id_roles_create_for- user_externalid.html Closes #3647 from vvondra/support_sts_external_id Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/449a7fd1 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/449a7fd1 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/449a7fd1 Branch: refs/heads/master Commit: 449a7fd1b72639f1eb2bbeb033e1642e8eaac96c Parents: a5fa1bf Author: Vojtech Vondra <[email protected]> Authored: Sat Jul 28 17:32:09 2018 +0200 Committer: Bolke de Bruin <[email protected]> Committed: Sat Jul 28 17:32:09 2018 +0200 ---------------------------------------------------------------------- airflow/contrib/hooks/aws_hook.py | 15 ++++++++++++--- tests/contrib/hooks/test_aws_hook.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/449a7fd1/airflow/contrib/hooks/aws_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/aws_hook.py b/airflow/contrib/hooks/aws_hook.py index c712d2d..69a1b0b 100644 --- a/airflow/contrib/hooks/aws_hook.py +++ b/airflow/contrib/hooks/aws_hook.py @@ -116,6 +116,7 @@ class AwsHook(BaseHook): region_name = connection_object.extra_dejson.get('region_name') role_arn = connection_object.extra_dejson.get('role_arn') + external_id = connection_object.extra_dejson.get('external_id') aws_account_id = connection_object.extra_dejson.get('aws_account_id') aws_iam_role = connection_object.extra_dejson.get('aws_iam_role') @@ -130,9 +131,17 @@ class AwsHook(BaseHook): region_name=region_name) sts_client = sts_session.client('sts') - sts_response = sts_client.assume_role( - RoleArn=role_arn, - RoleSessionName='Airflow_' + self.aws_conn_id) + + if external_id is None: + sts_response = sts_client.assume_role( + RoleArn=role_arn, + RoleSessionName='Airflow_' + self.aws_conn_id) + else: + sts_response = sts_client.assume_role( + RoleArn=role_arn, + RoleSessionName='Airflow_' + self.aws_conn_id, + ExternalId=external_id) + aws_access_key_id = sts_response['Credentials']['AccessKeyId'] aws_secret_access_key = sts_response['Credentials']['SecretAccessKey'] aws_session_token = sts_response['Credentials']['SessionToken'] http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/449a7fd1/tests/contrib/hooks/test_aws_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_aws_hook.py b/tests/contrib/hooks/test_aws_hook.py index 719d6d1..b27b37d 100644 --- a/tests/contrib/hooks/test_aws_hook.py +++ b/tests/contrib/hooks/test_aws_hook.py @@ -164,6 +164,25 @@ class TestAwsHook(unittest.TestCase): 'gRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15' 'fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE') + @unittest.skipIf(mock_sts is None, 'mock_sts package not present') + @mock.patch.object(AwsHook, 'get_connection') + @mock_sts + def test_get_credentials_from_role_arn_with_external_id(self, mock_get_connection): + mock_connection = Connection( + extra='{"role_arn":"arn:aws:iam::123456:role/role_arn",' + ' "external_id":"external_id"}') + mock_get_connection.return_value = mock_connection + hook = AwsHook() + credentials_from_hook = hook.get_credentials() + self.assertEqual(credentials_from_hook.access_key, 'AKIAIOSFODNN7EXAMPLE') + self.assertEqual(credentials_from_hook.secret_key, + 'aJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY') + self.assertEqual(credentials_from_hook.token, + 'BQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh' + '3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4I' + 'gRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15' + 'fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE') + if __name__ == '__main__': unittest.main()
