backend/worker/gmail_auth.py

from __future__ import annotations

import argparse
import os
from pathlib import Path

from dotenv import load_dotenv
from google_auth_oauthlib.flow import InstalledAppFlow

try:
    from backend.worker.hf_env_files import resolve_json_or_path
except ModuleNotFoundError:
    import sys

    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from backend.worker.hf_env_files import resolve_json_or_path


SCOPES = [
    "https://www.googleapis.com/auth/gmail.modify",
    "https://www.googleapis.com/auth/gmail.send",
]


def _repo_root() -> Path:
    return Path(__file__).resolve().parents[2]


def _alias_env(primary: str, fallback: str) -> None:
    if (os.environ.get(primary) or "").strip():
        return
    fb = (os.environ.get(fallback) or "").strip()
    if fb:
        os.environ[primary] = fb


def main() -> None:
    parser = argparse.ArgumentParser(description="Interactive Gmail OAuth token generator.")
    parser.add_argument(
        "--credentials",
        help="Path to OAuth client credentials JSON (overrides env).",
        default="",
    )
    parser.add_argument(
        "--token",
        help="Path to write token JSON (overrides env).",
        default="",
    )
    parser.add_argument(
        "--console",
        action="store_true",
        help="Use console-based auth (no local server).",
    )
    args = parser.parse_args()

    repo_root = _repo_root()

    # Load local env files if present (helps local dev; HF will use injected env vars).
    load_dotenv(repo_root / ".env", override=False)
    load_dotenv(repo_root / "backend" / ".env", override=False)

    if args.credentials:
        os.environ["GMAIL_CREDENTIALS_JSON"] = args.credentials
    if args.token:
        os.environ["GMAIL_TOKEN_JSON"] = args.token

    # Back-compat with older env var names.
    _alias_env("GMAIL_CREDENTIALS_JSON", "PDF_PIPELINE_GMAIL_CREDENTIALS_JSON")
    _alias_env("GMAIL_TOKEN_JSON", "PDF_PIPELINE_GMAIL_TOKEN_JSON")

    backend_dir = repo_root / "backend"
    default_creds = backend_dir / "credentials.json"
    default_token = backend_dir / "token.json"

    creds_path = resolve_json_or_path("GMAIL_CREDENTIALS_JSON", default_creds, Path("/tmp/credentials.json"))
    token_path = resolve_json_or_path("GMAIL_TOKEN_JSON", default_token, Path("/tmp/token.json"))

    if not creds_path.exists() and default_creds.exists():
        print(f"[gmail_auth] WARN: {creds_path} not found; using {default_creds}")
        creds_path = default_creds

    if not creds_path.exists():
        raise FileNotFoundError(f"Missing OAuth client json: {creds_path}")

    flow = InstalledAppFlow.from_client_secrets_file(str(creds_path), SCOPES)

    if args.console:
        creds = flow.run_console()
    else:
        creds = flow.run_local_server(port=0, access_type="offline", prompt="consent")

    token_path.parent.mkdir(parents=True, exist_ok=True)
    token_path.write_text(creds.to_json(), encoding="utf-8")
    print(f"[gmail_auth] Wrote token: {token_path}")


if __name__ == "__main__":
    main()