Support WIF Impersonation on AWS workloads (#2517)

Co-authored-by: Peter Mansour <peter.mansour@snowflake.com>

Author: 2 people

src/snowflake/connector/connection.py

@@ -1377,14 +1377,18 @@ def __open_connection(self):
                     )
                 if (
                     self._workload_identity_impersonation_path
-                    and self._workload_identity_provider != AttestationProvider.GCP
+                    and self._workload_identity_provider
+                    not in (
+                        AttestationProvider.GCP,
+                        AttestationProvider.AWS,
+                    )
                 ):
                     Error.errorhandler_wrapper(
                         self,
                         None,
                         ProgrammingError,
                         {
-                            "msg": "workload_identity_impersonation_path is currently only supported for GCP.",
+                            "msg": "workload_identity_impersonation_path is currently only supported for GCP and AWS.",
                             "errno": ER_INVALID_WIF_SETTINGS,
                         },
                     )

src/snowflake/connector/wif_util.py

@@ -148,8 +148,29 @@ def get_aws_sts_hostname(region: str, partition: str) -> str:
         )
 
 
+def get_aws_session(impersonation_path: list[str] | None = None):
+    """Creates a boto3 session with the appropriate credentials.
+
+    If impersonation_path is provided, this uses the role at the end of the path. Otherwise, this uses the role attached to the current workload.
+    """
+    session = boto3.session.Session()
+
+    impersonation_path = impersonation_path or []
+    for arn in impersonation_path:
+        response = session.client("sts").assume_role(
+            RoleArn=arn, RoleSessionName="identity-federation-session"
+        )
+        creds = response["Credentials"]
+        session = boto3.session.Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+        )
+    return session
+
+
 def create_aws_attestation(
-    session_manager: SessionManager | None = None,
+    impersonation_path: list[str] | None = None,
 ) -> WorkloadIdentityAttestation:
     """Tries to create a workload identity attestation for AWS.
 
@@ -162,7 +183,8 @@ def create_aws_attestation(
         )
 
     # TODO: SNOW-2223669 Investigate if our adapters - containing settings of http traffic - should be passed here as boto urllib3session. Those requests go to local servers, so they do not need Proxy setup or Headers customization in theory. But we may want to have all the traffic going through one class (e.g. Adapter or mixin).
-    session = boto3.session.Session()
+    session = get_aws_session(impersonation_path)
+
     aws_creds = session.get_credentials()
     if not aws_creds:
         raise ProgrammingError(
@@ -396,7 +418,7 @@ def create_attestation(
     )
 
     if provider == AttestationProvider.AWS:
-        return create_aws_attestation(session_manager)
+        return create_aws_attestation(impersonation_path)
     elif provider == AttestationProvider.AZURE:
         return create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:

test/csp_helpers.py

@@ -40,9 +40,8 @@ def gen_dummy_id_token(
     )
 
 
-def gen_dummy_access_token(sub="test-subject") -> str:
+def gen_dummy_access_token(sub="test-subject", key="secret") -> str:
     """Generates a dummy access token using the given subject."""
-    key = "secret"
     logger.debug(f"Generating dummy access token for subject {sub}")
     return (sub + key).encode("utf-8").hex()
 
@@ -368,6 +367,11 @@ class FakeAwsEnvironment:
     def __init__(self):
         # Defaults used for generating a token. Can be overriden in individual tests.
         self.arn = "arn:aws:sts::123456789:assumed-role/My-Role/i-34afe100cad287fab"
+        # Path of roles that can be assumed. Empty if no impersonation is allowed.
+        # Can be overriden in individual tests.
+        self.assumption_path = []
+        self.assume_role_call_count = 0
+
         self.caller_identity = {"Arn": self.arn}
         self.region = "us-east-1"
         self.credentials = Credentials(access_key="ak", secret_key="sk")
@@ -376,6 +380,25 @@ def __init__(self):
         )
         self.metadata_token = "test-token"
 
+    def assume_role(self, **kwargs):
+        if (
+            self.assumption_path
+            and kwargs["RoleArn"] == self.assumption_path[self.assume_role_call_count]
+        ):
+            arn = self.assumption_path[self.assume_role_call_count]
+            self.assume_role_call_count += 1
+            return {
+                "Credentials": {
+                    "AccessKeyId": "access_key",
+                    "SecretAccessKey": "secret_key",
+                    "SessionToken": "session_token",
+                    "Expiration": int(time()) + 60 * 60,
+                },
+                "AssumedRoleUser": {"AssumedRoleId": hash(arn), "Arn": arn},
+                "ResponseMetadata": {},
+            }
+        return {}
+
     def get_region(self):
         return self.region
 
@@ -401,6 +424,7 @@ def fetcher_fetch_metadata_token(self):
     def boto3_client(self, *args, **kwargs):
         mock_client = mock.Mock()
         mock_client.get_caller_identity.return_value = self.caller_identity
+        mock_client.assume_role = self.assume_role
         return mock_client
 
     def __enter__(self):
@@ -443,6 +467,9 @@ def __enter__(self):
                 side_effect=self.boto3_client,
             )
         )
+        self.patchers.append(
+            mock.patch("boto3.session.Session.client", side_effect=self.boto3_client)
+        )
         for patcher in self.patchers:
             patcher.__enter__()
         return self

test/unit/test_auth_workload_identity.py

@@ -274,6 +274,22 @@ def test_get_aws_sts_hostname_invalid_inputs(region, partition):
     assert "Invalid AWS partition" in str(excinfo.value)
 
 
+def test_aws_impersonation_calls_correct_apis_for_each_role_in_impersonation_path(
+    fake_aws_environment: FakeAwsEnvironment,
+):
+    impersonation_path = [
+        "arn:aws:iam::123456789:role/role2",
+        "arn:aws:iam::123456789:role/role3",
+    ]
+    fake_aws_environment.assumption_path = impersonation_path
+    auth_class = AuthByWorkloadIdentity(
+        provider=AttestationProvider.AWS, impersonation_path=impersonation_path
+    )
+    auth_class.prepare(conn=None)
+
+    assert fake_aws_environment.assume_role_call_count == 2
+
+
 # -- GCP Tests --
 
 

test/unit/test_connection.py

@@ -684,16 +684,14 @@ def test_workload_identity_provider_is_required_for_wif_authenticator(
     "provider_param",
     [
         # Strongly-typed values.
-        AttestationProvider.AWS,
         AttestationProvider.AZURE,
         AttestationProvider.OIDC,
         # String values.
-        "AWS",
         "AZURE",
         "OIDC",
     ],
 )
-def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
+def test_workload_identity_impersonation_path_errors_for_unsupported_providers(
     monkeypatch, provider_param
 ):
     with monkeypatch.context() as m:
@@ -711,20 +709,22 @@ def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
                 ],
             )
         assert (
-            "workload_identity_impersonation_path is currently only supported for GCP."
+            "workload_identity_impersonation_path is currently only supported for GCP and AWS."
             in str(excinfo.value)
         )
 
 
 @pytest.mark.parametrize(
-    "provider_param",
+    "provider_param,impersonation_path",
     [
-        AttestationProvider.GCP,
-        "GCP",
+        (AttestationProvider.GCP, ["sa2@project.iam.gserviceaccount.com"]),
+        (AttestationProvider.AWS, ["arn:aws:iam::1234567890:role/role2"]),
+        ("GCP", ["sa2@project.iam.gserviceaccount.com"]),
+        ("AWS", ["arn:aws:iam::1234567890:role/role2"]),
     ],
 )
-def test_workload_identity_impersonation_path_supported_for_gcp_provider(
-    monkeypatch, provider_param
+def test_workload_identity_impersonation_path_populates_auth_class_for_supported_provider(
+    monkeypatch, provider_param, impersonation_path
 ):
     with monkeypatch.context() as m:
         m.setattr(
@@ -735,14 +735,9 @@ def test_workload_identity_impersonation_path_supported_for_gcp_provider(
             account="account",
             authenticator="WORKLOAD_IDENTITY",
             workload_identity_provider=provider_param,
-            workload_identity_impersonation_path=[
-                "sa2@project.iam.gserviceaccount.com"
-            ],
+            workload_identity_impersonation_path=impersonation_path,
         )
-        assert conn.auth_class.provider == AttestationProvider.GCP
-        assert conn.auth_class.impersonation_path == [
-            "sa2@project.iam.gserviceaccount.com"
-        ]
+        assert conn.auth_class.impersonation_path == impersonation_path
 
 
 @pytest.mark.parametrize(