setup.py

#!/usr/bin/env python3
"""
gmail-agent-shield setup
Creates the AI/Quarantine label and filters from filters.json on the
authenticated Gmail account.

Usage:
    python setup.py           # create label + filters
    python setup.py --backfill  # also apply filters to existing matching mail

Requires credentials.json (OAuth client) in the repo root. See docs/oauth-setup.md.
"""

import argparse
import json
import sys
from pathlib import Path

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

SCOPES = [
    "https://www.googleapis.com/auth/gmail.labels",
    "https://www.googleapis.com/auth/gmail.settings.basic",
    "https://www.googleapis.com/auth/gmail.modify",  # only needed for --backfill
]

REPO_ROOT = Path(__file__).resolve().parent
CREDS_FILE = REPO_ROOT / "credentials.json"
TOKEN_FILE = REPO_ROOT / "token.json"
FILTERS_FILE = REPO_ROOT / "filters.json"


def authenticate():
    creds = None
    if TOKEN_FILE.exists():
        creds = Credentials.from_authorized_user_file(str(TOKEN_FILE), SCOPES)
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            if not CREDS_FILE.exists():
                sys.exit(
                    "credentials.json not found. See docs/oauth-setup.md for how "
                    "to create one in Google Cloud Console."
                )
            flow = InstalledAppFlow.from_client_secrets_file(str(CREDS_FILE), SCOPES)
            creds = flow.run_local_server(port=0)
        TOKEN_FILE.write_text(creds.to_json())
    return build("gmail", "v1", credentials=creds)


def ensure_label(service, label_name):
    existing = service.users().labels().list(userId="me").execute().get("labels", [])
    for lbl in existing:
        if lbl["name"] == label_name:
            print(f"[=] Label already exists: {label_name} ({lbl['id']})")
            return lbl["id"]
    created = service.users().labels().create(
        userId="me",
        body={
            "name": label_name,
            "labelListVisibility": "labelShow",
            "messageListVisibility": "show",
        },
    ).execute()
    print(f"[+] Created label: {label_name} ({created['id']})")
    return created["id"]


def filter_already_exists(service, query):
    existing = service.users().settings().filters().list(userId="me").execute().get("filter", [])
    for f in existing:
        if f.get("criteria", {}).get("query") == query:
            return f["id"]
    return None


def create_filter(service, label_id, spec):
    query = spec["query"]
    if filter_already_exists(service, query):
        print(f"[=] Filter already exists: {spec['name']}")
        return
    action = {"addLabelIds": [label_id], "removeLabelIds": ["INBOX"]}
    if spec.get("remove_important"):
        action["removeLabelIds"].append("IMPORTANT")
    if spec.get("action") == "trash":
        action = {"addLabelIds": ["TRASH"], "removeLabelIds": ["INBOX"]}
    body = {"criteria": {"query": query}, "action": action}
    service.users().settings().filters().create(userId="me", body=body).execute()
    print(f"[+] Created filter: {spec['name']}")


def backfill(service, label_id, filters):
    print("\n--- Backfilling existing matching messages ---")
    for spec in filters:
        query = spec["query"]
        print(f"[*] Searching for matches: {spec['name']}")
        ids = []
        page_token = None
        while True:
            resp = service.users().messages().list(
                userId="me", q=query, pageToken=page_token, maxResults=500
            ).execute()
            ids.extend(m["id"] for m in resp.get("messages", []))
            page_token = resp.get("nextPageToken")
            if not page_token:
                break
        if not ids:
            print(f"    0 matches")
            continue
        print(f"    {len(ids)} matches — applying label")
        for chunk in (ids[i : i + 1000] for i in range(0, len(ids), 1000)):
            body = {"ids": chunk, "addLabelIds": [label_id], "removeLabelIds": ["INBOX"]}
            service.users().messages().batchModify(userId="me", body=body).execute()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--backfill", action="store_true", help="apply filters to existing mail")
    args = parser.parse_args()

    spec = json.loads(FILTERS_FILE.read_text())
    service = authenticate()

    label_id = ensure_label(service, spec["label"])
    for f in spec["filters"]:
        try:
            create_filter(service, label_id, f)
        except HttpError as e:
            print(f"[!] Failed to create filter '{f['name']}': {e}")

    if args.backfill:
        backfill(service, label_id, spec["filters"])

    print("\nDone. Next steps:")
    print("  1. Add '-label:AI/Quarantine' to your agent's Gmail search queries.")
    print("  2. Paste skill-template.md into your agent's system prompt or skill definition.")
    print("  3. Review the AI/Quarantine label weekly for the first month — watch for false positives.")


if __name__ == "__main__":
    main()