replica-omnisciente/scripts/gws/gws.py

#!/usr/bin/env python3
"""Google Workspace ingestion tool for Aurelio research studies.

Supports two auth modes:
  1. OAuth2 user consent (for individual accounts)
  2. Service account with domain-wide delegation (preferred for Workspace-wide studies)

Usage:
    # Service account (place key at ~/.aurelio/gws/service_account.json)
    python3 scripts/gws/gws.py fetch-emails --emails fabio@savearth.io,claudio@savearth.io,joao@savearth.io --out-dir ./gws-output

    # OAuth2 (fallback)
    python3 scripts/gws/gws.py auth --email fabio@savearth.io
    python3 scripts/gws/gws.py auth-code --email fabio@savearth.io --code '<code>'
    python3 scripts/gws/gws.py fetch-emails --emails ... --out-dir ./gws-output
"""
import argparse
import base64
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google.oauth2.service_account import Credentials as ServiceAccountCredentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError

# Read-only scopes for research study
SCOPES = [
    "https://www.googleapis.com/auth/gmail.readonly",
    "https://www.googleapis.com/auth/calendar.readonly",
    "https://www.googleapis.com/auth/drive.readonly",
    "https://www.googleapis.com/auth/userinfo.profile",
    "https://www.googleapis.com/auth/contacts.readonly",
]

DEFAULT_CLIENT_SECRET = os.environ.get(
    "GWS_CLIENT_SECRET",
    str(Path.home() / ".aurelio" / "gws" / "client_secret.json"),
)
SERVICE_ACCOUNT_FILE = os.environ.get(
    "GWS_SERVICE_ACCOUNT",
    str(Path.home() / ".aurelio" / "gws" / "service_account.json"),
)
TOKEN_DIR = Path(os.environ.get("GWS_TOKEN_DIR", Path.home() / ".aurelio" / "gws" / "tokens"))


def _iso_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _client_secret_path() -> Path:
    p = Path(DEFAULT_CLIENT_SECRET)
    if not p.exists():
        raise SystemExit(
            f"OAuth client secret not found at {p}. "
            "Set GWS_CLIENT_SECRET or place it at ~/.aurelio/gws/client_secret.json"
        )
    return p


def _service_account_path() -> Path | None:
    p = Path(SERVICE_ACCOUNT_FILE)
    return p if p.exists() else None


def _token_path(email: str) -> Path:
    TOKEN_DIR.mkdir(parents=True, exist_ok=True)
    return TOKEN_DIR / f"{email}.json"


def _load_creds(email: str) -> Credentials | None:
    tp = _token_path(email)
    if tp.exists():
        return Credentials.from_authorized_user_file(str(tp), SCOPES)
    return None


def _save_creds(email: str, creds: Credentials) -> None:
    tp = _token_path(email)
    tp.write_text(creds.to_json())
    os.chmod(tp, 0o600)


AnyCredentials = Credentials | ServiceAccountCredentials


def _ensure_creds(email: str) -> AnyCredentials:
    sa_path = _service_account_path()
    if sa_path:
        creds = ServiceAccountCredentials.from_service_account_file(
            str(sa_path), scopes=SCOPES, subject=email
        )
        creds.refresh(Request())
        return creds

    creds = _load_creds(email)
    if creds and creds.expired and creds.refresh_token:
        creds.refresh(Request())
        _save_creds(email, creds)
        return creds
    if creds and creds.valid:
        return creds

    flow = InstalledAppFlow.from_client_secrets_file(str(_client_secret_path()), SCOPES)
    creds = flow.run_local_server(port=8080)
    _save_creds(email, creds)
    return creds


def cmd_auth(args: argparse.Namespace) -> int:
    email = args.email
    if _service_account_path():
        print(f"Service account mode active; no OAuth needed for {email}")
        return 0

    creds = _load_creds(email)
    if creds and creds.valid:
        print(f"Already authenticated for {email}")
        return 0
    if creds and creds.expired and creds.refresh_token:
        creds.refresh(Request())
        _save_creds(email, creds)
        print(f"Refreshed token for {email}")
        return 0

    flow = InstalledAppFlow.from_client_secrets_file(str(_client_secret_path()), SCOPES)
    auth_url, _ = flow.authorization_url(access_type="offline", prompt="consent")
    print(f"Authorize {email} here:\n{auth_url}\n")
    print("Then run:")
    print(f"  python3 scripts/gws/gws.py auth-code --email {email} --code '<code>'")
    return 0


def cmd_auth_code(args: argparse.Namespace) -> int:
    email = args.email
    code = args.code
    flow = InstalledAppFlow.from_client_secrets_file(str(_client_secret_path()), SCOPES)
    flow.fetch_token(code=code)
    creds = flow.credentials
    _save_creds(email, creds)
    print(f"Authenticated {email}; token saved to {_token_path(email)}")
    return 0


def _build(service_name: str, version: str, creds: AnyCredentials) -> Any:
    return build(service_name, version, credentials=creds, cache_discovery=False)


def _query_for_addresses(addresses: Iterable[str]) -> str:
    parts = []
    for a in addresses:
        parts.append(f"from:{a} OR to:{a} OR cc:{a}")
    return " OR ".join(f"({p})" for p in parts)


def cmd_fetch_emails(args: argparse.Namespace) -> int:
    emails = args.emails.split(",")
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    actor = args.actor or emails[0]

    creds = _ensure_creds(actor)
    service = _build("gmail", "v1", creds)
    query = _query_for_addresses(emails)
    if args.since:
        query += f" after:{args.since}"
    if args.before:
        query += f" before:{args.before}"

    results = []
    page_token = None
    while True:
        resp = (
            service.users()
            .messages()
            .list(userId="me", q=query, pageToken=page_token, maxResults=min(args.max_results, 500))
            .execute()
        )
        messages = resp.get("messages", [])
        for m in messages:
            msg = (
                service.users()
                .messages()
                .get(userId="me", id=m["id"], format="full")
                .execute()
            )
            payload = msg.get("payload", {})
            headers = {h["name"].lower(): h["value"] for h in payload.get("headers", [])}
            body = ""
            parts = payload.get("parts", [])
            for part in parts:
                if part.get("mimeType") == "text/plain" and "data" in part.get("body", {}):
                    body = base64.urlsafe_b64decode(part["body"]["data"]).decode("utf-8", errors="replace")
                    break
            results.append(
                {
                    "id": msg["id"],
                    "thread_id": msg["threadId"],
                    "internal_date": msg.get("internalDate"),
                    "subject": headers.get("subject", ""),
                    "from": headers.get("from", ""),
                    "to": headers.get("to", ""),
                    "cc": headers.get("cc", ""),
                    "body": body,
                }
            )
            if len(results) >= args.max_results:
                break
        page_token = resp.get("nextPageToken")
        if not page_token or len(results) >= args.max_results:
            break

    out_file = out_dir / f"emails-{actor}-{_iso_now()}.jsonl"
    with open(out_file, "w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"Wrote {len(results)} emails to {out_file}")
    return 0


def cmd_fetch_calendar(args: argparse.Namespace) -> int:
    emails = args.emails.split(",")
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    actor = args.actor or emails[0]

    creds = _ensure_creds(actor)
    service = _build("calendar", "v3", creds)

    time_min = args.since or "2020-01-01T00:00:00Z"
    time_max = args.before or "2030-01-01T00:00:00Z"

    results = []
    page_token = None
    while True:
        resp = (
            service.events()
            .list(
                calendarId="primary",
                timeMin=time_min,
                timeMax=time_max,
                pageToken=page_token,
                maxResults=min(args.max_results, 2500),
                showDeleted=False,
                singleEvents=True,
                orderBy="startTime",
            )
            .execute()
        )
        for item in resp.get("items", []):
            attendees = [a.get("email", "") for a in item.get("attendees", [])]
            if args.require_match and not any(e in attendees for e in emails):
                continue
            results.append(
                {
                    "id": item["id"],
                    "summary": item.get("summary", ""),
                    "description": item.get("description", ""),
                    "start": item.get("start", {}),
                    "end": item.get("end", {}),
                    "creator": item.get("creator", {}),
                    "attendees": attendees,
                    "link": item.get("htmlLink", ""),
                }
            )
            if len(results) >= args.max_results:
                break
        page_token = resp.get("nextPageToken")
        if not page_token or len(results) >= args.max_results:
            break

    out_file = out_dir / f"calendar-{actor}-{_iso_now()}.jsonl"
    with open(out_file, "w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"Wrote {len(results)} calendar events to {out_file}")
    return 0


def cmd_fetch_documents(args: argparse.Namespace) -> int:
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    actor = args.actor

    creds = _ensure_creds(actor)
    service = _build("drive", "v3", creds)

    results = []
    page_token = None
    query = "trashed = false"
    if args.mime_type:
        query += f" and mimeType = '{args.mime_type}'"
    while True:
        resp = (
            service.files()
            .list(
                q=query,
                pageToken=page_token,
                pageSize=min(args.max_results, 1000),
                fields="nextPageToken, files(id, name, mimeType, createdTime, modifiedTime, owners, webViewLink, size)",
            )
            .execute()
        )
        for item in resp.get("files", []):
            results.append(
                {
                    "id": item["id"],
                    "name": item.get("name", ""),
                    "mime_type": item.get("mimeType", ""),
                    "created_time": item.get("createdTime", ""),
                    "modified_time": item.get("modifiedTime", ""),
                    "owners": [o.get("emailAddress", "") for o in item.get("owners", [])],
                    "link": item.get("webViewLink", ""),
                    "size": item.get("size", ""),
                }
            )
            if len(results) >= args.max_results:
                break
        page_token = resp.get("nextPageToken")
        if not page_token or len(results) >= args.max_results:
            break

    out_file = out_dir / f"documents-{actor}-{_iso_now()}.jsonl"
    with open(out_file, "w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"Wrote {len(results)} documents to {out_file}")
    return 0


def cmd_to_brain(args: argparse.Namespace) -> int:
    """Normalize gws output into brain pages."""
    out_dir = Path(args.out_dir)
    study_dir = out_dir / "brain"
    study_dir.mkdir(parents=True, exist_ok=True)

    for f in sorted(out_dir.glob("emails-*.jsonl")):
        with open(f, encoding="utf-8") as fh:
            emails = [json.loads(line) for line in fh]
        actor = f.stem.split("-", 2)[1]
        page = {
            "title": f"Savearth communication study — {actor} emails",
            "source": "gws",
            "count": len(emails),
            "emails": emails[:50],
        }
        (study_dir / f"emails-{actor}.json").write_text(json.dumps(page, ensure_ascii=False, indent=2))
    print(f"Normalized brain pages to {study_dir}")
    return 0


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Google Workspace ingestion for Aurelio")
    sub = p.add_subparsers(dest="cmd", required=True)

    a = sub.add_parser("auth", help="Print authorization URL for an email")
    a.add_argument("--email", required=True)
    a.set_defaults(func=cmd_auth)

    ac = sub.add_parser("auth-code", help="Exchange authorization code for token")
    ac.add_argument("--email", required=True)
    ac.add_argument("--code", required=True)
    ac.set_defaults(func=cmd_auth_code)

    fe = sub.add_parser("fetch-emails", help="Fetch Gmail messages involving addresses")
    fe.add_argument("--emails", required=True, help="Comma-separated addresses")
    fe.add_argument("--actor", default=None, help="Workspace account to authenticate as")
    fe.add_argument("--since", default=None, help="YYYY/MM/DD")
    fe.add_argument("--before", default=None, help="YYYY/MM/DD")
    fe.add_argument("--max-results", type=int, default=500)
    fe.add_argument("--out-dir", default="./gws-output")
    fe.set_defaults(func=cmd_fetch_emails)

    fc = sub.add_parser("fetch-calendar", help="Fetch Calendar events")
    fc.add_argument("--emails", required=True)
    fc.add_argument("--actor", default=None)
    fc.add_argument("--since", default=None)
    fc.add_argument("--before", default=None)
    fc.add_argument("--max-results", type=int, default=500)
    fc.add_argument("--out-dir", default="./gws-output")
    fc.add_argument("--require-match", action="store_true", default=True)
    fc.set_defaults(func=cmd_fetch_calendar)

    fd = sub.add_parser("fetch-documents", help="Fetch Drive metadata")
    fd.add_argument("--actor", required=True)
    fd.add_argument("--mime-type", default=None, help="e.g. application/vnd.google-apps.document")
    fd.add_argument("--max-results", type=int, default=500)
    fd.add_argument("--out-dir", default="./gws-output")
    fd.set_defaults(func=cmd_fetch_documents)

    tb = sub.add_parser("to-brain", help="Convert gws output to brain pages")
    tb.add_argument("--out-dir", default="./gws-output")
    tb.set_defaults(func=cmd_to_brain)

    args = p.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())