skald-connectors/connectors/drive/drive_mcp_server.py

#!/usr/bin/env python3
"""Google Drive MCP server (JSON-RPC 2.0 over stdio).

Tools:
  status           — self-check: credentials, token refresh, API reachability
  list_files       — list files in a folder (or root) with pagination
  search_files     — full-text + name search across Drive
  get_file_metadata — get detailed file info by ID
  upload_file      — upload a local file to Drive
  download_file    — download a file from Drive to local path
  delete_file      — permanently delete a file
  create_folder    — create a folder in Drive

Skald mode: Skald injects credentials via DRIVE_CREDS_JSON env var (authorized_user JSON).
Standalone mode: reads from GOOGLE_CREDS_PATH or ./secrets/google_creds.json.

Required OAuth scopes:
  https://www.googleapis.com/auth/drive      (full access — read/write/delete all files)
  or https://www.googleapis.com/auth/drive.file (access only to files this app creates)
"""
from __future__ import annotations

import base64
import io
import json
import mimetypes
import os
import sys
import traceback
from typing import Any, Callable

import google.auth
from google.auth.transport.requests import Request as GoogleRequest
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload


# ── Logging ─────────────────────────────────────────────────────────────────
def log(msg: str) -> None:
    print(f"[drive_mcp] {msg}", file=sys.stderr, flush=True)


# ── Auth helpers ────────────────────────────────────────────────────────────
_creds: Credentials | None = None
_drive_service: Any = None


def _load_credentials() -> Credentials | None:
    """Load credentials from env (Skald mode) or file (standalone mode)."""
    # 1. Skald mode: DRIVE_CREDS_JSON env var (full authorized_user JSON)
    env_json = os.environ.get("DRIVE_CREDS_JSON", "").strip()
    if env_json:
        try:
            data = json.loads(env_json)
            creds = Credentials.from_authorized_user_info(data)
            log("Loaded credentials from DRIVE_CREDS_JSON env var")
            return creds
        except Exception as e:
            log(f"Failed to parse DRIVE_CREDS_JSON: {e}")
            return None

    # 2. Standalone: GOOGLE_CREDS_PATH or default path
    path = os.environ.get("GOOGLE_CREDS_PATH", "")
    if not path:
        path = "./secrets/google_creds.json"
    if os.path.isfile(path):
        try:
            with open(path) as f:
                data = json.load(f)
            creds = Credentials.from_authorized_user_info(data)
            log(f"Loaded credentials from {path}")
            return creds
        except Exception as e:
            log(f"Failed to load creds from {path}: {e}")

    return None


def _ensure_service() -> tuple[bool, str]:
    """Ensure we have valid credentials and a Drive service. Returns (ok, message)."""
    global _creds, _drive_service
    if _drive_service is not None:
        return True, "ok"

    _creds = _load_credentials()
    if _creds is None:
        return False, "No Google Drive credentials found. Set DRIVE_CREDS_JSON or configure OAuth."

    # Refresh if expired
    if not _creds.valid:
        if _creds.expired and _creds.refresh_token:
            try:
                _creds.refresh(GoogleRequest())
                log("Token refreshed successfully")
            except Exception as e:
                return False, f"Token refresh failed: {e}"
        else:
            return False, "Credentials expired and no refresh token available."

    try:
        _drive_service = build("drive", "v3", credentials=_creds)
        return True, "ok"
    except Exception as e:
        return False, f"Failed to build Drive service: {e}"


# ── Tool implementations ────────────────────────────────────────────────────

DRIVE_FILE_FIELDS = "id, name, mimeType, size, createdTime, modifiedTime, parents, webViewLink, description, ownedByMe, owners(emailAddress), thumbnailLink"


def _list_files(folder_id: str | None, page_size: int, page_token: str | None) -> dict:
    """List files in a folder."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    q = f"'{folder_id}' in parents and trashed=false" if folder_id else "trashed=false"
    try:
        results = _drive_service.files().list(
            q=q,
            pageSize=min(page_size, 100),
            pageToken=page_token or "",
            fields=f"files({DRIVE_FILE_FIELDS}), nextPageToken",
            orderBy="folder,modifiedTime desc,name",
        ).execute()
        return {
            "files": results.get("files", []),
            "next_page_token": results.get("nextPageToken", None),
        }
    except HttpError as e:
        return {"error": f"Drive API error: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _search_files(query: str, max_results: int) -> dict:
    """Search files across Drive."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    q = f"name contains '{query}' and trashed=false"
    try:
        results = _drive_service.files().list(
            q=q,
            pageSize=min(max_results, 100),
            fields=f"files({DRIVE_FILE_FIELDS})",
            orderBy="modifiedTime desc",
        ).execute()
        return {"files": results.get("files", [])}
    except HttpError as e:
        return {"error": f"Drive API error: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _get_file_metadata(file_id: str) -> dict:
    """Get detailed metadata for a specific file."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    try:
        file = _drive_service.files().get(
            fileId=file_id,
            fields=DRIVE_FILE_FIELDS,
            supportsAllDrives=True,
        ).execute()
        return {"file": file}
    except HttpError as e:
        return {"error": f"Drive API error: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _upload_file(local_path: str, parent_id: str | None, name: str | None) -> dict:
    """Upload a local file to Drive."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    resolved_path = os.path.expanduser(local_path)
    if not os.path.isfile(resolved_path):
        return {"error": f"File not found: {local_path}"}

    file_name = name or os.path.basename(resolved_path)
    mime_type = mimetypes.guess_type(resolved_path)[0] or "application/octet-stream"

    body = {"name": file_name, "mimeType": mime_type}
    if parent_id:
        body["parents"] = [parent_id]

    try:
        media = MediaFileUpload(resolved_path, mimetype=mime_type, resumable=True)
        file = _drive_service.files().create(
            body=body,
            media_body=media,
            fields=DRIVE_FILE_FIELDS,
            supportsAllDrives=True,
        ).execute()
        return {"file": file}
    except HttpError as e:
        return {"error": f"Upload failed: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _download_file(file_id: str, output_path: str | None) -> dict:
    """Download a file from Drive. If no output_path given, returns base64 content."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    try:
        # Get file metadata first
        file = _drive_service.files().get(fileId=file_id, fields="id, name, mimeType, size").execute()
        file_name = file.get("name", "unknown")
        mime_type = file.get("mimeType", "application/octet-stream")

        request = _drive_service.files().get_media(fileId=file_id)

        if output_path:
            resolved = os.path.expanduser(output_path)
            os.makedirs(os.path.dirname(os.path.abspath(resolved)) or ".", exist_ok=True)
            with io.FileIO(resolved, "wb") as fh:
                downloader = MediaIoBaseDownload(fh, request)
                done = False
                while not done:
                    _, done = downloader.next_chunk()
            file_size = os.path.getsize(resolved)
            return {
                "file": {"id": file_id, "name": file_name, "mimeType": mime_type, "size": file_size},
                "saved_to": resolved,
            }
        else:
            # Return base64 content inline
            buf = io.BytesIO()
            downloader = MediaIoBaseDownload(buf, request)
            done = False
            while not done:
                _, done = downloader.next_chunk()
            b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
            return {
                "file": {"id": file_id, "name": file_name, "mimeType": mime_type},
                "content_base64": b64,
                "content_length": len(buf.getvalue()),
            }
    except HttpError as e:
        return {"error": f"Download failed: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _delete_file(file_id: str) -> dict:
    """Permanently delete a file."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    try:
        _drive_service.files().delete(fileId=file_id, supportsAllDrives=True).execute()
        return {"deleted": True, "file_id": file_id}
    except HttpError as e:
        return {"error": f"Delete failed: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _create_folder(name: str, parent_id: str | None) -> dict:
    """Create a folder in Drive."""
    ok, msg = _ensure_service()
    if not ok:
        return {"error": msg}

    body = {"name": name, "mimeType": "application/vnd.google-apps.folder"}
    if parent_id:
        body["parents"] = [parent_id]

    try:
        folder = _drive_service.files().create(
            body=body,
            fields=DRIVE_FILE_FIELDS,
            supportsAllDrives=True,
        ).execute()
        return {"folder": folder}
    except HttpError as e:
        return {"error": f"Create folder failed: {e.reason or e}"}
    except Exception as e:
        return {"error": str(e)}


def _status() -> dict:
    """Health check: credentials, token expiry, API reachability."""
    global _creds, _drive_service
    # Reset so we re-check from scratch
    _drive_service = None
    ok, msg = _ensure_service()
    if not ok:
        return {
            "ok": False,
            "message": msg,
        }
    expiry = _creds.expiry.isoformat() if _creds and _creds.expiry else "unknown"
    scopes = _creds.scopes if _creds else []
    return {
        "ok": True,
        "message": "Google Drive API is reachable",
        "details": {
            "expiry": expiry,
            "scopes": scopes,
            "token_valid": _creds.valid if _creds else False,
            "has_refresh_token": bool(_creds and _creds.refresh_token),
        },
    }


# ── Tool registry ──────────────────────────────────────────────────────────
TOOLS: dict[str, tuple[Callable, str, dict, str]] = {
    "status": (
        _status,
        "Check if the Drive API is reachable and credentials are valid.",
        {
            "type": "object",
            "properties": {},
            "required": [],
        },
    "Status"),
    "list_files": (
        _list_files,
        "List files in a folder (or root if no folder_id given). Returns paginated results.",
        {
            "type": "object",
            "properties": {
                "folder_id": {
                    "type": "string",
                    "description": "Folder ID (omit or empty for root). Pass 'root' for My Drive root.",
                    "default": None,
                },
                "page_size": {
                    "type": "integer",
                    "description": "Results per page (max 100, default 50)",
                    "default": 50,
                },
                "page_token": {
                    "type": "string",
                    "description": "Token for the next page (from previous list_files response)",
                    "default": None,
                },
            },
            "required": [],
        },
    "List Files"),
    "search_files": (
        _search_files,
        "Search files across Drive by name or content.",
        {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query (matched against file name)",
                },
                "max_results": {
                    "type": "integer",
                    "description": "Max results (default 20, max 100)",
                    "default": 20,
                },
            },
            "required": ["query"],
        },
    "Search Files"),
    "get_file_metadata": (
        _get_file_metadata,
        "Get detailed metadata for a specific file by ID.",
        {
            "type": "object",
            "properties": {
                "file_id": {
                    "type": "string",
                    "description": "Google Drive file ID (from URL or search results)",
                },
            },
            "required": ["file_id"],
        },
    "Get File Metadata"),
    "upload_file": (
        _upload_file,
        "Upload a local file to Google Drive.",
        {
            "type": "object",
            "properties": {
                "local_path": {
                    "type": "string",
                    "description": "Local path to the file to upload",
                },
                "parent_id": {
                    "type": "string",
                    "description": "Optional folder ID to upload into (omit for root)",
                    "default": None,
                },
                "name": {
                    "type": "string",
                    "description": "Optional custom name for the file on Drive (defaults to filename)",
                    "default": None,
                },
            },
            "required": ["local_path"],
        },
    "Upload File"),
    "download_file": (
        _download_file,
        "Download a file from Drive. Returns base64 content if no output_path, or saves to disk.",
        {
            "type": "object",
            "properties": {
                "file_id": {
                    "type": "string",
                    "description": "Google Drive file ID to download",
                },
                "output_path": {
                    "type": "string",
                    "description": "Optional local path to save the file (omit to get base64 content)",
                    "default": None,
                },
            },
            "required": ["file_id"],
        },
    "Download File"),
    "delete_file": (
        _delete_file,
        "Permanently delete a file from Google Drive.",
        {
            "type": "object",
            "properties": {
                "file_id": {
                    "type": "string",
                    "description": "Google Drive file ID to delete",
                },
            },
            "required": ["file_id"],
        },
    "Delete File"),
    "create_folder": (
        _create_folder,
        "Create a new folder in Google Drive.",
        {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Folder name",
                },
                "parent_id": {
                    "type": "string",
                    "description": "Optional parent folder ID (omit for My Drive root)",
                    "default": None,
                },
            },
            "required": ["name"],
        },
    "Create Folder"),
}


# ── JSON-RPC handler ───────────────────────────────────────────────────────

def _send_response(id_val: int | str | None, result: dict | None, error: dict | None = None) -> None:
    msg = {"jsonrpc": "2.0", "id": id_val}
    if error:
        msg["error"] = error
    else:
        msg["result"] = result
    sys.stdout.write(json.dumps(msg) + "\n")
    sys.stdout.flush()


def _handle_request(req: dict) -> None:
    rid = req.get("id")
    method = req.get("method", "")

    # JSON-RPC initialize
    if method == "initialize":
        _send_response(rid, {
            "protocolVersion": "2025-03-26",
            "capabilities": {
                "tools": {},
                "roots": {"listChanged": False},
            },
            "serverInfo": {"name": "google-drive-mcp", "version": "1.0.0"},
        })
        return

    # List available tools
    if method == "tools/list":
        tool_list = []
        for name, (fn, desc, schema, title) in sorted(TOOLS.items()):
            tool_list.append({
                "name": name,
                "description": desc,
                "title": title,
                "inputSchema": schema,
            })
        _send_response(rid, {"tools": tool_list})
        return

    # Call a tool
    if method == "tools/call":
        tool_name = req.get("params", {}).get("name", "")
        args = req.get("params", {}).get("arguments", {})

        if tool_name not in TOOLS:
            _send_response(rid, None, {
                "code": -32601,
                "message": f"Unknown tool: {tool_name}",
            })
            return

        try:
            fn, _, schema = TOOLS[tool_name]
            # Extract params matching the schema
            kwargs = {}
            for prop_name in schema.get("properties", {}):
                if prop_name in args:
                    kwargs[prop_name] = args[prop_name]

            result = fn(**kwargs)
            _send_response(rid, result)
        except Exception as e:
            log(f"Error in {tool_name}: {traceback.format_exc()}")
            _send_response(rid, None, {
                "code": -32603,
                "message": f"Internal error: {e}",
            })
        return

    # Notifications (no id) — ignore
    if rid is None:
        return

    _send_response(rid, None, {
        "code": -32601,
        "message": f"Method not found: {method}",
    })


# ── Main loop ──────────────────────────────────────────────────────────────
def main() -> None:
    log("Google Drive MCP server starting (JSON-RPC 2.0 over stdio)")

    # Announce server info on stderr
    log(f"Python: {sys.version}")
    log(f"Tools: {', '.join(sorted(TOOLS.keys()))}")

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            _handle_request(req)
        except json.JSONDecodeError as e:
            log(f"Invalid JSON: {e}")
            # Only respond if there's an id we can extract
            try:
                partial = json.loads(line[:line.rfind("}")+1] if "}" in line else "{}")
                rid = partial.get("id") if isinstance(partial, dict) else None
            except Exception:
                rid = None
            if rid is not None:
                _send_response(rid, None, {
                    "code": -32700,
                    "message": f"Parse error: {e}",
                })
        except Exception as e:
            log(f"Fatal error: {traceback.format_exc()}")


if __name__ == "__main__":
    main()