smart-org/apps/api/scripts/agy_login.py

"""Add an agy (Google Antigravity) account through headless OAuth.

The admin console does this over HTTP; this script is the same flow for an
operator sitting at the machine that runs the API — useful because the login
needs a browser redirect to a loopback URL that deliberately never loads.

    # 1) start a login; prints the Google URL to open
    .venv/Scripts/python scripts/agy_login.py start

    # 2) after signing in, paste the browser's full address-bar URL back
    .venv/Scripts/python scripts/agy_login.py finish "http://localhost:51121/oauth-callback?code=...&state=..."

Run ``scripts/agy_login.py status`` to list stored accounts. The pending
challenge lives in ``.agy-login-state.json`` next to this script (gitignored)
for the ten minutes it is valid — never commit it.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path

# The script lives in <repo>/apps/api/scripts/; app/ is one level up.
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from app.core.config import settings  # noqa: E402
from app.core.crypto import encrypt  # noqa: E402
from app.core.database import SessionFactory  # noqa: E402
from app.domain.ai_control.models import (  # noqa: E402
    ModelCredential,
    Provider,
    ProviderType,
)
from app.infrastructure.adapters.agy import (  # noqa: E402
    CHALLENGE_TTL_SECONDS,
    AgyAuthClient,
    AgyLoginChallenge,
    AgyOAuthError,
)
from app.infrastructure.repositories.ai_control import (  # noqa: E402
    ModelCredentialRepository,
    ProviderRepository,
)

STATE_PATH = Path(__file__).resolve().parents[1] / ".agy-login-state.json"
AGY_PROVIDER_NAME = "agy"


def _save_challenge(challenge: AgyLoginChallenge, auth_url: str) -> None:
    STATE_PATH.write_text(
        json.dumps(
            {
                "login_id": challenge.login_id,
                "auth_url": auth_url,
                "redirect_uri": challenge.redirect_uri,
                "verifier": challenge.verifier,
                "state": challenge.state,
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )


def _load_challenge() -> tuple[AgyLoginChallenge, str]:
    if not STATE_PATH.exists():
        raise SystemExit(
            "هیچ لاگین در انتظاری نیست. اول «start» را اجرا کنید.\n"
            "No pending login. Run `start` first."
        )
    data = json.loads(STATE_PATH.read_text(encoding="utf-8"))
    challenge = AgyLoginChallenge(
        login_id=data["login_id"],
        auth_url=data["auth_url"],
        redirect_uri=data["redirect_uri"],
        verifier=data["verifier"],
        state=data["state"],
        created_at=0.0,
    )
    return challenge, data["auth_url"]


async def cmd_start(_: argparse.Namespace) -> None:
    challenge = AgyAuthClient().begin()
    _save_challenge(challenge, challenge.auth_url)
    print(challenge.auth_url)
    print()
    print("این لینک را در مرورگر باز کنید و با اکانت Google وارد شوید.")
    print("Open this URL in a browser and sign in with the Google account.")
    print(
        f"بعد از تأیید، آدرس کامل نوار آدرس مرورگر را کپی کنید "
        f"(با {challenge.redirect_uri} شروع می‌شود) و با «finish <URL>» بدهید."
    )
    print(f"این درخواست {CHALLENGE_TTL_SECONDS // 60} دقیقه معتبر است.")


async def cmd_finish(args: argparse.Namespace) -> None:
    challenge, _auth_url = _load_challenge()
    try:
        result = await AgyAuthClient().exchange(challenge, args.callback_url)
    except AgyOAuthError as exc:
        raise SystemExit(f"لاگین ناموفق بود: {exc}") from exc
    if not result.email:
        raise SystemExit("Google نشانی ایمیل اکانت را برنگرداند؛ لاگین لغو شد.")

    async with SessionFactory() as session:
        providers = ProviderRepository(session)
        credentials = ModelCredentialRepository(session)
        provider = await providers.get_by_name(AGY_PROVIDER_NAME)
        if provider is None:
            provider = await providers.create(
                Provider(name=AGY_PROVIDER_NAME, type=ProviderType.AGY, base_url="")
            )
        envelope = json.dumps(
            {
                "refresh_token": result.refresh_token,
                "project_id": result.project_id,
                "email": result.email,
            }
        )
        existing = await credentials.find_by_email(provider.id, result.email)
        if existing is not None:
            await credentials.replace_material(
                existing.id, encrypted_api_key=encrypt(envelope),
                project_id=result.project_id,
            )
            action = "به‌روزرسانی شد"
        else:
            await credentials.create(
                ModelCredential(
                    provider_id=provider.id,
                    name=result.email,
                    encrypted_api_key=encrypt(envelope),
                    email=result.email,
                    project_id=result.project_id,
                )
            )
            action = "اضافه شد"
        await session.commit()

    STATE_PATH.unlink(missing_ok=True)
    print(f"اکانت {result.email} {action}.")
    print(f"project_id: {result.project_id or '(auto)'}")
    print("توکن refresh به‌صورت رمزنگاری‌شده ذخیره شد و هرگز از API برنمی‌گردد.")
    print(
        "برای استفادهٔ واقعی، AI_DEFAULT_PROVIDER=agy را تنظیم و "
        "«db:seed» را دوباره اجرا کنید."
    )


async def cmd_status(_: argparse.Namespace) -> None:
    async with SessionFactory() as session:
        providers = ProviderRepository(session)
        provider = await providers.get_by_name(AGY_PROVIDER_NAME)
        if provider is None:
            print("provider agy وجود ندارد. یک لاگین شروع کنید.")
            return
        accounts = await ModelCredentialRepository(session).list_for_provider(provider.id)
    if not accounts:
        print("هیچ اکانتی ثبت نشده است.")
        return
    for account in accounts:
        state = account.status if account.is_active else "disabled"
        print(f"- {account.email or account.name} [{state}] last_used={account.last_used_at}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("start", help="شروع لاگین و چاپ لینک گوگل")
    finish = sub.add_parser("finish", help="تکمیل لاگین با URL بازگشت مرورگر")
    finish.add_argument("callback_url", help="آدرس کامل نوار آدرس بعد از لاگین")
    sub.add_parser("status", help="فهرست اکانت‌های ثبت‌شده")
    args = parser.parse_args()

    if settings.is_production and args.command in {"start", "finish"}:
        print(
            "هشدار: AI_DEFAULT_PROVIDER در محیط production است. "
            "ادامه می‌دهید چون لاگین فقط credential می‌سازد."
        )

    asyncio.run(
        {
            "start": cmd_start,
            "finish": cmd_finish,
            "status": cmd_status,
        }[args.command](args)
    )


if __name__ == "__main__":
    main()