ForcePilot/backend/package/yuxi/channel/extensions/googlechat/auth.py

from __future__ import annotations

import json
import logging
import re
import time
import hashlib

from google.oauth2 import service_account
from google.oauth2 import id_token as google_id_token
from google.auth.transport import requests as google_requests

import httpx
import jwt

from yuxi.channel.extensions.googlechat.types import ResolvedGoogleChatAccount

logger = logging.getLogger(__name__)

CHAT_API_SCOPE = "https://www.googleapis.com/auth/chat.bot"
CERTS_URL = "https://www.googleapis.com/service_accounts/v1/metadata/x509/chat@system.gserviceaccount.com"
CERTS_TTL = 600

_STANDARD_CHAT_ISSUER = "chat@system.gserviceaccount.com"
_ADDON_ISSUER_PATTERN = re.compile(r"^service-\d+@gcp-sa-gsuiteaddons\.iam\.gserviceaccount\.com$")

_auth_cache: dict[str, tuple[str, service_account.Credentials]] = {}
_MAX_AUTH_CACHE = 32

_cert_cache: tuple[float, dict[str, str]] | None = None
_http_client: httpx.AsyncClient | None = None


async def _get_http() -> httpx.AsyncClient:
    global _http_client
    if _http_client is None:
        _http_client = httpx.AsyncClient(timeout=httpx.Timeout(30.0))
    return _http_client


def _normalize_private_key(key: str) -> str:
    return key.replace("\\n", "\n")


def _build_auth_key(sa_info: dict, sa_file: str | None) -> str:
    if sa_file:
        return f"file:{sa_file}"
    if sa_info:
        key_hash = hashlib.sha256(
            json.dumps({"client_email": sa_info.get("client_email"), "project_id": sa_info.get("project_id")}).encode()
        ).hexdigest()[:16]
        return f"inline:{key_hash}"
    return "none"


def get_or_create_credentials(
    account_id: str, sa_info: dict, sa_file: str | None
) -> service_account.Credentials:
    key = _build_auth_key(sa_info, sa_file)
    cached = _auth_cache.get(account_id)
    if cached and cached[0] == key:
        return cached[1]

    private_key = _normalize_private_key(sa_info.get("private_key", ""))

    creds = service_account.Credentials.from_service_account_info(
        {
            **sa_info,
            "private_key": private_key,
        },
        scopes=[CHAT_API_SCOPE],
    )

    if len(_auth_cache) >= _MAX_AUTH_CACHE:
        oldest = next(iter(_auth_cache))
        _auth_cache.pop(oldest)

    _auth_cache[account_id] = (key, creds)
    return creds


async def get_access_token(account: ResolvedGoogleChatAccount) -> str | None:
    if not account.service_account_info:
        return None
    creds = get_or_create_credentials(
        account.account_id,
        account.service_account_info,
        account.service_account_file,
    )
    request = google_requests.Request()
    creds.refresh(request)
    return creds.token


def decode_unverified_jwt_payload(token: str) -> dict:
    import base64

    parts = token.split(".")
    if len(parts) < 2:
        return {}
    try:
        padded = parts[1] + "=" * (4 - len(parts[1]) % 4)
        decoded = base64.b64decode(padded.replace("-", "+").replace("_", "/"), validate=False)
        return json.loads(decoded)
    except Exception:
        return {}


async def _google_verify_id_token(token: str, audience: str) -> dict:
    request = google_requests.Request()
    return google_id_token.verify_oauth2_token(token, request, audience)


async def _fetch_chat_certs() -> dict[str, str]:
    global _cert_cache
    now = time.monotonic()
    if _cert_cache and now - _cert_cache[0] < CERTS_TTL:
        return _cert_cache[1]

    http = await _get_http()
    resp = await http.get(CERTS_URL)
    resp.raise_for_status()
    certs = resp.json()
    _cert_cache = (now, certs)
    return certs


async def _verify_signed_jwt_with_certs(token: str, audience: str, certs: dict[str, str]) -> dict:
    header = jwt.get_unverified_header(token)
    kid = header.get("kid", "")

    cert = certs.get(kid)
    if not cert:
        raise jwt.InvalidTokenError(f"No matching cert for kid: {kid}")

    return jwt.decode(token, cert, algorithms=["RS256"], audience=audience)


async def verify_webhook_token(
    token: str, audience_type: str, audience: str, app_principal: str | None
) -> tuple[bool, dict, str | None]:
    payload = decode_unverified_jwt_payload(token)
    iss = payload.get("iss", "")

    if iss == _STANDARD_CHAT_ISSUER:
        try:
            if audience_type == "app-url":
                verified = await _google_verify_id_token(token, audience)
            elif audience_type == "project-number":
                certs = await _fetch_chat_certs()
                verified = await _verify_signed_jwt_with_certs(token, audience, certs)
            else:
                return False, payload, f"unsupported audience_type: {audience_type}"
            return True, verified, None
        except Exception as e:
            return False, payload, f"token verification failed: {e}"

    if _ADDON_ISSUER_PATTERN.match(iss):
        if app_principal and payload.get("sub") != app_principal:
            return False, payload, "token sub does not match appPrincipal"
        try:
            certs = await _fetch_chat_certs()
            verified = await _verify_signed_jwt_with_certs(token, audience, certs)
            return True, verified, None
        except Exception as e:
            return False, payload, f"add-on token verification failed: {e}"

    return False, payload, f"invalid issuer: {iss}"


async def probe_account(account: ResolvedGoogleChatAccount) -> dict:
    token = await get_access_token(account)
    if not token:
        return {"ok": False, "status": "no_token", "error": "Failed to obtain access token"}

    http = await _get_http()
    try:
        resp = await http.get(
            "https://chat.googleapis.com/v1/spaces?pageSize=1",
            headers={"Authorization": f"Bearer {token}"},
            timeout=10.0,
        )
        resp.raise_for_status()
        return {"ok": True, "status": "connected"}
    except httpx.HTTPStatusError as e:
        return {"ok": False, "status": f"http_{e.response.status_code}", "error": str(e)}
    except Exception as e:
        return {"ok": False, "status": "error", "error": str(e)}


async def close_auth_http() -> None:
    global _http_client
    if _http_client:
        await _http_client.aclose()
        _http_client = None