deepcell-cli 0.6.1__py3-none-any.whl

deepcell_cli/commands/auth.py

@@ -0,0 +1,610 @@
+"""Authentication commands: login, register, logout, whoami."""
+
+from __future__ import annotations
+
+import os
+import time
+from urllib.parse import urlparse
+
+import click
+
+from deepcell_cli.config import (
+    clear_credentials,
+    client_headers,
+    get_api_url,
+    get_refresh_token,
+    load_config,
+    load_credentials,
+    save_config,
+    save_credentials,
+    warn_if_insecure,
+)
+from deepcell_cli.errors import APIError
+from deepcell_cli.errors import ConnectionError as CLIConnectionError
+from deepcell_cli.context import Ctx, pass_ctx
+from deepcell_cli.output import echo_error, echo_info, echo_success, echo_warning, output
+
+
+def _api_post(url: str, **kwargs: object) -> "httpx.Response":
+    """POST with user-friendly connection error handling.
+
+    Deliberately not routed through ``DeepCellClient``: these are the
+    pre-credential calls (register, login, the device flow), and that client
+    would try to mint an anonymous session to satisfy them. It must still
+    announce the calling surface, though — ``user_registered``,
+    ``demo_signup_converted`` and ``login_success`` are funnel milestones, and
+    they were the whole login stage reading NULL because this path built its
+    own request.
+    """
+    import httpx
+
+    headers = {**client_headers(), **(kwargs.pop("headers", None) or {})}  # type: ignore[dict-item]
+    try:
+        return httpx.post(url, timeout=30.0, headers=headers, **kwargs)
+    except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError) as exc:
+        base = url.rsplit("/auth/", 1)[0] if "/auth/" in url else url
+        raise CLIConnectionError(base, exc) from exc
+
+
+# ── anonymous-work claim ────────────────────────────────────
+
+
+# Key under which a failed claim's device_id survives in the (account)
+# credentials file, so the next `deepcell login` can retry the claim instead
+# of orphaning the anonymous workspace.
+_PENDING_CLAIM_KEY = "pending_anon_claim"
+
+_CLAIM_RETRY_NOTE = (
+    "Note: your earlier work could not be moved automatically. "
+    "It will be retried on your next `deepcell login`."
+)
+
+# Refresh tokens whose server-side revocation failed after local logout. They
+# live in the private (0600) config rather than credentials.json: keeping them
+# in credentials would let DeepCellClient refresh the supposedly logged-out
+# session on the next command. A later logout retries this list through the
+# unauthenticated RFC-7009-style endpoint.
+_PENDING_REVOCATIONS_KEY = "pending_refresh_revocations"
+
+
+def _snapshot_anon_credentials() -> dict | None:
+    """Capture the anonymous session (if any) before login overwrites it.
+
+    Also resurrects a claim that failed on a previous login: the stashed
+    ``pending_anon_claim`` device_id becomes a token-less snapshot, which
+    forces :func:`_claim_anonymous_work` down its re-mint path.
+    """
+    creds = load_credentials()
+    if creds.get("anonymous") and creds.get("device_id"):
+        return creds
+    pending = creds.get(_PENDING_CLAIM_KEY)
+    if pending:
+        return {"device_id": pending, "expires_at": 0}
+    return None
+
+
+def _stash_pending_claim(device_id: str | None) -> None:
+    """Persist the anon device_id onto the just-saved account credentials.
+
+    The in-memory snapshot is the only copy once login has overwritten the
+    credentials file — without this, a transiently failed claim would orphan
+    the anonymous workspace permanently.
+    """
+    if not device_id:
+        return
+    creds = load_credentials()
+    creds[_PENDING_CLAIM_KEY] = device_id
+    save_credentials(creds)
+
+
+def _register_claim_proof(prior: dict | None) -> tuple[str, str] | None:
+    """``(device_id, live anon access token)`` for an in-place register claim.
+
+    Same two factors ``POST /auth/claim-demo`` requires in its body variant,
+    gathered before login overwrites the credentials file. Returns None when
+    there is nothing to claim or no live token can be obtained — register then
+    proceeds as an ordinary signup and the post-hoc claim path still runs, so a
+    failure here costs the promotion, never the work.
+    """
+    if not prior:
+        return None
+    device_id = prior.get("device_id")
+    if not device_id:
+        return None
+    anon_token = prior.get("access_token")
+    if prior.get("expires_at", 0) < time.time() + 30:
+        anon_token = None
+        try:
+            resp = _api_post(
+                f"{get_api_url()}/demo/session",
+                json={"client": "cli", "device_id": device_id},
+            )
+            if resp.status_code == 200:
+                anon_token = resp.json()["access_token"]
+        except Exception:
+            return None
+    if not anon_token:
+        return None
+    return str(device_id), str(anon_token)
+
+
+def _claim_anonymous_work(prior: dict | None) -> None:
+    """Move a prior anonymous session's work into the just-authenticated account.
+
+    Proof mirrors POST /auth/claim-demo's body variant: the stored
+    ``device_id`` plus a live anon access token — re-minted first when the
+    stored one is stale. Every failure is non-fatal (login already
+    succeeded), but never silent: without a live token the claim is skipped
+    entirely (a stale token would 404 indistinguishably from "nothing to
+    claim"), and any failure stashes the device_id back into the credentials
+    file so the next login retries the claim.
+    """
+    if not prior:
+        return
+    device_id = prior.get("device_id")
+    anon_token = prior.get("access_token")
+
+    if prior.get("expires_at", 0) < time.time() + 30:
+        anon_token = None
+        try:
+            resp = _api_post(
+                f"{get_api_url()}/demo/session",
+                json={"client": "cli", "device_id": device_id},
+            )
+            if resp.status_code == 200:
+                anon_token = resp.json()["access_token"]
+        except Exception:
+            pass
+
+    if not anon_token:
+        _stash_pending_claim(device_id)
+        echo_info(_CLAIM_RETRY_NOTE)
+        return
+
+    from deepcell_cli.client import DeepCellClient
+
+    client = DeepCellClient()
+    try:
+        data = client.post(
+            "/auth/claim-demo",
+            json={"device_id": device_id, "anon_access_token": anon_token},
+        )
+    except APIError as exc:
+        if exc.status_code == 404:
+            # Presented a live token, so 404 genuinely means nothing to
+            # claim (row already merged or expired server-side) — stay quiet.
+            return
+        _stash_pending_claim(device_id)
+        echo_info(_CLAIM_RETRY_NOTE)
+        return
+    except Exception:
+        _stash_pending_claim(device_id)
+        echo_info(_CLAIM_RETRY_NOTE)
+        return
+    finally:
+        client.close()
+
+    moved = int(data.get("workspaces_moved", 0) or 0)
+    echo_success(
+        f"Moved {moved} workspace{'s' if moved != 1 else ''} into your account."
+    )
+
+
+# ── login ───────────────────────────────────────────────────
+
+
+@click.command()
+@click.option("--email", help="Email for direct login (headless/CI).")
+@click.option("--password", help="Password for direct login (headless/CI).")
+@pass_ctx
+def login(ctx: Ctx, email: str | None, password: str | None) -> None:
+    """Authenticate with the DeepCell API.
+
+    Without flags, opens a browser-based authorization flow (like `gh auth login`).
+    No account yet? Sign up right on that page (email or Google) — no separate
+    `register` step needed.
+
+    Use `--email` (password will be prompted securely) for headless environments.
+    """
+    if password and not email:
+        raise click.UsageError("--password requires --email.")
+    prior_anon = _snapshot_anon_credentials()
+    if email:
+        if not password:
+            password = click.prompt("Password", hide_input=True)
+        _login_password(ctx, email, password)
+    else:
+        _login_device_flow(ctx)
+    _claim_anonymous_work(prior_anon)
+
+
+def _login_password(ctx: Ctx, email: str, password: str) -> None:
+    """Direct email/password login."""
+    api_url = get_api_url()
+    warn_if_insecure(api_url)
+
+    resp = _api_post(
+        f"{api_url}/auth/login",
+        json={"email": email, "password": password},
+    )
+    if resp.status_code != 200:
+        try:
+            detail = resp.json().get("detail", resp.text)
+        except Exception:
+            detail = resp.text
+        raise APIError(resp.status_code, str(detail))
+
+    data = resp.json()
+    creds: dict[str, object] = {
+        "access_token": data["access_token"],
+        "expires_at": time.time() + data["expires_in"],
+        "user_id": data["user"]["id"],
+        "email": data["user"]["email"],
+    }
+    if data.get("refresh_token"):
+        creds["refresh_token"] = data["refresh_token"]
+    save_credentials(creds)
+    echo_success(f"Authenticated as {data['user']['email']}")
+    echo_info("  Token saved to ~/.deepcell/credentials.json")
+
+
+def _is_safe_verification_url(url: str, api_url: str) -> bool:
+    """Return True if *url* is safe to open in a browser.
+
+    Rules:
+    - Scheme must be ``https`` (or ``http`` only for localhost).
+    - Host must match the API host.
+    """
+    try:
+        parsed = urlparse(url)
+        api_parsed = urlparse(api_url)
+    except Exception:
+        return False
+
+    if parsed.scheme not in ("http", "https"):
+        return False
+
+    url_host = parsed.hostname or ""
+    api_host = api_parsed.hostname or ""
+
+    if url_host != api_host:
+        return False
+
+    # Allow http only for localhost
+    if parsed.scheme == "http" and url_host not in ("localhost", "127.0.0.1", "::1"):
+        return False
+
+    return True
+
+
+def _login_device_flow(ctx: Ctx) -> None:
+    """Browser-based device code login flow."""
+    api_url = get_api_url()
+    warn_if_insecure(api_url)
+
+    # Step 1: Request device code, carrying this session's anonymous identity
+    # when it has one. If the user signs UP in the browser during this flow,
+    # that lets their new account adopt the row whose workspace already holds
+    # their work — otherwise they land on an empty personal workspace with the
+    # work beside it. Nothing here is required: an unproven or absent identity
+    # simply yields an ordinary device flow.
+    proof = _register_claim_proof(_snapshot_anon_credentials())
+    device_payload: dict[str, object] = {}
+    device_headers: dict[str, str] = {}
+    if proof:
+        device_payload["device_id"] = proof[0]
+        device_headers["Authorization"] = f"Bearer {proof[1]}"
+
+    resp = _api_post(
+        f"{api_url}/auth/device/code",
+        json=device_payload,
+        headers=device_headers or None,
+    )
+    if resp.status_code != 200:
+        try:
+            detail = resp.json().get("detail", resp.text)
+        except Exception:
+            detail = resp.text
+        raise APIError(resp.status_code, str(detail))
+
+    data = resp.json()
+    device_code = data["device_code"]
+    verification_url = data["verification_url"]
+    expires_in = data.get("expires_in", 900)
+    interval = data.get("interval", 5)
+
+    # Step 2: Show URL and try to open browser
+    echo_info("To authenticate, visit this URL in your browser:\n")
+    echo_info(f"  {verification_url}\n")
+
+    if _is_safe_verification_url(verification_url, api_url):
+        try:
+            import webbrowser
+
+            webbrowser.open(verification_url)
+            echo_info("(Browser opened automatically)")
+        except Exception:
+            pass
+    else:
+        echo_info("(URL not opened automatically — copy and paste it into your browser)")
+
+    # Step 3: Poll for authorization
+    deadline = time.time() + expires_in
+    spinner_chars = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"
+    spin_idx = 0
+
+    while time.time() < deadline:
+        click.echo(
+            f"\rWaiting for authorization... {spinner_chars[spin_idx % len(spinner_chars)]}",
+            nl=False,
+            err=True,
+        )
+        spin_idx += 1
+        time.sleep(interval)
+
+        try:
+            poll_resp = _api_post(
+                f"{api_url}/auth/device/token",
+                json={"device_code": device_code},
+            )
+        except CLIConnectionError:
+            continue
+
+        if poll_resp.status_code == 200:
+            poll_data = poll_resp.json()
+            if poll_data.get("status") == "pending":
+                continue
+
+            # Got tokens
+            click.echo("\r" + " " * 50 + "\r", nl=False, err=True)
+            creds: dict[str, object] = {
+                "access_token": poll_data["access_token"],
+                "expires_at": time.time() + poll_data["expires_in"],
+                "user_id": poll_data["user"]["id"],
+                "email": poll_data["user"]["email"],
+            }
+            if poll_data.get("refresh_token"):
+                creds["refresh_token"] = poll_data["refresh_token"]
+            save_credentials(creds)
+            echo_success(f"Authenticated as {poll_data['user']['email']}")
+            echo_info("  Token saved to ~/.deepcell/credentials.json")
+            return
+        elif poll_resp.status_code == 404:
+            click.echo("\r" + " " * 50 + "\r", nl=False, err=True)
+            echo_error("Device code expired or not found.")
+            raise SystemExit(1)
+
+    click.echo("\r" + " " * 50 + "\r", nl=False, err=True)
+    echo_error("Authorization timed out. Please try again.")
+    raise SystemExit(1)
+
+
+# ── verify-email ──────────────────────────────────────────────
+
+
+def _run_verification_flow(ctx: Ctx) -> bool:
+    """Run the email verification flow interactively.
+
+    Returns True on success, False on failure.
+    """
+    echo_info("Sending verification email…")
+    try:
+        data = ctx.client.post("/auth/send-verification")
+    except APIError as exc:
+        echo_error(f"Failed to send verification email: {exc.detail}")
+        return False
+
+    if data.get("message") == "Email is already verified":
+        echo_success("Email is already verified!")
+        return True
+
+    # In dev mode the server returns the token directly
+    token_from_server = data.get("token")
+    if token_from_server:
+        echo_info(f"  [dev mode] Verification token: {token_from_server}")
+
+    token = click.prompt("Enter the verification token from your email")
+
+    try:
+        ctx.client.post("/auth/verify-email", json={"token": token}, auth=False)
+    except APIError as exc:
+        echo_error(f"Verification failed: {exc.detail}")
+        return False
+
+    echo_success("Email verified successfully!")
+    return True
+
+
+@click.command("verify-email")
+@pass_ctx
+def verify_email(ctx: Ctx) -> None:
+    """Verify your email address.
+
+    Sends a verification token to your registered email. In development mode
+    the token is displayed directly. Enter the token when prompted.
+    """
+    if not _run_verification_flow(ctx):
+        raise SystemExit(1)
+
+
+# ── register ─────────────────────────────────────────────────
+
+
+@click.command()
+@click.option("--email", prompt=True, help="Account email.")
+@click.option(
+    "--password",
+    prompt=True,
+    hide_input=True,
+    confirmation_prompt=True,
+    help="Account password (min 8 chars).",
+)
+@click.option("--name", "display_name", prompt="Display name", help="Your display name.")
+@click.option("--locale", default="zh-CN", help="Locale (default: zh-CN).")
+def register(email: str, password: str, display_name: str, locale: str) -> None:
+    """Create a new DeepCell account (headless/CI).
+
+    For interactive use, prefer `deepcell login` — it opens the browser,
+    where you can sign up (email or Google) and authorize the CLI in one flow.
+    """
+    api_url = get_api_url()
+    warn_if_insecure(api_url)
+    prior_anon = _snapshot_anon_credentials()
+
+    # Hand the anonymous identity to /auth/register so it upgrades that row IN
+    # PLACE, which is what makes the scratch workspace become the new account's
+    # personal workspace. Without it register creates an empty personal
+    # workspace first and the scratch one arrives afterwards as an ordinary
+    # second workspace — the work survives, beside a home that is not it.
+    #
+    # Both halves are required server-side, so send both or neither: the
+    # device_id proves which row, the live anon token proves we held its
+    # session. A stale token is re-minted first, exactly as the claim path does.
+    claim_proof = _register_claim_proof(prior_anon)
+    payload: dict[str, object] = {
+        "email": email,
+        "password": password,
+        "display_name": display_name,
+        "locale": locale,
+    }
+    headers: dict[str, str] = {}
+    if claim_proof:
+        payload["device_id"] = claim_proof[0]
+        headers["Authorization"] = f"Bearer {claim_proof[1]}"
+
+    resp = _api_post(
+        f"{api_url}/auth/register",
+        json=payload,
+        headers=headers or None,
+    )
+    if resp.status_code == 409:
+        echo_error("An account with that email already exists.")
+        raise SystemExit(1)
+    if resp.status_code != 201:
+        try:
+            detail = resp.json().get("detail", resp.text)
+        except Exception:
+            detail = resp.text
+        raise APIError(resp.status_code, str(detail))
+
+    data = resp.json()
+    creds: dict[str, object] = {
+        "access_token": data["access_token"],
+        "expires_at": time.time() + data["expires_in"],
+        "user_id": data["user"]["id"],
+        "email": data["user"]["email"],
+    }
+    if data.get("refresh_token"):
+        creds["refresh_token"] = data["refresh_token"]
+    save_credentials(creds)
+    echo_success(f"Account created! Authenticated as {data['user']['email']}")
+    echo_info("  Token saved to ~/.deepcell/credentials.json")
+    if claim_proof:
+        # Register already merged it, in place. Calling claim-demo now would
+        # find no anonymous row for this device_id and — because the stored
+        # token is stale by then — re-mint a BRAND NEW anonymous user via
+        # /demo/session just to claim its empty self. Harmless but it litters,
+        # so skip it when we know the work already moved.
+        echo_success("Moved your workspace and its files into your account.")
+    else:
+        _claim_anonymous_work(prior_anon)
+
+    if not data["user"].get("email_verified", False):
+        # Don't block onboarding on verification — features that need a
+        # verified email will say so; `deepcell verify-email` handles it then.
+        echo_info(
+            "\nSome features require a verified email. "
+            "Run `deepcell verify-email` when you need them."
+        )
+
+
+# ── logout ──────────────────────────────────────────────────
+
+
+@click.command()
+@pass_ctx
+def logout(ctx: Ctx) -> None:
+    """Log out and remove stored credentials."""
+    from deepcell_cli.config import is_anonymous_session
+
+    if is_anonymous_session():
+        # The device_id is the only key to the anonymous demo workspace —
+        # discarding it orphans that work permanently.
+        click.confirm(
+            "This discards your local session and its workspace. Continue?",
+            abort=True,
+        )
+        # No server-side revocation: there is no anon-audience logout route,
+        # and the anon access token is short-lived (~15 min) — discarding the
+        # device_id locally is what actually severs access to the workspace.
+        clear_credentials()
+        echo_success("Session discarded.")
+        return
+    # /auth/logout reads the refresh token from a COOKIE, and the CLI holds no
+    # cookie jar (login posts with a bare httpx call), so it revoked nothing
+    # and returned 204 — `logout` reported success while a copied
+    # credentials.json stayed valid. /auth/revoke-token takes the token in the
+    # body precisely for non-browser callers.
+    config = load_config()
+    pending = config.get(_PENDING_REVOCATIONS_KEY) or []
+    if not isinstance(pending, list):
+        pending = []
+    refresh_token = get_refresh_token()
+    tokens = list(dict.fromkeys(
+        [str(token) for token in pending if token]
+        + ([refresh_token] if refresh_token else [])
+    ))
+
+    failed_tokens: list[str] = []
+    failure_messages: list[str] = []
+    for token in tokens:
+        try:
+            # The endpoint treats possession of the refresh token as
+            # authorization and intentionally needs no Bearer token. Using
+            # ctx.client here made a retry impossible after credentials were
+            # cleared because the client tried to authenticate first.
+            response = _api_post(
+                f"{get_api_url()}/auth/revoke-token",
+                json={"refresh_token": token},
+            )
+            if response.status_code != 200:
+                raise APIError(response.status_code, response.text)
+        except Exception as exc:  # noqa: BLE001 — retained and reported
+            failed_tokens.append(token)
+            failure_messages.append(str(exc))
+    clear_credentials()
+
+    if failed_tokens:
+        config[_PENDING_REVOCATIONS_KEY] = failed_tokens
+        save_config(config)
+        echo_warning(
+            "Local credentials removed, but the server did not confirm "
+            f"revocation ({'; '.join(failure_messages)}) — {len(failed_tokens)} "
+            "refresh token(s) may still be valid. Retry `deepcell logout` once "
+            "you are back online; the pending token(s) were retained privately "
+            "for that retry."
+        )
+        raise click.exceptions.Exit(1)
+    else:
+        if _PENDING_REVOCATIONS_KEY in config:
+            config.pop(_PENDING_REVOCATIONS_KEY, None)
+            save_config(config)
+        echo_success("Logged out. Credentials removed.")
+
+
+# ── whoami ──────────────────────────────────────────────────
+
+
+@click.command()
+@pass_ctx
+def whoami(ctx: Ctx) -> None:
+    """Show the currently authenticated user."""
+    from deepcell_cli.config import is_anonymous_session
+
+    if is_anonymous_session():
+        # /auth/me is full-user only; report the anonymous state locally.
+        click.echo("Not signed in.")
+        click.echo("Run `deepcell login` to sign in.")
+        return
+    data = ctx.client.get("/auth/me")
+    output(data, ctx.fmt)