Skip to content

Commit

Permalink
Partial compatibility with Google security update 2021 (#350)
Browse files Browse the repository at this point in the history 
* Partial compatibility with Google security update 2021

Google Drive added a new "resourceKey" attribute required to access documents shared by links.
This resourceKey must be passed through HTTP header, aside with the document ID.
resourceKey can be retrieved from a previous list operation on containing folder.
Partial implementation; only for basic methods: GetContentFile, GetContentIOBuffer and FetchMetadata.

* lint

---------

Co-authored-by: Ivan Shcheklein <[email protected]>
  • Loading branch information
@agrenott @shcheklein
agrenott and shcheklein authored Jul 15, 2024
1 parent 66cf7bc commit a84dedd
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 9 deletions.
33 changes: 24 additions & 9 deletions pydrive2/files.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
178 changes: 178 additions & 0 deletions pydrive2/test/test_file.py
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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, "")
Expand Down

0 comments on commit a84dedd

Please sign in to comment.