rose-ash/market/scrape/auth.py

"""OAuth device flow authentication for the market scraper.

Same flow as the artdag CLI client:
1. Request device code from account server
2. User approves in browser
3. Poll for access token
4. Save token to ~/.artdag/token.json (shared with artdag CLI)
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import httpx

TOKEN_DIR = Path.home() / ".artdag"
TOKEN_FILE = TOKEN_DIR / "token.json"

_DEFAULT_ACCOUNT_SERVER = "https://account.rose-ash.com"


def _get_account_server() -> str:
    return os.getenv("ARTDAG_ACCOUNT", _DEFAULT_ACCOUNT_SERVER)


def load_token() -> dict:
    """Load saved token from ~/.artdag/token.json."""
    if TOKEN_FILE.exists():
        try:
            with open(TOKEN_FILE) as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            return {}
    return {}


def save_token(token_data: dict):
    """Save token to ~/.artdag/token.json (mode 0600)."""
    TOKEN_DIR.mkdir(parents=True, exist_ok=True)
    with open(TOKEN_FILE, "w") as f:
        json.dump(token_data, f, indent=2)
    TOKEN_FILE.chmod(0o600)


def get_access_token(require: bool = True) -> str | None:
    """Return the saved access token, or None.

    When *require* is True and no token is found, triggers interactive
    device-flow login (same as ``artdag login``).
    """
    token = load_token().get("access_token")
    if not token and require:
        print("No saved token found — starting device-flow login...")
        token_data = login()
        token = token_data.get("access_token")
    return token


def auth_headers() -> dict[str, str]:
    """Return Authorization header dict for HTTP requests."""
    token = get_access_token(require=True)
    return {"Authorization": f"Bearer {token}"}


def login() -> dict:
    """Interactive device-flow login (blocking, prints to stdout).

    Returns the token data dict on success. Exits on failure.
    """
    account = _get_account_server()

    # 1. Request device code
    try:
        with httpx.Client(timeout=10) as client:
            resp = client.post(
                f"{account}/auth/device/authorize",
                json={"client_id": "artdag"},
            )
            resp.raise_for_status()
            data = resp.json()
    except httpx.HTTPError as e:
        print(f"Login failed: {e}", file=sys.stderr)
        sys.exit(1)

    device_code = data["device_code"]
    user_code = data["user_code"]
    verification_uri = data["verification_uri"]
    expires_in = data.get("expires_in", 900)
    interval = data.get("interval", 5)

    print("To sign in, open this URL in your browser:")
    print(f"  {verification_uri}")
    print(f"  and enter code:  {user_code}")
    print()

    # Try to open browser
    try:
        import webbrowser
        webbrowser.open(verification_uri)
    except Exception:
        pass

    # 2. Poll for approval
    print("Waiting for authorization", end="", flush=True)
    deadline = time.time() + expires_in

    with httpx.Client(timeout=10) as client:
        while time.time() < deadline:
            time.sleep(interval)
            print(".", end="", flush=True)

            try:
                resp = client.post(
                    f"{account}/auth/device/token",
                    json={"device_code": device_code, "client_id": "artdag"},
                )
                data = resp.json()
            except httpx.HTTPError:
                continue

            error = data.get("error")
            if error == "authorization_pending":
                continue
            elif error == "expired_token":
                print("\nCode expired. Please try again.", file=sys.stderr)
                sys.exit(1)
            elif error == "access_denied":
                print("\nAuthorization denied.", file=sys.stderr)
                sys.exit(1)
            elif error:
                print(f"\nLogin failed: {error}", file=sys.stderr)
                sys.exit(1)

            # Success
            token_data = {
                "access_token": data["access_token"],
                "username": data.get("username", ""),
                "display_name": data.get("display_name", ""),
            }
            save_token(token_data)
            print(f"\nLogged in as {token_data['username'] or token_data['display_name']}")
            return token_data

    print("\nTimed out waiting for authorization.", file=sys.stderr)
    sys.exit(1)