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
wf --config wf.config.json admin auth oauth-login google \
--id google.drive.personal \
--authorization-response "<redirected URL>"
--id google.drive.personal
```
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
intended for local/dev use only.
+19 -4
View File
@@ -149,16 +149,28 @@ async def _login_with_pasted_response(
provider,
client_id: str,
client_secret: str | None,
authorization_response: str,
authorization_response: str | None,
) -> OAuthLoginResult:
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]
return await flow.login_with_authorization_response(
provider=provider,
client_id=client_id,
client_secret=client_secret,
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.")],
auth_ref: Annotated[str, typer.Option("--id", help="Auth record id/ref to save.")],
authorization_response: Annotated[
str,
str | None,
typer.Option(
"--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:
"""Run an OAuth login flow and save the resulting refresh token as an auth record."""
from wf_config import load_workflow_config
+13 -2
View File
@@ -67,15 +67,26 @@ class OAuthCodeLoginFlow:
provider: OAuthProviderConfig,
client_id: str,
client_secret: str | None,
authorization_response: str,
authorization_response: str | None,
authorization_url_callback: Callable[[str, str], str | None] | None = None,
) -> OAuthLoginResult:
client = self._client_factory(
client_id=client_id,
client_secret=client_secret,
scope=" ".join(provider.scopes),
redirect_uri=provider.redirect_uri,
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(
str(provider.token_url),
authorization_response=authorization_response,
+53 -19
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import json
import pytest
from pydantic import AnyHttpUrl
from typer.testing import CliRunner
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
def test_build_oauth_record_creates_refresh_token_auth() -> None:
provider = OAuthProviderConfig(
def _oauth_provider(
*,
scopes: tuple[str, ...] = (),
) -> OAuthProviderConfig:
return OAuthProviderConfig(
kind="oauth_authorization_code_pkce",
auth_url="https://accounts.google.com/o/oauth2/v2/auth",
token_url="https://oauth2.googleapis.com/token",
auth_url=AnyHttpUrl("https://accounts.google.com/o/oauth2/v2/auth"),
token_url=AnyHttpUrl("https://oauth2.googleapis.com/token"),
client_id_env="GOOGLE_OAUTH_CLIENT_ID",
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",),
)
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:
provider = OAuthProviderConfig(
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",
)
provider = _oauth_provider()
with pytest.raises(ValueError, match="refresh token"):
build_oauth_record(
@@ -64,12 +69,15 @@ def test_build_oauth_record_rejects_missing_refresh_token() -> None:
class _FakeOAuthClient:
def __init__(self) -> None:
def __init__(self, **kwargs: object) -> None:
self.authorization_url = "https://auth.example/authorize?state=abc"
self.init_kwargs = kwargs
self.auth_kwargs: dict[str, object] = {}
self.fetch_calls: list[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"
self.auth_kwargs = dict(kwargs)
return self.authorization_url, "state-123"
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:
provider = OAuthProviderConfig(
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",
provider = _oauth_provider(
scopes=("https://www.googleapis.com/auth/drive.readonly",),
)
client = _FakeOAuthClient()
flow = OAuthCodeLoginFlow(client_factory=lambda **kwargs: client)
clients: list[_FakeOAuthClient] = []
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(
provider=provider,
@@ -100,11 +110,35 @@ async def test_oauth_code_login_flow_uses_injected_client() -> None:
assert result.refresh_token == "refresh"
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"]
def test_auth_oauth_login_saves_record_from_provider_profile(monkeypatch, tmp_path) -> None:
saved: list[object] = []
saved: list[dict[str, object]] = []
class _FakeAdmin:
async def save_auth_record(self, **kwargs: object) -> dict[str, object]: