feat(storage): add support for signing URLs using token (#9889)

HemangChothani · web-flow · commit e7105475a0d9 · 2020-01-30T15:32:16.000-08:00
* feat(storage): add support for signing URLs using token

* feat(bigquery): add system tests for the feature

* feat(storage): add iam dependency in nox file and cosmetic changes
diff --git a/storage/google/cloud/storage/_signing.py b/storage/google/cloud/storage/_signing.py
@@ -19,10 +19,14 @@
 import datetime
 import hashlib
 import re
+import json
 
 import six
 
 import google.auth.credentials
+
+from google.auth import exceptions
+from google.auth.transport import requests
 from google.cloud import _helpers
 
 
@@ -265,6 +269,8 @@ def generate_signed_url_v2(
     generation=None,
     headers=None,
     query_parameters=None,
+    service_account_email=None,
+    access_token=None,
 ):
     """Generate a V2 signed URL to provide query-string auth'n to a resource.
 
@@ -340,6 +346,12 @@ def generate_signed_url_v2(
         Requests using the signed URL *must* pass the specified header
         (name and value) with each request for the URL.
 
+    :type service_account_email: str
+    :param service_account_email: (Optional) E-mail address of the service account.
+
+    :type access_token: str
+    :param access_token: (Optional) Access token for a service account.
+
     :type query_parameters: dict
     :param query_parameters:
         (Optional) Additional query paramtersto be included as part of the
@@ -370,9 +382,17 @@ def generate_signed_url_v2(
     string_to_sign = "\n".join(elements_to_sign)
 
     # Set the right query parameters.
-    signed_query_params = get_signed_query_params_v2(
-        credentials, expiration_stamp, string_to_sign
-    )
+    if access_token and service_account_email:
+        signature = _sign_message(string_to_sign, access_token, service_account_email)
+        signed_query_params = {
+            "GoogleAccessId": service_account_email,
+            "Expires": str(expiration),
+            "Signature": signature,
+        }
+    else:
+        signed_query_params = get_signed_query_params_v2(
+            credentials, expiration_stamp, string_to_sign
+        )
 
     if response_type is not None:
         signed_query_params["response-content-type"] = response_type
@@ -409,6 +429,8 @@ def generate_signed_url_v4(
     generation=None,
     headers=None,
     query_parameters=None,
+    service_account_email=None,
+    access_token=None,
     _request_timestamp=None,  # for testing only
 ):
     """Generate a V4 signed URL to provide query-string auth'n to a resource.
@@ -492,6 +514,12 @@ def generate_signed_url_v4(
         signed URLs.  See:
         https://cloud.google.com/storage/docs/xml-api/reference-headers#query
 
+    :type service_account_email: str
+    :param service_account_email: (Optional) E-mail address of the service account.
+
+    :type access_token: str
+    :param access_token: (Optional) Access token for a service account.
+
     :raises: :exc:`TypeError` when expiration is not a valid type.
     :raises: :exc:`AttributeError` if credentials is not an instance
             of :class:`google.auth.credentials.Signing`.
@@ -583,9 +611,58 @@ def generate_signed_url_v4(
     ]
     string_to_sign = "\n".join(string_elements)
 
-    signature_bytes = credentials.sign_bytes(string_to_sign.encode("ascii"))
-    signature = binascii.hexlify(signature_bytes).decode("ascii")
+    if access_token and service_account_email:
+        signature = _sign_message(string_to_sign, access_token, service_account_email)
+        signature_bytes = base64.b64decode(signature)
+        signature = binascii.hexlify(signature_bytes).decode("ascii")
+    else:
+        signature_bytes = credentials.sign_bytes(string_to_sign.encode("ascii"))
+        signature = binascii.hexlify(signature_bytes).decode("ascii")
 
     return "{}{}?{}&X-Goog-Signature={}".format(
         api_access_endpoint, resource, canonical_query_string, signature
     )
+
+
+def _sign_message(message, access_token, service_account_email):
+
+    """Signs a message.
+
+    :type message: str
+    :param message: The message to be signed.
+
+    :type access_token: str
+    :param access_token: Access token for a service account.
+
+
+    :type service_account_email: str
+    :param service_account_email: E-mail address of the service account.
+
+    :raises: :exc:`TransportError` if an `access_token` is unauthorized.
+
+    :rtype: str
+    :returns: The signature of the message.
+
+    """
+    message = _helpers._to_bytes(message)
+
+    method = "POST"
+    url = "https://iam.googleapis.com/v1/projects/-/serviceAccounts/{}:signBlob?alt=json".format(
+        service_account_email
+    )
+    headers = {
+        "Authorization": "Bearer " + access_token,
+        "Content-type": "application/json",
+    }
+    body = json.dumps({"bytesToSign": base64.b64encode(message).decode("utf-8")})
+
+    request = requests.Request()
+    response = request(url=url, method=method, body=body, headers=headers)
+
+    if response.status != six.moves.http_client.OK:
+        raise exceptions.TransportError(
+            "Error calling the IAM signBytes API: {}".format(response.data)
+        )
+
+    data = json.loads(response.data.decode("utf-8"))
+    return data["signature"]
diff --git a/storage/google/cloud/storage/blob.py b/storage/google/cloud/storage/blob.py
@@ -358,6 +358,8 @@ def generate_signed_url(
         client=None,
         credentials=None,
         version=None,
+        service_account_email=None,
+        access_token=None,
     ):
         """Generates a signed URL for this blob.
 
@@ -445,6 +447,12 @@ def generate_signed_url(
         :param version: (Optional) The version of signed credential to create.
                         Must be one of 'v2' | 'v4'.
 
+        :type service_account_email: str
+        :param service_account_email: (Optional) E-mail address of the service account.
+
+        :type access_token: str
+        :param access_token: (Optional) Access token for a service account.
+
         :raises: :exc:`ValueError` when version is invalid.
         :raises: :exc:`TypeError` when expiration is not a valid type.
         :raises: :exc:`AttributeError` if credentials is not an instance
@@ -497,6 +505,8 @@ def generate_signed_url(
             generation=generation,
             headers=headers,
             query_parameters=query_parameters,
+            service_account_email=service_account_email,
+            access_token=access_token,
         )
 
     def exists(self, client=None):
diff --git a/storage/noxfile.py b/storage/noxfile.py
@@ -112,7 +112,7 @@ def system(session):
     session.install("mock", "pytest")
     for local_dep in LOCAL_DEPS:
         session.install("-e", local_dep)
-    systest_deps = ["../test_utils/", "../pubsub", "../kms"]
+    systest_deps = ["../test_utils/", "../pubsub", "../kms", "../iam"]
     for systest_dep in systest_deps:
         session.install("-e", systest_dep)
     session.install("-e", ".")
diff --git a/storage/tests/system.py b/storage/tests/system.py
@@ -27,13 +27,13 @@
 import six
 
 from google.cloud import exceptions
+from google.cloud import iam_credentials_v1
 from google.cloud import storage
 from google.cloud.storage._helpers import _base64_md5hash
 from google.cloud.storage.bucket import LifecycleRuleDelete
 from google.cloud.storage.bucket import LifecycleRuleSetStorageClass
 from google.cloud import kms
 import google.oauth2
-
 from test_utils.retry import RetryErrors
 from test_utils.system import unique_resource_id
 from test_utils.vpcsc_config import vpcsc_config
@@ -109,7 +109,6 @@ def tearDown(self):
 
     def test_get_service_account_email(self):
         domain = "gs-project-accounts.iam.gserviceaccount.com"
-
         email = Config.CLIENT.get_service_account_email()
 
         new_style = re.compile(r"service-(?P<projnum>[^@]+)@" + domain)
@@ -962,6 +961,8 @@ def _create_signed_read_url_helper(
         payload=None,
         expiration=None,
         encryption_key=None,
+        service_account_email=None,
+        access_token=None,
     ):
         expiration = self._morph_expiration(version, expiration)
 
@@ -972,7 +973,12 @@ def _create_signed_read_url_helper(
             blob = self.blob
 
         signed_url = blob.generate_signed_url(
-            expiration=expiration, method=method, client=Config.CLIENT, version=version
+            expiration=expiration,
+            method=method,
+            client=Config.CLIENT,
+            version=version,
+            service_account_email=None,
+            access_token=None,
         )
 
         headers = {}
@@ -1045,6 +1051,29 @@ def test_create_signed_read_url_v4_w_csek(self):
             version="v4",
         )
 
+    def test_create_signed_read_url_v2_w_access_token(self):
+        client = iam_credentials_v1.IAMCredentialsClient()
+        service_account_email = Config.CLIENT._credentials.service_account_email
+        name = client.service_account_path("-", service_account_email)
+        scope = ["https://www.googleapis.com/auth/devstorage.read_write"]
+        response = client.generate_access_token(name, scope)
+        self._create_signed_read_url_helper(
+            service_account_email=service_account_email,
+            access_token=response.access_token,
+        )
+
+    def test_create_signed_read_url_v4_w_access_token(self):
+        client = iam_credentials_v1.IAMCredentialsClient()
+        service_account_email = Config.CLIENT._credentials.service_account_email
+        name = client.service_account_path("-", service_account_email)
+        scope = ["https://www.googleapis.com/auth/devstorage.read_write"]
+        response = client.generate_access_token(name, scope)
+        self._create_signed_read_url_helper(
+            version="v4",
+            service_account_email=service_account_email,
+            access_token=response.access_token,
+        )
+
     def _create_signed_delete_url_helper(self, version="v2", expiration=None):
         expiration = self._morph_expiration(version, expiration)
 
diff --git a/storage/tests/unit/test__signing.py b/storage/tests/unit/test__signing.py
@@ -390,6 +390,8 @@ def _generate_helper(
             generation=generation,
             headers=headers,
             query_parameters=query_parameters,
+            service_account_email=None,
+            access_token=None,
         )
 
         # Check the mock was called.
@@ -504,6 +506,22 @@ def test_with_google_credentials(self):
         with self.assertRaises(AttributeError):
             self._call_fut(credentials, resource=resource, expiration=expiration)
 
+    def test_with_access_token(self):
+        resource = "/name/path"
+        credentials = _make_credentials()
+        expiration = int(time.time() + 5)
+        email = mock.sentinel.service_account_email
+        with mock.patch(
+            "google.cloud.storage._signing._sign_message", return_value=b"DEADBEEF"
+        ):
+            self._call_fut(
+                credentials,
+                resource=resource,
+                expiration=expiration,
+                service_account_email=email,
+                access_token="token",
+            )
+
 
 class Test_generate_signed_url_v4(unittest.TestCase):
     DEFAULT_EXPIRATION = 1000
@@ -638,6 +656,51 @@ def test_w_custom_query_parameters_w_string_value(self):
     def test_w_custom_query_parameters_w_none_value(self):
         self._generate_helper(query_parameters={"qux": None})
 
+    def test_with_access_token(self):
+        resource = "/name/path"
+        signer_email = "service@example.com"
+        credentials = _make_credentials(signer_email=signer_email)
+        with mock.patch(
+            "google.cloud.storage._signing._sign_message", return_value=b"DEADBEEF"
+        ):
+            self._call_fut(
+                credentials,
+                resource=resource,
+                expiration=datetime.timedelta(days=5),
+                service_account_email=signer_email,
+                access_token="token",
+            )
+
+
+class Test_sign_message(unittest.TestCase):
+    @staticmethod
+    def _call_fut(*args, **kwargs):
+        from google.cloud.storage._signing import _sign_message
+
+        return _sign_message(*args, **kwargs)
+
+    def test_sign_bytes(self):
+        signature = "DEADBEEF"
+        data = {"signature": signature}
+        request = make_request(200, data)
+        with mock.patch("google.auth.transport.requests.Request", return_value=request):
+            returned_signature = self._call_fut(
+                "123", service_account_email="service@example.com", access_token="token"
+            )
+            assert returned_signature == signature
+
+    def test_sign_bytes_failure(self):
+        from google.auth import exceptions
+
+        request = make_request(401)
+        with mock.patch("google.auth.transport.requests.Request", return_value=request):
+            with pytest.raises(exceptions.TransportError):
+                self._call_fut(
+                    "123",
+                    service_account_email="service@example.com",
+                    access_token="token",
+                )
+
 
 _DUMMY_SERVICE_ACCOUNT = None
 
@@ -697,3 +760,16 @@ def _make_credentials(signer_email=None):
         return credentials
     else:
         return mock.Mock(spec=google.auth.credentials.Credentials)
+
+
+def make_request(status, data=None):
+    from google.auth import transport
+
+    response = mock.create_autospec(transport.Response, instance=True)
+    response.status = status
+    if data is not None:
+        response.data = json.dumps(data).encode("utf-8")
+
+    request = mock.create_autospec(transport.Request)
+    request.return_value = response
+    return request
diff --git a/storage/tests/unit/test_blob.py b/storage/tests/unit/test_blob.py
@@ -391,6 +391,8 @@ def _generate_signed_url_helper(
         credentials=None,
         expiration=None,
         encryption_key=None,
+        access_token=None,
+        service_account_email=None,
     ):
         from six.moves.urllib import parse
         from google.cloud._helpers import UTC
@@ -432,6 +434,8 @@ def _generate_signed_url_helper(
                 headers=headers,
                 query_parameters=query_parameters,
                 version=version,
+                access_token=access_token,
+                service_account_email=service_account_email,
             )
 
         self.assertEqual(signed_uri, signer.return_value)
@@ -464,6 +468,8 @@ def _generate_signed_url_helper(
             "generation": generation,
             "headers": expected_headers,
             "query_parameters": query_parameters,
+            "access_token": access_token,
+            "service_account_email": service_account_email,
         }
         signer.assert_called_once_with(expected_creds, **expected_kwargs)
 
