Partial compatibility with Google security update 2021

pydrive2/files.py

@@ -453,16 +453,14 @@ def FetchMetadata(self, fields=None, fetch_all=False):
 
         if file_id:
             try:
-                metadata = (
-                    self.auth.service.files()
-                    .get(
-                        fileId=file_id,
-                        fields=fields,
-                        # Teamdrive support
-                        supportsAllDrives=True,
-                    )
-                    .execute(http=self.http)
+                request = self.auth.service.files().get(
+                    fileId=file_id,
+                    fields=fields,
+                    # Teamdrive support
+                    supportsAllDrives=True,
                 )
+                request = self._AddResourceKeyHeaders(request)
+                metadata = request.execute(http=self.http)
             except errors.HttpError as error:
                 raise ApiRequestError(error)
             else:
@@ -687,6 +685,23 @@ def _WrapRequest(self, request):
         """
         if self.http:
             request.http = self.http
+        request = self._AddResourceKeyHeaders(request)
+        return request
+
+    def _AddResourceKeyHeaders(self, request):
+        """Add resourceKey headers to request if file is secured with resourceKey and
+        its available (from a list for example).
+
+        :param request: request to add headers to.
+        :type request: googleapiclient.http.HttpRequest
+        """
+        file_id = self.metadata.get("id") or self.get("id")
+        if file_id:
+            resourceKey = self.get("resourceKey")
+            if resourceKey:
+                request.headers[
+                    "X-Goog-Drive-Resource-Keys"
+                ] = f"{file_id}/{resourceKey}"
         return request
 
     @LoadAuth

pydrive2/test/test_file.py

@@ -1,6 +1,9 @@
 import filecmp
+import json
 import os
 import unittest
+from unittest.mock import MagicMock
+import httplib2
 import pytest
 import sys
 from io import BytesIO
@@ -25,6 +28,59 @@
 )
 
 
+def auth_with_resource_key_mock() -> GoogleAuth:
+    """
+    Create GoogleAuth with mocked inner httplib2.Http simulating need
+    for resourceKey header.
+    """
+    http_mock = MagicMock()
+
+    def resource_key_request(
+        uri,
+        method="GET",
+        body=None,
+        headers=None,
+        redirections=1,
+        connection_type=None,
+    ):
+        """httplib2.Http.request mock."""
+        # If "media" is not requested, it means we expect metadata
+        fetch_meta_data = "alt=media" not in uri
+        if (
+            headers
+            and "X-Goog-Drive-Resource-Keys" in headers
+            and headers["X-Goog-Drive-Resource-Keys"]
+            == "0BxphPoRgwhnodHNjS3JESnFNS1E/0-vjzOveuin3fnf4LUlfsD3A"
+        ):
+            if fetch_meta_data:
+                # Fake meta data query response
+                content = json.dumps({"title": "N48E012.zip"}).encode()
+            else:
+                # Fake file content query response
+                content = b"some content"
+            return (
+                httplib2.Response(
+                    {"status": "200", "content-length": str(len(content))}
+                ),
+                content,
+            )
+        # Simulate 404 response for file not found; body must be valid error JSON
+        return (
+            httplib2.Response({"status": "404"}),
+            json.dumps({"error": {"code": 404}}).encode(),
+        )
+
+    http_mock.request.side_effect = resource_key_request
+    ga = GoogleAuth(
+        settings_file_path(
+            "default.yaml", os.path.join(os.path.dirname(__file__), "")
+        )
+    )
+    ga.thread_local.http = http_mock
+    ga.ServiceAuth()
+    return ga
+
+
 class GoogleDriveFileTest(unittest.TestCase):
     """Tests basic file operations of files.GoogleDriveFile.
     Upload and download of contents and metadata, and thread-safety checks.
@@ -356,6 +412,128 @@ def test_Files_Get_Content_Buffer(self):
 
         self.DeleteUploadedFiles(drive, [file1["id"]])
 
+    def test_Files_Get_Content_Buffer_resourceKey_missing(self):
+        """404 expected for file secured with resourceKey when not provided."""
+
+        ga = auth_with_resource_key_mock()
+
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+            }
+        )
+        with self.assertRaisesRegex(
+            ApiRequestError, "HttpError 404 when requesting"
+        ):
+            pydrive_retry(file1.GetContentIOBuffer)
+
+    def test_Files_Get_Content_Buffer_resourceKey(self):
+        """End to end scenario with real file."""
+        ga = auth_with_resource_key_mock()
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+                "resourceKey": "0-vjzOveuin3fnf4LUlfsD3A",
+            }
+        )
+
+        buffer1 = pydrive_retry(file1.GetContentIOBuffer)
+
+        self.assertEqual(len(buffer1), 12)
+
+    @pytest.mark.manual
+    def test_Files_Get_Content_Buffer_resourceKey_missing_real(self):
+        """
+        404 expected for file secured with resourceKey when not provided.
+        End to end scenario with real public file.
+        """
+        drive = GoogleDrive(self.ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+            }
+        )
+        with self.assertRaisesRegex(
+            ApiRequestError, "HttpError 404 when requesting"
+        ):
+            pydrive_retry(file1.GetContentIOBuffer)
+
+    @pytest.mark.manual
+    def test_Files_Get_Content_Buffer_resourceKey_real(self):
+        """End to end scenario with real public file."""
+        drive = GoogleDrive(self.ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+                "resourceKey": "0-vjzOveuin3fnf4LUlfsD3A",
+            }
+        )
+
+        buffer1 = pydrive_retry(file1.GetContentIOBuffer)
+
+        self.assertEqual(len(buffer1), 6128902)
+
+    def test_Files_Get_Content_File_resourceKey_missing(self):
+        """404 expected for file secured with resourceKey when not provided."""
+        ga = auth_with_resource_key_mock()
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+            }
+        )
+        fileOut = self.getTempFile()
+        with self.assertRaisesRegex(
+            ApiRequestError, "HttpError 404 when requesting"
+        ):
+            pydrive_retry(file1.GetContentFile, fileOut)
+
+    def test_Files_Get_Content_File_resourceKey(self):
+        ga = auth_with_resource_key_mock()
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+                "resourceKey": "0-vjzOveuin3fnf4LUlfsD3A",
+            }
+        )
+
+        fileOut = self.getTempFile()
+        pydrive_retry(file1.GetContentFile, fileOut)
+
+        with open(fileOut, "rb") as f:
+            self.assertEqual(len(f.read()), 12)
+
+    def test_Files_Fetch_Metadata_resourceKey_missing(self):
+        """404 expected for file secured with resourceKey when not provided."""
+        ga = auth_with_resource_key_mock()
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+            }
+        )
+        with self.assertRaisesRegex(
+            ApiRequestError, "HttpError 404 when requesting"
+        ):
+            pydrive_retry(file1.FetchMetadata)
+
+    def test_Files_Fetch_Metadata_resourceKey(self):
+        ga = auth_with_resource_key_mock()
+        drive = GoogleDrive(ga)
+        file1 = drive.CreateFile(
+            {
+                "id": "0BxphPoRgwhnodHNjS3JESnFNS1E",
+                "resourceKey": "0-vjzOveuin3fnf4LUlfsD3A",
+            }
+        )
+
+        pydrive_retry(file1.FetchMetadata)
+
+        self.assertEqual(file1.metadata["title"], "N48E012.zip")
+
     def test_Upload_Download_Empty_File(self):
         filename = os.path.join(self.tmpdir, str(time()))
         create_file(filename, "")