exec-aide/src/shared/gmail.py

import logging
import re
from datetime import datetime, timezone

import google.auth.transport.requests
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError

from shared.secrets import get_gmail_oauth, save_gmail_tokens

logger = logging.getLogger(__name__)

_service = None

SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]


def get_authenticated_service():
    global _service
    oauth = get_gmail_oauth()
    creds = Credentials(
        token=oauth.get("access_token"),
        refresh_token=oauth["refresh_token"],
        token_uri="https://oauth2.googleapis.com/token",
        client_id=oauth["client_id"],
        client_secret=oauth["client_secret"],
        scopes=SCOPES,
    )
    if not creds.valid:
        creds.refresh(google.auth.transport.requests.Request())
        save_gmail_tokens(
            creds.token,
            creds.expiry.isoformat() if creds.expiry else "",
        )
    _service = build("gmail", "v1", credentials=creds, cache_discovery=False)
    return _service


def fetch_history(service, history_id):
    message_ids = []
    new_history_id = history_id
    try:
        page_token = None
        while True:
            resp = (
                service.users()
                .history()
                .list(
                    userId="me",
                    startHistoryId=history_id,
                    historyTypes=["messageAdded"],
                    labelId="INBOX",
                    pageToken=page_token,
                )
                .execute()
            )
            new_history_id = resp.get("historyId", new_history_id)
            for record in resp.get("history", []):
                for added in record.get("messagesAdded", []):
                    message_ids.append(added["message"]["id"])
            page_token = resp.get("nextPageToken")
            if not page_token:
                break
    except HttpError as e:
        if e.resp.status == 404:
            logger.warning("History ID expired, signaling re-sync")
            return None, None
        raise
    return message_ids, new_history_id


def fetch_message(service, message_id):
    msg = (
        service.users()
        .messages()
        .get(
            userId="me",
            id=message_id,
            format="metadata",
            metadataHeaders=["From", "To", "Cc", "Subject", "Date"],
        )
        .execute()
    )
    return _parse_message(msg)


def fetch_thread(service, thread_id):
    thread = (
        service.users()
        .threads()
        .get(userId="me", id=thread_id, format="metadata")
        .execute()
    )
    return [_parse_message(m) for m in thread.get("messages", [])]


def initial_sync(service):
    message_ids = []
    page_token = None
    while True:
        resp = (
            service.users()
            .messages()
            .list(
                userId="me",
                q="newer_than:2d",
                labelIds=["INBOX"],
                pageToken=page_token,
            )
            .execute()
        )
        for m in resp.get("messages", []):
            message_ids.append(m["id"])
        page_token = resp.get("nextPageToken")
        if not page_token:
            break

    if not message_ids:
        profile = service.users().getProfile(userId="me").execute()
        return [], str(profile["historyId"])

    first_msg = (
        service.users()
        .messages()
        .get(userId="me", id=message_ids[0], format="minimal")
        .execute()
    )
    return message_ids, str(first_msg["historyId"])


def _parse_message(msg):
    headers = {h["name"]: h["value"] for h in msg.get("payload", {}).get("headers", [])}
    from_raw = headers.get("From", "")
    from_name, from_email = _parse_from(from_raw)
    return {
        "message_id": msg["id"],
        "thread_id": msg["threadId"],
        "label_ids": msg.get("labelIds", []),
        "from_email": from_email,
        "from_name": from_name,
        "to_emails": _parse_address_list(headers.get("To", "")),
        "cc_emails": _parse_address_list(headers.get("Cc", "")),
        "subject": headers.get("Subject", "(no subject)"),
        "snippet": msg.get("snippet", ""),
        "internal_date": int(msg.get("internalDate", 0)),
    }


def _parse_from(raw):
    match = re.match(r"^(.+?)\s*<(.+?)>$", raw)
    if match:
        return match.group(1).strip().strip('"'), match.group(2).lower()
    return "", raw.strip().lower()


def _parse_address_list(raw):
    if not raw:
        return []
    addresses = []
    for part in raw.split(","):
        part = part.strip()
        match = re.match(r".*<(.+?)>", part)
        if match:
            addresses.append(match.group(1).lower())
        elif "@" in part:
            addresses.append(part.lower())
    return addresses