fix: make oauth login flow interactive

This commit is contained in:
lda
2026-06-13 03:57:34 +07:00 Verified
parent f17fe18779
commit 85b2ca8bc7
4 changed files with 90 additions and 27 deletions
+5 -2
View File
@@ -445,10 +445,13 @@ Run OAuth login:
```bash ```bash
wf --config wf.config.json admin auth oauth-login google \ wf --config wf.config.json admin auth oauth-login google \
--id google.drive.personal \ --id google.drive.personal
--authorization-response "<redirected URL>"
``` ```
The command prints the authorization URL, then prompts for the full redirected
callback URL after browser login. For non-interactive use, pass the callback URL
directly with `--authorization-response "<redirected URL>"`.
Refresh tokens are sensitive. The local file auth store is plaintext and Refresh tokens are sensitive. The local file auth store is plaintext and
intended for local/dev use only. intended for local/dev use only.
+19 -4
View File
@@ -149,16 +149,28 @@ async def _login_with_pasted_response(
provider, provider,
client_id: str, client_id: str,
client_secret: str | None, client_secret: str | None,
authorization_response: str, authorization_response: str | None,
) -> OAuthLoginResult: ) -> OAuthLoginResult:
from authlib.integrations.httpx_client import AsyncOAuth2Client from authlib.integrations.httpx_client import AsyncOAuth2Client
def prompt_for_authorization_response(
authorization_url: str,
state: str,
) -> str | None:
if authorization_response is not None:
return None
typer.echo("Open this URL in your browser to authorize access:")
typer.echo(authorization_url)
typer.echo(f"Expected OAuth state: {state}")
return typer.prompt("Paste the full redirected callback URL")
flow = OAuthCodeLoginFlow(client_factory=AsyncOAuth2Client) # type: ignore[arg-type] flow = OAuthCodeLoginFlow(client_factory=AsyncOAuth2Client) # type: ignore[arg-type]
return await flow.login_with_authorization_response( return await flow.login_with_authorization_response(
provider=provider, provider=provider,
client_id=client_id, client_id=client_id,
client_secret=client_secret, client_secret=client_secret,
authorization_response=authorization_response, authorization_response=authorization_response,
authorization_url_callback=prompt_for_authorization_response,
) )
@@ -168,12 +180,15 @@ def oauth_login(
provider_name: Annotated[str, typer.Argument(help="Auth provider profile name.")], provider_name: Annotated[str, typer.Argument(help="Auth provider profile name.")],
auth_ref: Annotated[str, typer.Option("--id", help="Auth record id/ref to save.")], auth_ref: Annotated[str, typer.Option("--id", help="Auth record id/ref to save.")],
authorization_response: Annotated[ authorization_response: Annotated[
str, str | None,
typer.Option( typer.Option(
"--authorization-response", "--authorization-response",
help="Full redirected callback URL after login.", help=(
"Full redirected callback URL after login. If omitted, prints "
"the authorization URL and prompts for the callback URL."
),
), ),
], ] = None,
) -> None: ) -> None:
"""Run an OAuth login flow and save the resulting refresh token as an auth record.""" """Run an OAuth login flow and save the resulting refresh token as an auth record."""
from wf_config import load_workflow_config from wf_config import load_workflow_config
+13 -2
View File
@@ -67,15 +67,26 @@ class OAuthCodeLoginFlow:
provider: OAuthProviderConfig, provider: OAuthProviderConfig,
client_id: str, client_id: str,
client_secret: str | None, client_secret: str | None,
authorization_response: str, authorization_response: str | None,
authorization_url_callback: Callable[[str, str], str | None] | None = None,
) -> OAuthLoginResult: ) -> OAuthLoginResult:
client = self._client_factory( client = self._client_factory(
client_id=client_id, client_id=client_id,
client_secret=client_secret, client_secret=client_secret,
scope=" ".join(provider.scopes), scope=" ".join(provider.scopes),
redirect_uri=provider.redirect_uri,
code_challenge_method="S256", code_challenge_method="S256",
) )
client.create_authorization_url(str(provider.auth_url)) authorization_url, state = client.create_authorization_url(
str(provider.auth_url),
redirect_uri=provider.redirect_uri,
)
if authorization_url_callback is not None:
callback_response = authorization_url_callback(authorization_url, state)
if authorization_response is None:
authorization_response = callback_response
if authorization_response is None:
raise ValueError("OAuth authorization response is required")
token = await client.fetch_token( token = await client.fetch_token(
str(provider.token_url), str(provider.token_url),
authorization_response=authorization_response, authorization_response=authorization_response,
+53 -19
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import json import json
import pytest import pytest
from pydantic import AnyHttpUrl
from typer.testing import CliRunner from typer.testing import CliRunner
from wf_api.auth import OAuthRefreshTokenAuth from wf_api.auth import OAuthRefreshTokenAuth
@@ -11,13 +12,22 @@ from wf_cli.oauth import OAuthCodeLoginFlow, OAuthLoginResult, build_oauth_recor
from wf_config import OAuthProviderConfig from wf_config import OAuthProviderConfig
def test_build_oauth_record_creates_refresh_token_auth() -> None: def _oauth_provider(
provider = OAuthProviderConfig( *,
scopes: tuple[str, ...] = (),
) -> OAuthProviderConfig:
return OAuthProviderConfig(
kind="oauth_authorization_code_pkce", kind="oauth_authorization_code_pkce",
auth_url="https://accounts.google.com/o/oauth2/v2/auth", auth_url=AnyHttpUrl("https://accounts.google.com/o/oauth2/v2/auth"),
token_url="https://oauth2.googleapis.com/token", token_url=AnyHttpUrl("https://oauth2.googleapis.com/token"),
client_id_env="GOOGLE_OAUTH_CLIENT_ID", client_id_env="GOOGLE_OAUTH_CLIENT_ID",
client_secret_env="GOOGLE_OAUTH_CLIENT_SECRET", client_secret_env="GOOGLE_OAUTH_CLIENT_SECRET",
scopes=scopes,
)
def test_build_oauth_record_creates_refresh_token_auth() -> None:
provider = _oauth_provider(
scopes=("https://www.googleapis.com/auth/drive.readonly",), scopes=("https://www.googleapis.com/auth/drive.readonly",),
) )
result = OAuthLoginResult( result = OAuthLoginResult(
@@ -45,12 +55,7 @@ def test_build_oauth_record_creates_refresh_token_auth() -> None:
def test_build_oauth_record_rejects_missing_refresh_token() -> None: def test_build_oauth_record_rejects_missing_refresh_token() -> None:
provider = OAuthProviderConfig( provider = _oauth_provider()
kind="oauth_authorization_code_pkce",
auth_url="https://accounts.google.com/o/oauth2/v2/auth",
token_url="https://oauth2.googleapis.com/token",
client_id_env="GOOGLE_OAUTH_CLIENT_ID",
)
with pytest.raises(ValueError, match="refresh token"): with pytest.raises(ValueError, match="refresh token"):
build_oauth_record( build_oauth_record(
@@ -64,12 +69,15 @@ def test_build_oauth_record_rejects_missing_refresh_token() -> None:
class _FakeOAuthClient: class _FakeOAuthClient:
def __init__(self) -> None: def __init__(self, **kwargs: object) -> None:
self.authorization_url = "https://auth.example/authorize?state=abc" self.authorization_url = "https://auth.example/authorize?state=abc"
self.init_kwargs = kwargs
self.auth_kwargs: dict[str, object] = {}
self.fetch_calls: list[str] = [] self.fetch_calls: list[str] = []
def create_authorization_url(self, auth_url: str, **kwargs: object) -> tuple[str, str]: def create_authorization_url(self, auth_url: str, **kwargs: object) -> tuple[str, str]:
assert auth_url == "https://accounts.google.com/o/oauth2/v2/auth" assert auth_url == "https://accounts.google.com/o/oauth2/v2/auth"
self.auth_kwargs = dict(kwargs)
return self.authorization_url, "state-123" return self.authorization_url, "state-123"
async def fetch_token(self, token_url: str, authorization_response: str) -> dict[str, object]: async def fetch_token(self, token_url: str, authorization_response: str) -> dict[str, object]:
@@ -81,15 +89,17 @@ class _FakeOAuthClient:
async def test_oauth_code_login_flow_uses_injected_client() -> None: async def test_oauth_code_login_flow_uses_injected_client() -> None:
provider = OAuthProviderConfig( provider = _oauth_provider(
kind="oauth_authorization_code_pkce",
auth_url="https://accounts.google.com/o/oauth2/v2/auth",
token_url="https://oauth2.googleapis.com/token",
client_id_env="GOOGLE_OAUTH_CLIENT_ID",
scopes=("https://www.googleapis.com/auth/drive.readonly",), scopes=("https://www.googleapis.com/auth/drive.readonly",),
) )
client = _FakeOAuthClient() clients: list[_FakeOAuthClient] = []
flow = OAuthCodeLoginFlow(client_factory=lambda **kwargs: client)
def client_factory(**kwargs: object) -> _FakeOAuthClient:
client = _FakeOAuthClient(**kwargs)
clients.append(client)
return client
flow = OAuthCodeLoginFlow(client_factory=client_factory)
result = await flow.login_with_authorization_response( result = await flow.login_with_authorization_response(
provider=provider, provider=provider,
@@ -100,11 +110,35 @@ async def test_oauth_code_login_flow_uses_injected_client() -> None:
assert result.refresh_token == "refresh" assert result.refresh_token == "refresh"
assert result.scopes == ("https://www.googleapis.com/auth/drive.readonly",) assert result.scopes == ("https://www.googleapis.com/auth/drive.readonly",)
client = clients[0]
assert client.init_kwargs["redirect_uri"] == provider.redirect_uri
assert client.auth_kwargs["redirect_uri"] == provider.redirect_uri
assert client.fetch_calls == ["http://127.0.0.1/callback?code=abc&state=state-123"]
async def test_oauth_code_login_flow_callback_can_supply_response() -> None:
provider = _oauth_provider(
scopes=("https://www.googleapis.com/auth/drive.readonly",),
)
client = _FakeOAuthClient()
flow = OAuthCodeLoginFlow(client_factory=lambda **kwargs: client)
result = await flow.login_with_authorization_response(
provider=provider,
client_id="client",
client_secret=None,
authorization_response=None,
authorization_url_callback=lambda url, state: (
"http://127.0.0.1/callback?code=abc&state=state-123"
),
)
assert result.refresh_token == "refresh"
assert client.fetch_calls == ["http://127.0.0.1/callback?code=abc&state=state-123"] assert client.fetch_calls == ["http://127.0.0.1/callback?code=abc&state=state-123"]
def test_auth_oauth_login_saves_record_from_provider_profile(monkeypatch, tmp_path) -> None: def test_auth_oauth_login_saves_record_from_provider_profile(monkeypatch, tmp_path) -> None:
saved: list[object] = [] saved: list[dict[str, object]] = []
class _FakeAdmin: class _FakeAdmin:
async def save_auth_record(self, **kwargs: object) -> dict[str, object]: async def save_auth_record(self, **kwargs: object) -> dict[str, object]: