Google Docs

keyring>=24.0.0

#!/usr/bin/env python3
"""
OAuth token management for Google Docs API.
Standalone authentication - does not require the MCP server.
"""

import http.server
import json
import keyring
import secrets
import socket
import sys
import time
import webbrowser
from base64 import b64encode
from dataclasses import dataclass
from typing import Optional
from urllib.parse import urlencode, parse_qs, urlparse
import urllib.request
import urllib.error

# OAuth Configuration - uses same cloud function as MCP server
CLIENT_ID = "338689075775-o75k922vn5fdl18qergr96rp8g63e4d7.apps.googleusercontent.com"
CLOUD_FUNCTION_URL = "https://google-workspace-extension.geminicli.com"
REFRESH_ENDPOINT = f"{CLOUD_FUNCTION_URL}/refreshToken"

# Keyring configuration (cross-platform)
KEYCHAIN_SERVICE = "google-docs-skill-oauth"
KEYCHAIN_ACCOUNT = "main-account"

# Google Docs requires these scopes
SCOPES = [
    "https://www.googleapis.com/auth/documents",
    "https://www.googleapis.com/auth/drive",
    "https://www.googleapis.com/auth/userinfo.profile",
]

TOKEN_EXPIRY_BUFFER_MS = 5 * 60 * 1000  # 5 minutes


@dataclass
class TokenInfo:
    access_token: str
    refresh_token: Optional[str]
    expires_at: Optional[int]
    scope: Optional[str]


def get_available_port() -> int:
    """Find an available port for the OAuth callback server."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('localhost', 0))
        return s.getsockname()[1]


# =============================================================================
# Cross-Platform Token Storage using keyring
# =============================================================================

def get_tokens_from_keychain() -> Optional[TokenInfo]:
    """Read OAuth tokens from keyring."""
    try:
        data_str = keyring.get_password(KEYCHAIN_SERVICE, KEYCHAIN_ACCOUNT)
        if not data_str:
            return None
        data = json.loads(data_str)
        token = data.get("token", {})
        return TokenInfo(
            access_token=token.get("accessToken", ""),
            refresh_token=token.get("refreshToken"),
            expires_at=token.get("expiresAt"),
            scope=token.get("scope")
        )
    except (json.JSONDecodeError, keyring.errors.KeyringError):
        return None


def save_tokens_to_keychain(token_info: TokenInfo) -> bool:
    """Save OAuth tokens to keyring."""
    data = {
        "serverName": KEYCHAIN_ACCOUNT,
        "token": {
            "accessToken": token_info.access_token,
            "refreshToken": token_info.refresh_token,
            "tokenType": "Bearer",
            "scope": token_info.scope,
            "expiresAt": token_info.expires_at
        },
        "updatedAt": int(time.time() * 1000)
    }
    try:
        keyring.set_password(KEYCHAIN_SERVICE, KEYCHAIN_ACCOUNT, json.dumps(data))
        return True
    except keyring.errors.KeyringError:
        return False


def clear_tokens_from_keychain() -> bool:
    """Clear OAuth tokens from keyring."""
    try:
        keyring.delete_password(KEYCHAIN_SERVICE, KEYCHAIN_ACCOUNT)
        return True
    except (keyring.errors.PasswordDeleteError, keyring.errors.KeyringError):
        return False


# =============================================================================
# Token Management Functions
# =============================================================================

def get_tokens() -> Optional[TokenInfo]:
    """Get OAuth tokens from keyring."""
    return get_tokens_from_keychain()


def save_tokens(token_info: TokenInfo) -> bool:
    """Save OAuth tokens to keyring."""
    return save_tokens_to_keychain(token_info)


def clear_tokens() -> bool:
    """Clear OAuth tokens from keyring."""
    return clear_tokens_from_keychain()


def is_token_expired(token_info: TokenInfo) -> bool:
    """Check if the access token is expired or expiring soon."""
    if not token_info.expires_at:
        return False
    current_time_ms = int(time.time() * 1000)
    return token_info.expires_at < (current_time_ms + TOKEN_EXPIRY_BUFFER_MS)


def refresh_access_token(refresh_token: str) -> Optional[dict]:
    """Refresh access token using the cloud function."""
    try:
        data = json.dumps({"refresh_token": refresh_token}).encode('utf-8')
        req = urllib.request.Request(
            REFRESH_ENDPOINT,
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST"
        )

        with urllib.request.urlopen(req, timeout=30) as response:
            return json.loads(response.read().decode('utf-8'))
    except (urllib.error.URLError, urllib.error.HTTPError, json.JSONDecodeError) as e:
        print(f"Error refreshing token: {e}", file=sys.stderr)
        return None


class OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
    """HTTP request handler for OAuth callback."""

    token_info: Optional[TokenInfo] = None
    csrf_token: str = ""
    error: Optional[str] = None

    def log_message(self, format: str, *args) -> None:  # noqa: A002
        """Suppress default logging."""
        _ = format, args  # Unused

    def do_GET(self):
        """Handle the OAuth callback."""
        parsed = urlparse(self.path)

        if not parsed.path.startswith("/oauth2callback"):
            self.send_error(404)
            return

        params = parse_qs(parsed.query)

        # Validate CSRF token
        state = params.get("state", [""])[0]
        if state != OAuthCallbackHandler.csrf_token:
            OAuthCallbackHandler.error = "State mismatch - possible CSRF attack"
            self.send_response(400)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(b"<html><body><h1>Error: State mismatch</h1></body></html>")
            return

        # Check for errors
        if "error" in params:
            OAuthCallbackHandler.error = params.get("error_description", params["error"])[0]
            self.send_response(400)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(f"<html><body><h1>Error: {OAuthCallbackHandler.error}</h1></body></html>".encode())
            return

        # Extract tokens
        access_token = params.get("access_token", [""])[0]
        refresh_token = params.get("refresh_token", [""])[0]
        expiry_date = params.get("expiry_date", [""])[0]
        scope = params.get("scope", [""])[0]

        if access_token:
            OAuthCallbackHandler.token_info = TokenInfo(
                access_token=access_token,
                refresh_token=refresh_token or None,
                expires_at=int(expiry_date) if expiry_date else None,
                scope=scope or None
            )

            self.send_response(200)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(b"""<html><body>
                <h1>Authentication successful!</h1>
                <p>You can close this window and return to the terminal.</p>
                <script>window.close();</script>
            </body></html>""")
        else:
            OAuthCallbackHandler.error = "No access token received"
            self.send_response(400)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            self.wfile.write(b"<html><body><h1>Error: No access token received</h1></body></html>")


def perform_oauth_flow() -> Optional[TokenInfo]:
    """Perform the full OAuth flow with browser-based authentication."""
    port = get_available_port()
    redirect_uri = f"http://localhost:{port}/oauth2callback"
    csrf_token = secrets.token_hex(32)

    # Build state payload for cloud function
    state_payload = {
        "uri": redirect_uri,
        "manual": False,
        "csrf": csrf_token
    }
    state = b64encode(json.dumps(state_payload).encode()).decode()

    # Build auth URL
    auth_params = {
        "client_id": CLIENT_ID,
        "redirect_uri": CLOUD_FUNCTION_URL,  # Cloud function handles the secret
        "response_type": "code",
        "scope": " ".join(SCOPES),
        "access_type": "offline",
        "prompt": "consent",
        "state": state
    }
    auth_url = f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(auth_params)}"

    # Set up callback handler
    OAuthCallbackHandler.csrf_token = csrf_token
    OAuthCallbackHandler.token_info = None
    OAuthCallbackHandler.error = None

    # Start local server
    server = http.server.HTTPServer(('localhost', port), OAuthCallbackHandler)
    server.timeout = 300  # 5 minute timeout

    print(f"Opening browser for authentication...", file=sys.stderr)
    print(f"If browser doesn't open, visit: {auth_url}", file=sys.stderr)

    webbrowser.open(auth_url)

    # Wait for callback
    print("Waiting for authentication...", file=sys.stderr)
    while OAuthCallbackHandler.token_info is None and OAuthCallbackHandler.error is None:
        server.handle_request()

    server.server_close()

    if OAuthCallbackHandler.error:
        print(f"Authentication failed: {OAuthCallbackHandler.error}", file=sys.stderr)
        return None

    return OAuthCallbackHandler.token_info


def get_valid_access_token(interactive: bool = True) -> Optional[str]:
    """
    Get a valid access token.

    Args:
        interactive: If True, will prompt for OAuth flow if no tokens exist.
                    If False, will return None if no valid tokens.
    """
    token_info = get_tokens()

    # No tokens - need to authenticate
    if not token_info or not token_info.access_token:
        if not interactive:
            return None

        print("No OAuth tokens found. Starting authentication...", file=sys.stderr)
        token_info = perform_oauth_flow()

        if not token_info:
            return None

        save_tokens(token_info)
        print("Authentication successful!", file=sys.stderr)
        return token_info.access_token

    # Check if token needs refresh
    if is_token_expired(token_info) and token_info.refresh_token:
        print("Access token expired, refreshing...", file=sys.stderr)
        new_tokens = refresh_access_token(token_info.refresh_token)

        if new_tokens:
            token_info.access_token = new_tokens.get("access_token", token_info.access_token)
            token_info.expires_at = new_tokens.get("expiry_date", token_info.expires_at)
            save_tokens(token_info)
            print("Token refreshed successfully.", file=sys.stderr)
        else:
            # Refresh failed - try re-authenticating if interactive
            if interactive:
                print("Token refresh failed. Re-authenticating...", file=sys.stderr)
                token_info = perform_oauth_flow()
                if token_info:
                    save_tokens(token_info)
                else:
                    return None
            else:
                print("Warning: Failed to refresh token.", file=sys.stderr)

    return token_info.access_token if token_info else None


def logout() -> bool:
    """Clear stored OAuth tokens."""
    return clear_tokens()


def main():
    """CLI for auth operations."""
    import argparse

    parser = argparse.ArgumentParser(description="Google Docs OAuth management")
    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser("login", help="Authenticate with Google")
    subparsers.add_parser("logout", help="Clear stored tokens")
    subparsers.add_parser("token", help="Print current access token")
    subparsers.add_parser("status", help="Check authentication status")

    args = parser.parse_args()

    if args.command == "login":
        token_info = perform_oauth_flow()
        if token_info:
            save_tokens(token_info)
            print("Login successful!")
        else:
            print("Login failed.")
            sys.exit(1)

    elif args.command == "logout":
        if logout():
            print("Logged out successfully.")
        else:
            print("No tokens to clear.")

    elif args.command == "token":
        token = get_valid_access_token(interactive=False)
        if token:
            print(token)
        else:
            print("Not authenticated. Run: python auth.py login", file=sys.stderr)
            sys.exit(1)

    elif args.command == "status":
        token_info = get_tokens()
        if token_info and token_info.access_token:
            expired = is_token_expired(token_info)
            print(f"Status: Authenticated")
            print(f"Token expired: {expired}")
            if token_info.expires_at:
                expires_in = (token_info.expires_at - int(time.time() * 1000)) / 1000 / 60
                print(f"Expires in: {expires_in:.1f} minutes")
            print(f"Storage: keyring ({KEYCHAIN_SERVICE})")
        else:
            print("Status: Not authenticated")
            print("Run: python auth.py login")
            sys.exit(1)

    else:
        parser.print_help()


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Google Docs API operations.
Lightweight alternative to the full Google Workspace MCP server.
"""

import argparse
import json
import re
import sys
import urllib.request
import urllib.error
import urllib.parse
from typing import Optional

from auth import get_valid_access_token

DOCS_API_BASE = "https://docs.googleapis.com/v1"
DRIVE_API_BASE = "https://www.googleapis.com/drive/v3"


def api_request(method: str, url: str, data: Optional[dict] = None, params: Optional[dict] = None) -> dict:
    """Make an authenticated request to Google APIs."""
    token = get_valid_access_token()
    if not token:
        return {"error": "Failed to get access token"}

    if params:
        url += "?" + urllib.parse.urlencode(params)

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }

    body = json.dumps(data).encode('utf-8') if data else None

    try:
        req = urllib.request.Request(url, data=body, headers=headers, method=method)
        with urllib.request.urlopen(req, timeout=30) as response:
            return json.loads(response.read().decode('utf-8'))
    except urllib.error.HTTPError as e:
        error_body = e.read().decode('utf-8') if e.fp else str(e)
        return {"error": f"HTTP {e.code}: {error_body}"}
    except urllib.error.URLError as e:
        return {"error": f"Request failed: {e.reason}"}
    except json.JSONDecodeError:
        return {"error": "Invalid JSON response"}


def extract_doc_id(doc_id_or_url: str) -> str:
    """Extract document ID from a URL or return the ID as-is."""
    # Pattern to match Google Docs URLs
    patterns = [
        r'docs\.google\.com/document/d/([a-zA-Z0-9_-]+)',
        r'^([a-zA-Z0-9_-]+)$'
    ]
    for pattern in patterns:
        match = re.search(pattern, doc_id_or_url)
        if match:
            return match.group(1)
    return doc_id_or_url


def create_doc(title: str, markdown: Optional[str] = None, folder_name: Optional[str] = None) -> dict:
    """Create a new Google Doc, optionally with markdown content."""
    if markdown:
        # Create doc via Drive API with HTML content (converted from markdown)
        # For simplicity, we'll create a blank doc and then insert text
        # Full markdown support would require HTML conversion like the MCP server
        result = api_request(
            "POST",
            f"{DOCS_API_BASE}/documents",
            data={"title": title}
        )

        if "error" in result:
            return result

        doc_id = result.get("documentId")
        if doc_id and markdown:
            # Insert the markdown as plain text (basic implementation)
            insert_result = insert_text(doc_id, markdown)
            if "error" in insert_result:
                return {"documentId": doc_id, "title": title, "warning": "Doc created but text insertion failed"}

        return {"documentId": doc_id, "title": result.get("title", title)}
    else:
        # Create blank doc
        result = api_request(
            "POST",
            f"{DOCS_API_BASE}/documents",
            data={"title": title}
        )

        if "error" in result:
            return result

        return {"documentId": result.get("documentId"), "title": result.get("title", title)}


def find_docs(query: str, page_size: int = 10, page_token: Optional[str] = None) -> dict:
    """Find Google Docs by title search using Drive API."""
    # Build the query for Drive API
    # Search for documents with the query in the name
    drive_query = f"mimeType='application/vnd.google-apps.document' and name contains '{query}' and trashed=false"

    params = {
        "q": drive_query,
        "pageSize": page_size,
        "fields": "nextPageToken, files(id, name, modifiedTime, createdTime)"
    }
    if page_token:
        params["pageToken"] = page_token

    result = api_request("GET", f"{DRIVE_API_BASE}/files", params=params)

    if "error" in result:
        return result

    return {
        "files": result.get("files", []),
        "nextPageToken": result.get("nextPageToken")
    }


def get_text(document_id: str) -> dict:
    """Get the text content of a Google Doc."""
    doc_id = extract_doc_id(document_id)

    params = {
        "fields": "title,body.content"
    }

    result = api_request("GET", f"{DOCS_API_BASE}/documents/{doc_id}", params=params)

    if "error" in result:
        return result

    # Extract text from the document structure
    text = ""
    body_content = result.get("body", {}).get("content", [])

    for element in body_content:
        text += _read_structural_element(element)

    return {
        "documentId": doc_id,
        "title": result.get("title", ""),
        "text": text
    }


def _read_structural_element(element: dict) -> str:
    """Extract text from a structural element."""
    text = ""
    if "paragraph" in element:
        for p_element in element["paragraph"].get("elements", []):
            if "textRun" in p_element:
                text += p_element["textRun"].get("content", "")
    elif "table" in element:
        for row in element["table"].get("tableRows", []):
            for cell in row.get("tableCells", []):
                for cell_content in cell.get("content", []):
                    text += _read_structural_element(cell_content)
    return text


def append_text(document_id: str, text: str) -> dict:
    """Append text to the end of a Google Doc."""
    doc_id = extract_doc_id(document_id)

    # First, get the document to find the end index
    params = {"fields": "body.content"}
    doc_result = api_request("GET", f"{DOCS_API_BASE}/documents/{doc_id}", params=params)

    if "error" in doc_result:
        return doc_result

    # Find the end index
    body_content = doc_result.get("body", {}).get("content", [])
    if body_content:
        last_element = body_content[-1]
        end_index = last_element.get("endIndex", 1)
    else:
        end_index = 1

    # Insert at the end (before the final newline)
    location_index = max(1, end_index - 1)

    # Batch update to insert text
    requests = [
        {
            "insertText": {
                "location": {"index": location_index},
                "text": text
            }
        }
    ]

    result = api_request(
        "POST",
        f"{DOCS_API_BASE}/documents/{doc_id}:batchUpdate",
        data={"requests": requests}
    )

    if "error" in result:
        return result

    return {"success": True, "documentId": doc_id, "message": "Text appended successfully"}


def insert_text(document_id: str, text: str) -> dict:
    """Insert text at the beginning of a Google Doc."""
    doc_id = extract_doc_id(document_id)

    # Insert at index 1 (beginning of document content)
    requests = [
        {
            "insertText": {
                "location": {"index": 1},
                "text": text
            }
        }
    ]

    result = api_request(
        "POST",
        f"{DOCS_API_BASE}/documents/{doc_id}:batchUpdate",
        data={"requests": requests}
    )

    if "error" in result:
        return result

    return {"success": True, "documentId": doc_id, "message": "Text inserted successfully"}


def replace_text(document_id: str, find_text: str, replace_with: str) -> dict:
    """Replace all occurrences of text in a Google Doc."""
    doc_id = extract_doc_id(document_id)

    # Use replaceAllText request
    requests = [
        {
            "replaceAllText": {
                "containsText": {
                    "text": find_text,
                    "matchCase": True
                },
                "replaceText": replace_with
            }
        }
    ]

    result = api_request(
        "POST",
        f"{DOCS_API_BASE}/documents/{doc_id}:batchUpdate",
        data={"requests": requests}
    )

    if "error" in result:
        return result

    # Get the number of replacements made
    replies = result.get("replies", [])
    occurrences_changed = 0
    if replies:
        occurrences_changed = replies[0].get("replaceAllText", {}).get("occurrencesChanged", 0)

    return {
        "success": True,
        "documentId": doc_id,
        "occurrencesChanged": occurrences_changed,
        "message": f"Replaced {occurrences_changed} occurrence(s)"
    }


def main():
    parser = argparse.ArgumentParser(description="Google Docs API operations")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # create
    create_parser = subparsers.add_parser("create", help="Create a new Google Doc")
    create_parser.add_argument("title", help="Title for the new document")
    create_parser.add_argument("--content", help="Initial content (plain text or markdown)")

    # find
    find_parser = subparsers.add_parser("find", help="Find docs by title search")
    find_parser.add_argument("query", help="Search query")
    find_parser.add_argument("--limit", type=int, default=10, help="Max results to return")
    find_parser.add_argument("--page-token", help="Pagination token")

    # get-text
    get_text_parser = subparsers.add_parser("get-text", help="Get text content of a doc")
    get_text_parser.add_argument("document_id", help="Document ID or URL")

    # append-text
    append_text_parser = subparsers.add_parser("append-text", help="Append text to a doc")
    append_text_parser.add_argument("document_id", help="Document ID or URL")
    append_text_parser.add_argument("text", help="Text to append")

    # insert-text
    insert_text_parser = subparsers.add_parser("insert-text", help="Insert text at beginning of a doc")
    insert_text_parser.add_argument("document_id", help="Document ID or URL")
    insert_text_parser.add_argument("text", help="Text to insert")

    # replace-text
    replace_text_parser = subparsers.add_parser("replace-text", help="Replace text in a doc")
    replace_text_parser.add_argument("document_id", help="Document ID or URL")
    replace_text_parser.add_argument("find", help="Text to find")
    replace_text_parser.add_argument("replace", help="Replacement text")

    args = parser.parse_args()

    if args.command == "create":
        result = create_doc(args.title, args.content)
    elif args.command == "find":
        result = find_docs(args.query, args.limit, args.page_token)
    elif args.command == "get-text":
        result = get_text(args.document_id)
    elif args.command == "append-text":
        result = append_text(args.document_id, args.text)
    elif args.command == "insert-text":
        result = insert_text(args.document_id, args.text)
    elif args.command == "replace-text":
        result = replace_text(args.document_id, args.find, args.replace)
    else:
        result = {"error": f"Unknown command: {args.command}"}

    print(json.dumps(result, indent=2))

    if isinstance(result, dict) and "error" in result:
        sys.exit(1)


if __name__ == "__main__":
    main()

keyring>=24.0.0