Gdrive file type fileter - Fixing problems detected by coderabbitai

Author: fmgs31

gdrive/drive_helpers.py

@@ -218,6 +218,9 @@ def build_drive_list_params(
 SHORTCUT_MIME_TYPE = "application/vnd.google-apps.shortcut"
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
 
+# RFC 6838 token-style MIME type validation (safe for Drive query interpolation).
+MIME_TYPE_PATTERN = re.compile(r"^[A-Za-z0-9!#$&^_.+-]+/[A-Za-z0-9!#$&^_.+-]+$")
+
 # Mapping from friendly type names to Google Drive MIME types.
 # Raw MIME type strings (containing '/') are always accepted as-is.
 FILE_TYPE_MIME_MAP: Dict[str, str] = {
@@ -256,9 +259,14 @@ def resolve_file_type_mime(file_type: str) -> str:
     Raises:
         ValueError: If the value is not a recognised friendly name and contains no '/'.
     """
-    if "/" in file_type:
-        return file_type
-    lower = file_type.lower()
+    normalized = file_type.strip()
+    if "/" in normalized:
+        if not MIME_TYPE_PATTERN.fullmatch(normalized):
+            raise ValueError(
+                f"Invalid MIME type '{file_type}'. Expected format like 'application/pdf'."
+            )
+        return normalized
+    lower = normalized.lower()
     if lower not in FILE_TYPE_MIME_MAP:
         valid = ", ".join(sorted(FILE_TYPE_MIME_MAP.keys()))
         raise ValueError(

gdrive/drive_tools.py

@@ -76,8 +76,10 @@ async def search_drive_files(
                                  If 'drive_id' is specified and 'corpora' is None, it defaults to 'drive'.
                                  Otherwise, Drive API default behavior applies. Prefer 'user' or 'drive' over 'allDrives' for efficiency.
         file_type (Optional[str]): Restrict results to a specific file type. Accepts a friendly
-                                   name ('folder', 'document', 'doc', 'spreadsheet', 'sheet', 'presentation', 'slides', 'form', 'drawing', 'pdf', 'shortcut') or 
-                                   any raw MIME type string (e.g. 'application/pdf'). Defaults to None (all types).
+                                   name ('folder', 'document'/'doc', 'spreadsheet'/'sheet',
+                                   'presentation'/'slides', 'form', 'drawing', 'pdf', 'shortcut',
+                                   'script', 'site', 'jam'/'jamboard') or any raw MIME type
+                                   string (e.g. 'application/pdf'). Defaults to None (all types).
 
     Returns:
         str: A formatted list of found files/folders with their details (ID, name, type, size, modified time, link).
@@ -105,7 +107,7 @@ async def search_drive_files(
 
     if file_type is not None:
         mime = resolve_file_type_mime(file_type)
-        final_query += f" and mimeType = '{mime}'"
+        final_query = f"({final_query}) and mimeType = '{mime}'"
         logger.info(f"[search_drive_files] Added mimeType filter: '{mime}'")
 
     list_params = build_drive_list_params(
@@ -439,10 +441,10 @@ async def list_drive_items(
         include_items_from_all_drives (bool): Whether items from all accessible shared drives should be included if `drive_id` is not set. Defaults to True.
         corpora (Optional[str]): Corpus to query ('user', 'drive', 'allDrives'). If `drive_id` is set and `corpora` is None, 'drive' is used. If None and no `drive_id`, API defaults apply.
         file_type (Optional[str]): Restrict results to a specific file type. Accepts a friendly
-                                   name ('folder', 'document', 'doc', 'spreadsheet', 'sheet',
-                                   'presentation', 'slides', 'form', 'drawing', 'pdf',
-                                   'shortcut') or any raw MIME type string
-                                   (e.g. 'application/pdf'). Defaults to None (all types).
+                                   name ('folder', 'document'/'doc', 'spreadsheet'/'sheet',
+                                   'presentation'/'slides', 'form', 'drawing', 'pdf', 'shortcut',
+                                   'script', 'site', 'jam'/'jamboard') or any raw MIME type
+                                   string (e.g. 'application/pdf'). Defaults to None (all types).
 
     Returns:
         str: A formatted list of files/folders in the specified folder.
@@ -456,7 +458,7 @@ async def list_drive_items(
 
     if file_type is not None:
         mime = resolve_file_type_mime(file_type)
-        final_query += f" and mimeType = '{mime}'"
+        final_query = f"({final_query}) and mimeType = '{mime}'"
         logger.info(f"[list_drive_items] Added mimeType filter: '{mime}'")
 
     list_params = build_drive_list_params(

tests/gdrive/test_drive_tools.py

@@ -55,9 +55,10 @@ async def test_create_drive_folder():
 # Helpers
 # ---------------------------------------------------------------------------
 
-def _unwrap(tool):
-    """Unwrap a FunctionTool + decorator chain to the original async function."""
-    fn = tool.fn
+def _unwrap(fn):
+    """Unwrap a FunctionTool or plain-function decorator chain to the original async function."""
+    if hasattr(fn, "fn"):
+        fn = fn.fn  # FunctionTool wrapper (some server versions)
     while hasattr(fn, "__wrapped__"):
         fn = fn.__wrapped__
     return fn
@@ -396,3 +397,49 @@ async def test_list_items_file_type_unknown_raises(mock_resolve_folder):
             folder_id="root",
             file_type="unknowntype",
         )
+
+
+# ---------------------------------------------------------------------------
+# OR-precedence grouping
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_search_or_query_is_grouped_before_mime_filter():
+    """An OR structured query is wrapped in parentheses so the MIME filter binds correctly."""
+    mock_service = Mock()
+    mock_service.files().list().execute.return_value = {"files": []}
+
+    await _unwrap(search_drive_files)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        query="name contains 'a' or name contains 'b'",
+        file_type="document",
+    )
+
+    q = mock_service.files.return_value.list.call_args.kwargs["q"]
+    # Without grouping this would be:  name contains 'a' or name contains 'b' and mimeType = ...
+    # The 'and' would only bind to the second term, leaking the first term through unfiltered.
+    assert q.startswith("(")
+    assert "name contains 'a' or name contains 'b'" in q
+    assert ") and mimeType = 'application/vnd.google-apps.document'" in q
+
+
+# ---------------------------------------------------------------------------
+# MIME type validation
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_file_type_mime_invalid_mime_raises():
+    """A raw string with '/' but containing quotes raises ValueError."""
+    from gdrive.drive_helpers import resolve_file_type_mime
+
+    with pytest.raises(ValueError, match="Invalid MIME type"):
+        resolve_file_type_mime("application/pdf' or '1'='1")
+
+
+def test_resolve_file_type_mime_strips_whitespace():
+    """Leading/trailing whitespace is stripped from raw MIME strings."""
+    from gdrive.drive_helpers import resolve_file_type_mime
+
+    assert resolve_file_type_mime("  application/pdf  ") == "application/pdf"