fix: secure login flow with copy-paste API key exchange

src/runpod_flash/cli/commands/login.py

@@ -1,6 +1,4 @@
 import asyncio
-import datetime as dt
-from typing import Optional
 
 import typer
 from rich.console import Console
@@ -15,20 +13,8 @@
 
 console = Console()
 
-POLL_INTERVAL_SECONDS = 2.0
-DEFAULT_TIMEOUT_SECONDS = 600.0
 
-
-def _parse_expires_at(value: Optional[str]) -> Optional[dt.datetime]:
-    if not value:
-        return None
-    try:
-        return dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
-    except ValueError:
-        return None
-
-
-async def _login(open_browser: bool, timeout_seconds: float) -> None:
+async def _login(open_browser: bool) -> None:
     async with RunpodGraphQLClient(require_api_key=False) as client:
         request = await client.create_flash_auth_request()
         request_id = request.get("id")
@@ -45,46 +31,40 @@ async def _login(open_browser: bool, timeout_seconds: float) -> None:
         if open_browser:
             typer.launch(auth_url)
 
-        expires_at = _parse_expires_at(request.get("expiresAt"))
-        deadline = dt.datetime.now(dt.timezone.utc) + dt.timedelta(
-            seconds=timeout_seconds
-        )
-        if expires_at and expires_at < deadline:
-            deadline = expires_at
+    api_key = console.input(
+        "Paste the API key shown after authorization: ", password=True
+    ).strip()
 
-        with console.status("[dim]Waiting for authorization...[/dim]"):
-            while True:
-                status_payload = await client.get_flash_auth_request_status(request_id)
-                status = status_payload.get("status")
-                api_key = status_payload.get("apiKey")
-
-                if api_key and status in {"APPROVED", "CONSUMED"}:
-                    check_and_migrate_legacy_credentials()
-                    path = save_api_key(api_key)
-                    console.print(
-                        f"[green]Logged in.[/green] Credentials saved to [dim]{path}[/dim]"
-                    )
-                    console.print()
-                    return
-
-                if status in {"DENIED", "EXPIRED", "CONSUMED"}:
-                    raise RuntimeError(f"login failed: {status.lower()}")
+    if not api_key:
+        raise RuntimeError(
+            "no api key provided. copy the key shown in the browser after approval "
+            "and paste it here."
+        )
+    if not api_key.startswith("rpa_"):
+        raise RuntimeError("invalid api key format. Runpod API keys start with rpa_.")
 
-                if dt.datetime.now(dt.timezone.utc) >= deadline:
-                    raise RuntimeError("login timed out")
+    async with RunpodGraphQLClient(api_key=api_key) as client:
+        if not await client.verify_api_key():
+            raise RuntimeError(
+                "api key could not be verified. Check the pasted key and try again."
+            )
 
-                await asyncio.sleep(POLL_INTERVAL_SECONDS)
+    check_and_migrate_legacy_credentials()
+    path = save_api_key(api_key)
+    console.print(f"[green]Logged in.[/green] Credentials saved to [dim]{path}[/dim]")
+    console.print()
 
 
 def login_command(
     no_open: bool = typer.Option(False, "--no-open", help="do not open the browser"),
-    timeout: float = typer.Option(
-        DEFAULT_TIMEOUT_SECONDS, "--timeout", help="max wait time in seconds"
-    ),
 ):
     """Authenticate and save a Runpod API key for flash."""
     try:
-        asyncio.run(_login(open_browser=not no_open, timeout_seconds=timeout))
+        asyncio.run(_login(open_browser=not no_open))
     except RuntimeError as exc:
         print_error(console, str(exc))
         raise typer.Exit(code=1)
+    except KeyboardInterrupt:
+        console.print()
+        print_error(console, "login cancelled")
+        raise typer.Exit(code=130)

src/runpod_flash/core/api/runpod.py

@@ -860,19 +860,26 @@ async def create_flash_auth_request(self) -> Dict[str, Any]:
         result = await self._execute_graphql(mutation)
         return result.get("createFlashAuthRequest", {})
 
-    async def get_flash_auth_request_status(self, request_id: str) -> Dict[str, Any]:
+    async def verify_api_key(self) -> bool:
         query = """
-        query flashAuthRequestStatus($flashAuthRequestId: String!) {
-            flashAuthRequestStatus(flashAuthRequestId: $flashAuthRequestId) {
-                id
-                status
-                expiresAt
-                apiKey
+        query verifyApiKey {
+            myself {
+                __typename
             }
         }
         """
-        result = await self._execute_graphql(query, {"flashAuthRequestId": request_id})
-        return result.get("flashAuthRequestStatus", {})
+        try:
+            result = await self._execute_graphql(query)
+        except (
+            _GraphQLHTTPStatusError,
+            _GraphQLNetworkError,
+            _GraphQLErrorResponse,
+            asyncio.TimeoutError,
+        ) as exc:
+            log.debug("Failed to verify Runpod API key: %s", exc)
+            return False
+
+        return isinstance(result.get("myself"), dict)
 
     async def close(self):
         """Close the HTTP session."""

tests/unit/test_login.py

@@ -5,24 +5,6 @@
 
 import pytest
 
-from runpod_flash.cli.commands.login import _parse_expires_at
-
-
-class TestParseExpiresAt:
-    def test_iso_format(self):
-        result = _parse_expires_at("2026-03-01T12:00:00Z")
-        assert result is not None
-        assert result.year == 2026
-
-    def test_none_input(self):
-        assert _parse_expires_at(None) is None
-
-    def test_empty_string(self):
-        assert _parse_expires_at("") is None
-
-    def test_invalid_string(self):
-        assert _parse_expires_at("not-a-date") is None
-
 
 class TestGraphQLClientNoKeyForLogin:
     """Login mutations must not send stored credentials."""
@@ -61,16 +43,14 @@ def test_require_api_key_true_loads_key(self):
             assert client.api_key == "loaded-key"
 
 
-def _make_mock_client(**status_return):
+def _make_mock_client():
     """Build an AsyncMock that works as an async context manager."""
     client = AsyncMock()
     client.create_flash_auth_request.return_value = {
         "id": "req-123",
         "expiresAt": None,
     }
-    client.get_flash_auth_request_status.return_value = status_return
-    # _login uses `async with RunpodGraphQLClient(...) as client:`,
-    # so __aenter__ must return the same mock instance
+    client.verify_api_key.return_value = True
     client.__aenter__.return_value = client
     return client
 
@@ -85,42 +65,73 @@ def _get_login_fn():
 
 
 class TestLoginFlow:
-    async def test_login_denied(self):
-        mock_client = _make_mock_client(status="DENIED", apiKey=None)
+    async def test_login_saves_pasted_key(self, isolate_credentials_file):
+        mock_client = _make_mock_client()
         _login = _get_login_fn()
 
-        with patch(
-            "runpod_flash.cli.commands.login.RunpodGraphQLClient",
-            return_value=mock_client,
+        with (
+            patch(
+                "runpod_flash.cli.commands.login.RunpodGraphQLClient",
+                return_value=mock_client,
+            ),
+            patch("runpod_flash.cli.commands.login.console") as mock_console,
         ):
-            with pytest.raises(RuntimeError, match="login failed: denied"):
-                await _login(open_browser=False, timeout_seconds=5)
+            mock_console.input.return_value = "rpa_pasted-api-key"
+            await _login(open_browser=False)
+            assert isolate_credentials_file.exists()
+            assert "rpa_pasted-api-key" in isolate_credentials_file.read_text()
+            mock_console.input.assert_called_once_with(
+                "Paste the API key shown after authorization: ", password=True
+            )
 
-    async def test_login_approved_saves_key(self, isolate_credentials_file):
-        mock_client = _make_mock_client(status="APPROVED", apiKey="fresh-api-key")
+    async def test_login_empty_key_raises(self):
+        mock_client = _make_mock_client()
         _login = _get_login_fn()
 
-        with patch(
-            "runpod_flash.cli.commands.login.RunpodGraphQLClient",
-            return_value=mock_client,
+        with (
+            patch(
+                "runpod_flash.cli.commands.login.RunpodGraphQLClient",
+                return_value=mock_client,
+            ),
+            patch("runpod_flash.cli.commands.login.console") as mock_console,
         ):
-            await _login(open_browser=False, timeout_seconds=5)
-            assert isolate_credentials_file.exists()
-            assert "fresh-api-key" in isolate_credentials_file.read_text()
+            mock_console.input.return_value = "  "
+            with pytest.raises(RuntimeError, match="copy the key shown"):
+                await _login(open_browser=False)
 
-    async def test_login_expired(self):
-        mock_client = _make_mock_client(status="EXPIRED", apiKey=None)
+    async def test_login_invalid_key_format_raises(self):
+        mock_client = _make_mock_client()
         _login = _get_login_fn()
 
-        with patch(
-            "runpod_flash.cli.commands.login.RunpodGraphQLClient",
-            return_value=mock_client,
+        with (
+            patch(
+                "runpod_flash.cli.commands.login.RunpodGraphQLClient",
+                return_value=mock_client,
+            ),
+            patch("runpod_flash.cli.commands.login.console") as mock_console,
+        ):
+            mock_console.input.return_value = "not-a-runpod-key"
+            with pytest.raises(RuntimeError, match="Runpod API keys start with rpa_"):
+                await _login(open_browser=False)
+
+    async def test_login_unverified_key_raises(self):
+        mock_client = _make_mock_client()
+        mock_client.verify_api_key.return_value = False
+        _login = _get_login_fn()
+
+        with (
+            patch(
+                "runpod_flash.cli.commands.login.RunpodGraphQLClient",
+                return_value=mock_client,
+            ),
+            patch("runpod_flash.cli.commands.login.console") as mock_console,
         ):
-            with pytest.raises(RuntimeError, match="login failed: expired"):
-                await _login(open_browser=False, timeout_seconds=5)
+            mock_console.input.return_value = "rpa_bad-key"
+            with pytest.raises(RuntimeError, match="api key could not be verified"):
+                await _login(open_browser=False)
 
     async def test_no_request_id_raises(self):
-        mock_client = _make_mock_client(status="APPROVED", apiKey="key")
+        mock_client = _make_mock_client()
         mock_client.create_flash_auth_request.return_value = {}
         _login = _get_login_fn()
 
@@ -129,4 +140,4 @@ async def test_no_request_id_raises(self):
             return_value=mock_client,
         ):
             with pytest.raises(RuntimeError, match="auth request failed"):
-                await _login(open_browser=False, timeout_seconds=5)
+                await _login(open_browser=False)