fix: refactor tech debt in aws and identity pool credentials (#1501)

* fix: implement fixes suggested in suppliers PR * Add back elif * update const name to include unit
googleapis · Mar 18, 2024 · ce435b0 · ce435b0
1 parent edbd22e
commit ce435b0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 16 deletions.
diff --git a/google/auth/aws.py b/google/auth/aws.py
@@ -69,6 +69,8 @@
 _DEFAULT_AWS_REGIONAL_CREDENTIAL_VERIFICATION_URL = (
     "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15"
 )
+# IMDSV2 session token lifetime. This is set to a low value because the session token is used immediately.
+_IMDSV2_SESSION_TOKEN_TTL_SECONDS = "300"
 
 
 class RequestSigner(object):
@@ -476,9 +478,9 @@ def get_aws_region(self, context, request):
             else response.data
         )
 
-        if response.status != 200:
+        if response.status != http_client.OK:
             raise exceptions.RefreshError(
-                "Unable to retrieve AWS region", response_body
+                "Unable to retrieve AWS region: {}".format(response_body)
             )
 
         # This endpoint will return the region in format: us-east-2b.
@@ -487,16 +489,19 @@ def get_aws_region(self, context, request):
 
     def _get_imdsv2_session_token(self, request):
         if request is not None and self._imdsv2_session_token_url is not None:
-            headers = {"X-aws-ec2-metadata-token-ttl-seconds": "300"}
+            headers = {
+                "X-aws-ec2-metadata-token-ttl-seconds": _IMDSV2_SESSION_TOKEN_TTL_SECONDS
+            }
 
             imdsv2_session_token_response = request(
                 url=self._imdsv2_session_token_url, method="PUT", headers=headers
             )
 
-            if imdsv2_session_token_response.status != 200:
+            if imdsv2_session_token_response.status != http_client.OK:
                 raise exceptions.RefreshError(
-                    "Unable to retrieve AWS Session Token",
-                    imdsv2_session_token_response.data,
+                    "Unable to retrieve AWS Session Token: {}".format(
+                        imdsv2_session_token_response.data
+                    )
                 )
 
             return imdsv2_session_token_response.data
@@ -545,7 +550,7 @@ def _get_metadata_security_credentials(
 
         if response.status != http_client.OK:
             raise exceptions.RefreshError(
-                "Unable to retrieve AWS security credentials", response_body
+                "Unable to retrieve AWS security credentials: {}".format(response_body)
             )
 
         credentials_response = json.loads(response_body)
@@ -593,7 +598,7 @@ def _get_metadata_role_name(self, request, imdsv2_session_token):
 
         if response.status != http_client.OK:
             raise exceptions.RefreshError(
-                "Unable to retrieve AWS role name", response_body
+                "Unable to retrieve AWS role name {}".format(response_body)
             )
 
         return response_body
@@ -690,7 +695,7 @@ def __init__(
                 "regional_cred_verification_url"
             )
 
-            # Get the environment ID. Currently, only one version supported (v1).
+            # Get the environment ID, i.e. "aws1". Currently, only one version supported (1).
             matches = re.match(r"^(aws)([\d]+)$", environment_id)
             if matches:
                 env_id, env_version = matches.groups()
@@ -701,7 +706,7 @@ def __init__(
                 raise exceptions.InvalidResource(
                     "No valid AWS 'credential_source' provided"
                 )
-            elif int(env_version or "") != 1:
+            elif env_version is None or int(env_version) != 1:
                 raise exceptions.InvalidValue(
                     "aws version '{}' is not supported in the current build.".format(
                         env_version
@@ -784,15 +789,12 @@ def retrieve_subject_token(self, request):
         request_headers["x-goog-cloud-target-resource"] = self._target_resource
 
         # Serialize AWS signed request.
-        # Keeping inner keys in sorted order makes testing easier for Python
-        # versions <=3.5 as the stringified JSON string would have a predictable
-        # key order.
         aws_signed_req = {}
         aws_signed_req["url"] = request_options.get("url")
         aws_signed_req["method"] = request_options.get("method")
         aws_signed_req["headers"] = []
         # Reformat header to GCP STS expected format.
-        for key in sorted(request_headers.keys()):
+        for key in request_headers.keys():
             aws_signed_req["headers"].append(
                 {"key": key, "value": request_headers[key]}
             )

diff --git a/google/auth/identity_pool.py b/google/auth/identity_pool.py
@@ -41,7 +41,6 @@
 except ImportError:  # pragma: NO COVER
     from collections import Mapping
 import abc
-import io
 import json
 import os
 from typing import NamedTuple
@@ -104,7 +103,7 @@ def get_subject_token(self, context, request):
         if not os.path.exists(self._path):
             raise exceptions.RefreshError("File '{}' was not found.".format(self._path))
 
-        with io.open(self._path, "r", encoding="utf-8") as file_obj:
+        with open(self._path, "r", encoding="utf-8") as file_obj:
             token_content = _TokenContent(file_obj.read(), self._path)
 
         return _parse_token_data(