diff --git a/README.md b/README.md index 1534e09..4c9b6de 100644 --- a/README.md +++ b/README.md @@ -403,7 +403,8 @@ gauth = GoogleAuth() gauth.LocalWebserverAuth() drive = GoogleDrive(gauth) -file_list = drive.ListFile({"q": "'folder_id' in parents"}).GetList() +folder_id = "YOUR_FOLDER_ID" # a Drive id: [A-Za-z0-9_-]; quote any other value +file_list = drive.ListFile({"q": f"'{folder_id}' in parents"}).GetList() for file in file_list: content = file.GetContentString() diff --git a/pydrivedol/base.py b/pydrivedol/base.py index 97d060b..2b757b9 100644 --- a/pydrivedol/base.py +++ b/pydrivedol/base.py @@ -61,9 +61,9 @@ def _extract_file_id(url: str) -> Optional[str]: 'ABC123' """ patterns = [ - r"drive\.google\.com/file/d/([^/]+)", - r"drive\.google\.com/open\?id=([^&]+)", - r"drive\.google\.com/uc\?.*id=([^&]+)", + r"drive\.google\.com/file/d/([A-Za-z0-9_-]+)", + r"drive\.google\.com/open\?id=([A-Za-z0-9_-]+)", + r"drive\.google\.com/uc\?(?:[^#]*&)?id=([A-Za-z0-9_-]+)", ] for pattern in patterns: match = re.search(pattern, url) @@ -72,6 +72,20 @@ def _extract_file_id(url: str) -> Optional[str]: return None +def _q(value) -> str: + """Quote *value* as a string literal in the Drive search-query language. + + Backslashes and single quotes are escaped, so a file or folder name such as + ``x' or title contains '`` stays one literal instead of rewriting the query + (which would let a key select -- and ``__setitem__`` overwrite -- files + outside the store's folder). + + >>> _q("it's") + "'it\\\\'s'" + """ + return "'" + str(value).replace("\\", "\\\\").replace("'", "\\'") + "'" + + def _extract_folder_id(url: str) -> Optional[str]: """ Extract Google Drive folder ID from URL. @@ -79,7 +93,7 @@ def _extract_folder_id(url: str) -> Optional[str]: >>> _extract_folder_id('https://drive.google.com/drive/folders/ABC123') 'ABC123' """ - pattern = r"drive\.google\.com/drive/(?:u/\d+/)?folders/([^?]+)" + pattern = r"drive\.google\.com/drive/(?:u/\d+/)?folders/([A-Za-z0-9_-]+)" match = re.search(pattern, url) return match.group(1) if match else None @@ -103,7 +117,7 @@ def _resolve_file_id(url_or_id: str) -> str: file_id = _extract_file_id(url_or_id) if file_id: return file_id - if _FILE_ID_PATTERN.match(url_or_id): + if _FILE_ID_PATTERN.fullmatch(url_or_id): return url_or_id raise ValueError( f"Not a Google Drive file URL or file id: {url_or_id!r}. Expected something like " @@ -475,7 +489,7 @@ def _iter_folder_files( if max_levels is not None and level > max_levels: return - query = f"'{folder_id}' in parents and trashed=false" + query = f"{_q(folder_id)} in parents and trashed=false" for item in drive.ListFile({"q": query}).GetList(): name = item["title"] if not include_hidden and name.startswith("."): @@ -866,7 +880,7 @@ def upload( convert = self.convert_office parent_id = self._get_or_create_folders(key) filename = os.path.basename(key) - query = f"'{parent_id}' in parents and title='{filename}' and trashed=false" + query = f"{_q(parent_id)} in parents and title={_q(filename)} and trashed=false" existing = self._drive.ListFile({"q": query}).GetList() file_id = existing[0]["id"] if existing else None gfile = _upload_converting( @@ -899,8 +913,8 @@ def _get_or_create_folders(self, key: str) -> str: for folder_name in folder_parts: query = ( - f"'{current_id}' in parents " - f"and title='{folder_name}' " + f"{_q(current_id)} in parents " + f"and title={_q(folder_name)} " f"and mimeType='application/vnd.google-apps.folder' " f"and trashed=false" ) @@ -938,7 +952,7 @@ def __setitem__(self, key: str, value: bytes): filename = os.path.basename(key) # Check if file exists - query = f"'{parent_id}' in parents and title='{filename}' and trashed=false" + query = f"{_q(parent_id)} in parents and title={_q(filename)} and trashed=false" files = self._drive.ListFile({"q": query}).GetList() if files: diff --git a/tests/test_query_escaping.py b/tests/test_query_escaping.py new file mode 100644 index 0000000..f1aa82e --- /dev/null +++ b/tests/test_query_escaping.py @@ -0,0 +1,115 @@ +"""Keys are quoted as literals in Drive search queries, never spliced in. + +A key such as ``x' or title contains '`` used to rewrite the ``q`` query, so +lookups (and the "update existing file" branch of ``__setitem__``) could match +files outside the store's folder. File/folder ids taken from URLs are likewise +limited to Drive's id alphabet, so they cannot carry ``../`` into cache paths. +""" + +import pytest + +from pydrivedol.base import GDStore, _extract_file_id, _extract_folder_id, _q + + +class _File(dict): + def SetContentString(self, s): + self["content"] = s + + def SetContentFile(self, path): + self["content_file"] = path + + def Upload(self, param=None): + pass + + +class _List: + def GetList(self): + return [] + + +class _Drive: + def __init__(self): + self.queries = [] + + def ListFile(self, q): + self.queries.append(q["q"]) + return _List() + + def CreateFile(self, meta=None): + return _File( + meta or {}, id="new-id", alternateLink="https://drive.example/new-id" + ) + + +def _store(): + s = object.__new__(GDStore) + s.folder_id = "ROOTFOLDER123" + s._drive = _Drive() + s._refresh_cache = lambda: None + return s + + +HOSTILE = "x' or title contains '" + + +def test_hostile_filename_stays_one_literal(): + s = _store() + s[HOSTILE] = b"data" + q = s._drive.queries[-1] + assert q == ( + "'ROOTFOLDER123' in parents and title='x\\' or title contains \\'' " + "and trashed=false" + ) + + +def test_hostile_folder_name_stays_one_literal(): + s = _store() + s[f"{HOSTILE}/f.txt"] = b"data" + assert "title='x\\' or title contains \\'' " in s._drive.queries[0] + + +@pytest.mark.parametrize( + "value, literal", + [ + ("plain", "'plain'"), + ("it's", "'it\\'s'"), + ("a\\b", "'a\\\\b'"), + ("\\'", "'\\\\\\''"), + ], +) +def test_q(value, literal): + assert _q(value) == literal + + +@pytest.mark.parametrize( + "url, expected", + [ + ("https://drive.google.com/open?id=../secret/s.txt", None), + ("https://drive.google.com/file/d/AbC_12-x/view", "AbC_12-x"), + ("https://drive.google.com/uc?export=download&id=AbC123", "AbC123"), + ("https://drive.google.com/uc?id=GOOD123456&x_id=EVIL999999", "GOOD123456"), + ("https://drive.google.com/uc?x_id=EVIL999999&id=GOOD123456", "GOOD123456"), + ], +) +def test_file_ids_are_confined_to_the_id_alphabet(url, expected): + assert _extract_file_id(url) == expected + + +def test_folder_id_stops_at_the_id(): + assert ( + _extract_folder_id("https://drive.google.com/drive/folders/AbC123/") == "AbC123" + ) + + +def test_bare_id_with_trailing_newline_is_refused(): + from pydrivedol.base import _resolve_file_id + + with pytest.raises(ValueError): + _resolve_file_id("AAAAAAAAAAAA\n") + + +def test_upload_quotes_the_filename(): + s = _store() + s.convert_office = False + s.upload(HOSTILE, b"data") + assert "title='x\\' or title contains \\'' " in s._drive.queries[-1]