feat: add typed auth record variants
This commit is contained in:
+101
-1
@@ -3,7 +3,9 @@ from __future__ import annotations
|
|||||||
import re
|
import re
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Protocol
|
from typing import Annotated, Any, Literal, Protocol
|
||||||
|
|
||||||
|
from pydantic import AnyUrl, BaseModel, Field
|
||||||
|
|
||||||
AUTH_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
|
AUTH_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
|
||||||
|
|
||||||
@@ -49,9 +51,107 @@ class AuthStore(Protocol):
|
|||||||
def load_auth(self, auth_ref: str) -> AuthRecord | None: ...
|
def load_auth(self, auth_ref: str) -> AuthRecord | None: ...
|
||||||
|
|
||||||
|
|
||||||
|
class BearerAuth(BaseModel):
|
||||||
|
kind: Literal["bearer"] = "bearer"
|
||||||
|
access_token: str
|
||||||
|
|
||||||
|
|
||||||
|
class HeaderAuth(BaseModel):
|
||||||
|
kind: Literal["headers"] = "headers"
|
||||||
|
headers: dict[str, str]
|
||||||
|
|
||||||
|
|
||||||
|
class EnvAuth(BaseModel):
|
||||||
|
kind: Literal["env"] = "env"
|
||||||
|
env: dict[str, str]
|
||||||
|
|
||||||
|
|
||||||
|
class OAuthRefreshTokenAuth(BaseModel):
|
||||||
|
kind: Literal["oauth_refresh_token"] = "oauth_refresh_token"
|
||||||
|
client_id: str
|
||||||
|
client_secret: str
|
||||||
|
refresh_token: str
|
||||||
|
token_url: AnyUrl
|
||||||
|
scopes: tuple[str, ...] = ()
|
||||||
|
|
||||||
|
|
||||||
|
class OpaqueAuth(BaseModel):
|
||||||
|
kind: Literal["opaque"] = "opaque"
|
||||||
|
scheme: str
|
||||||
|
payload: dict[str, object] = Field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
AuthVariant = Annotated[
|
||||||
|
BearerAuth | HeaderAuth | EnvAuth | OAuthRefreshTokenAuth | OpaqueAuth,
|
||||||
|
Field(discriminator="kind"),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class StoredAuthRecord(BaseModel):
|
||||||
|
id: str
|
||||||
|
auth: AuthVariant
|
||||||
|
metadata: dict[str, object] = Field(default_factory=dict)
|
||||||
|
|
||||||
|
def model_post_init(self, __context: Any) -> None:
|
||||||
|
validate_auth_id(self.id)
|
||||||
|
|
||||||
|
|
||||||
|
def auth_record_from_compat(
|
||||||
|
*,
|
||||||
|
id: str,
|
||||||
|
scheme: str,
|
||||||
|
payload: Mapping[str, object],
|
||||||
|
metadata: Mapping[str, object] | None = None,
|
||||||
|
) -> StoredAuthRecord:
|
||||||
|
"""Create a StoredAuthRecord from legacy scheme + payload shape."""
|
||||||
|
payload_dict = dict(payload)
|
||||||
|
metadata_dict = dict(metadata or {})
|
||||||
|
match scheme:
|
||||||
|
case "bearer":
|
||||||
|
token = payload_dict.get("token") or payload_dict.get("access_token")
|
||||||
|
if not isinstance(token, str) or not token:
|
||||||
|
raise ValueError("bearer token is required")
|
||||||
|
auth: AuthVariant = BearerAuth(access_token=token)
|
||||||
|
case "headers":
|
||||||
|
raw_headers = payload_dict.get("headers", {})
|
||||||
|
headers = (
|
||||||
|
{
|
||||||
|
str(key): str(value)
|
||||||
|
for key, value in raw_headers.items()
|
||||||
|
if isinstance(key, str) and isinstance(value, str)
|
||||||
|
}
|
||||||
|
if isinstance(raw_headers, dict)
|
||||||
|
else {}
|
||||||
|
)
|
||||||
|
auth = HeaderAuth(headers=headers)
|
||||||
|
case "env":
|
||||||
|
raw_env = payload_dict.get("env", {})
|
||||||
|
env = (
|
||||||
|
{
|
||||||
|
str(key): str(value)
|
||||||
|
for key, value in raw_env.items()
|
||||||
|
if isinstance(key, str) and isinstance(value, str)
|
||||||
|
}
|
||||||
|
if isinstance(raw_env, dict)
|
||||||
|
else {}
|
||||||
|
)
|
||||||
|
auth = EnvAuth(env=env)
|
||||||
|
case _:
|
||||||
|
auth = OpaqueAuth(scheme=scheme, payload=payload_dict)
|
||||||
|
return StoredAuthRecord(id=id, auth=auth, metadata=metadata_dict)
|
||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"AUTH_ID_PATTERN",
|
"AUTH_ID_PATTERN",
|
||||||
"AuthRecord",
|
"AuthRecord",
|
||||||
"AuthStore",
|
"AuthStore",
|
||||||
|
"AuthVariant",
|
||||||
|
"BearerAuth",
|
||||||
|
"EnvAuth",
|
||||||
|
"HeaderAuth",
|
||||||
|
"OAuthRefreshTokenAuth",
|
||||||
|
"OpaqueAuth",
|
||||||
|
"StoredAuthRecord",
|
||||||
|
"auth_record_from_compat",
|
||||||
"validate_auth_id",
|
"validate_auth_id",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -50,3 +50,82 @@ def test_auth_store_protocol_is_read_only_lookup() -> None:
|
|||||||
|
|
||||||
assert store.load_auth("github.work") is record
|
assert store.load_auth("github.work") is record
|
||||||
assert store.load_auth("missing") is None
|
assert store.load_auth("missing") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_stored_auth_record_accepts_bearer_variant() -> None:
|
||||||
|
from wf_api.auth import BearerAuth, StoredAuthRecord
|
||||||
|
|
||||||
|
record = StoredAuthRecord(
|
||||||
|
id="google.drive.personal",
|
||||||
|
auth=BearerAuth(access_token="access-token"),
|
||||||
|
metadata={"provider": "google"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert record.id == "google.drive.personal"
|
||||||
|
assert record.auth.kind == "bearer"
|
||||||
|
assert record.metadata["provider"] == "google"
|
||||||
|
|
||||||
|
|
||||||
|
def test_stored_auth_record_accepts_oauth_refresh_token_variant() -> None:
|
||||||
|
from pydantic import AnyUrl
|
||||||
|
|
||||||
|
from wf_api.auth import OAuthRefreshTokenAuth, StoredAuthRecord
|
||||||
|
|
||||||
|
record = StoredAuthRecord(
|
||||||
|
id="google.drive.personal",
|
||||||
|
auth=OAuthRefreshTokenAuth(
|
||||||
|
client_id="client-id",
|
||||||
|
client_secret="client-secret",
|
||||||
|
refresh_token="refresh-token",
|
||||||
|
token_url=AnyUrl("https://oauth2.googleapis.com/token"),
|
||||||
|
scopes=("https://www.googleapis.com/auth/drive.readonly",),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert record.auth.kind == "oauth_refresh_token"
|
||||||
|
assert str(record.auth.token_url) == "https://oauth2.googleapis.com/token/"
|
||||||
|
assert record.auth.scopes == ("https://www.googleapis.com/auth/drive.readonly",)
|
||||||
|
|
||||||
|
|
||||||
|
def test_auth_record_from_compat_maps_existing_scheme_payload_shape() -> None:
|
||||||
|
from wf_api.auth import BearerAuth, StoredAuthRecord, auth_record_from_compat
|
||||||
|
|
||||||
|
record = auth_record_from_compat(
|
||||||
|
id="demo.default",
|
||||||
|
scheme="bearer",
|
||||||
|
payload={"token": "abc"},
|
||||||
|
metadata={"source": "test"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(record, StoredAuthRecord)
|
||||||
|
assert isinstance(record.auth, BearerAuth)
|
||||||
|
assert record.auth.access_token == "abc"
|
||||||
|
assert record.metadata["source"] == "test"
|
||||||
|
|
||||||
|
|
||||||
|
def test_auth_record_from_compat_preserves_unknown_as_opaque() -> None:
|
||||||
|
from wf_api.auth import OpaqueAuth, StoredAuthRecord, auth_record_from_compat
|
||||||
|
|
||||||
|
record = auth_record_from_compat(
|
||||||
|
id="demo.default",
|
||||||
|
scheme="custom",
|
||||||
|
payload={"x": "y"},
|
||||||
|
metadata={},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(record, StoredAuthRecord)
|
||||||
|
assert isinstance(record.auth, OpaqueAuth)
|
||||||
|
assert record.auth.scheme == "custom"
|
||||||
|
assert record.auth.payload == {"x": "y"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_typed_auth_rejects_missing_bearer_token() -> None:
|
||||||
|
from wf_api.auth import auth_record_from_compat
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="bearer token"):
|
||||||
|
auth_record_from_compat(
|
||||||
|
id="demo.default",
|
||||||
|
scheme="bearer",
|
||||||
|
payload={},
|
||||||
|
metadata={},
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user