"""Persisted, encrypted OAuth connections for the API-access case.
The connected-service case (QuickBooks, and any "call the provider's
API on the user's behalf later" integration) keeps the provider's
tokens around. This module supplies the storage, following the same
codegen-database-flavoured idiom as :mod:`fsh_lib.files` / :mod:`fsh_lib.links`
-- the consumer owns the table, we own the columns -- plus the
refresh-on-read helper.
Access and refresh tokens are long-lived credentials, so they are
**encrypted at rest, transparently at the column layer**, via
:class:`codegen_database.types.EncryptedText` (Fernet). The ORM
attributes are plain strings in Python; ciphertext is what hits the
database. The Fernet key is read -- lazily, at the moment a value
is written or read -- from an environment variable (default
``OAUTH_TOKEN_KEY``), so no key is ever embedded in generated source,
matching :mod:`fsh_lib.auth`'s ``secret_env`` convention. Owning
the type in codegen-database also makes Alembic autogeneration
render these columns correctly (its ``render_item`` knows the
ciphertext is plain ``TEXT``).
Two helpers operate on a connection row:
* :func:`store_tokens` writes a :class:`~fsh_lib.oauth.client.TokenResponse`
onto a row (computing ``expires_at``), after a code exchange or a
refresh.
* :func:`ensure_fresh` returns a valid access token, refreshing in
place first when the stored one has expired.
Both mutate the row in place and leave the commit to the caller --
the same "caller commits" contract as :func:`fsh_lib.links.resolve`.
Because providers may rotate refresh tokens, concurrent refreshes of
one row must be serialized by the caller -- see the warning on
:func:`ensure_fresh`.
Requires the ``oauth`` extra (``pip install 'fsh-lib[oauth]'``)
for ``codegen-database`` (the ``EncryptedText`` column type) +
``cryptography``.
"""
from __future__ import annotations
import datetime
import os
import secrets
from typing import TYPE_CHECKING
from codegen_database.types import EncryptedText
from sqlalchemy import DateTime, Text
from sqlalchemy.orm import Mapped, mapped_column
from fsh_lib.oauth.client import OAuthError
if TYPE_CHECKING:
import uuid
from collections.abc import Callable
from fsh_lib.oauth.client import OAuthClient, TokenResponse
#: Name of the env var holding the token-encryption key. A consumer
#: that needs a different one redeclares the token columns with
#: ``EncryptedText(key=...)`` directly.
DEFAULT_KEY_ENV = "OAUTH_TOKEN_KEY"
#: Default grace window -- refresh a token this long *before* its
#: stated expiry so an access token doesn't lapse mid-request.
DEFAULT_REFRESH_LEEWAY = datetime.timedelta(seconds=60)
[docs]
class OAuthTokenKeyError(RuntimeError):
"""The token-encryption key env var is unset when a token is touched."""
[docs]
def generate_key() -> str:
"""Return a fresh, high-entropy value for the token-encryption env var.
The Fernet engine hashes whatever string the env var holds into
its actual key, so any sufficiently-random secret works -- this
returns a 256-bit URL-safe one. Generate it once per
environment and store it as ``OAUTH_TOKEN_KEY`` (see
:data:`DEFAULT_KEY_ENV`); rotating it makes every already-stored
token undecryptable, so treat it like a database credential.
"""
return secrets.token_urlsafe(32)
def _token_key_resolver(key_env: str) -> Callable[[], str]:
"""Return a callable that reads the encryption key from *key_env*.
The callable is handed to
:class:`~codegen_database.types.EncryptedText`, which invokes it
at bind/result time -- so the key is resolved lazily, never
captured at import.
Raises (when invoked):
OAuthTokenKeyError: *key_env* is unset or empty.
"""
def _resolve() -> str:
key = os.environ.get(key_env)
if not key:
msg = (
f"{key_env} environment variable is required to "
f"encrypt/decrypt stored OAuth tokens"
)
raise OAuthTokenKeyError(msg)
return key
return _resolve
[docs]
class OAuthConnectionMixin:
"""codegen-database mixin: the columns of a stored OAuth connection.
Subclass on a codegen-database-mapped model alongside a PK plugin
(the plugin owns ``id``), the same way as
:class:`fsh_lib.links.ShortLinkMixin`:
.. code-block:: python
from fsh_lib.oauth.connection import OAuthConnectionMixin
from codegen_database.factory import CodegenDatabaseSimple
from codegen_database.plugins.pk import UUIDV7PKPlugin
class OAuthConnection(Base, OAuthConnectionMixin):
__tablename__ = "oauth_connections"
__factory__ = CodegenDatabaseSimple
__plugins__ = [UUIDV7PKPlugin()]
# consumers typically add their own FK to the owning user
Like the file / short-link mixins it deliberately doesn't declare
``id`` (the consumer's PK plugin owns it) or ``created_at``
(``CodegenDatabaseSimple`` auto-adds the timestamp plugin). The
user<->connection link is intentionally left to the consumer (a
``user_id`` FK, a unique ``(user_id, provider)`` constraint,
etc.) since that shape is application-specific.
:attr:`access_token` and :attr:`refresh_token` are encrypted at
rest (:class:`~codegen_database.types.EncryptedText`, keyed from
the :data:`DEFAULT_KEY_ENV` env var); the attributes themselves
are plaintext.
"""
if TYPE_CHECKING:
# Type-only -- the real column comes from the consumer's
# codegen-database PK plugin (see ShortLinkMixin for the rationale).
id: Mapped[uuid.UUID]
provider: Mapped[str] = mapped_column(Text)
"""Identifier of the provider this connection is with, e.g.
``"google"`` / ``"quickbooks"`` -- the consumer's own label,
used to pick the right :class:`~fsh_lib.oauth.client.OAuthClient`
when refreshing."""
account_id: Mapped[str | None] = mapped_column(Text, nullable=True)
"""The provider-side account the tokens are scoped to -- the OIDC
``sub``, or QuickBooks' ``realmId`` (which arrives as a callback
query param, not in the token response). Nullable for providers
that expose no such id."""
access_token: Mapped[str] = mapped_column(
EncryptedText(key=_token_key_resolver(DEFAULT_KEY_ENV)),
)
"""The current access token (encrypted at rest)."""
refresh_token: Mapped[str | None] = mapped_column(
EncryptedText(key=_token_key_resolver(DEFAULT_KEY_ENV)),
nullable=True,
)
"""The refresh token (encrypted at rest), or ``NULL`` when the
provider didn't issue one."""
token_type: Mapped[str] = mapped_column(Text, default="Bearer")
"""The token type, normally ``"Bearer"``."""
scope: Mapped[str | None] = mapped_column(Text, nullable=True)
"""Space-delimited scopes actually granted, or ``NULL``."""
expires_at: Mapped[datetime.datetime | None] = mapped_column(
DateTime(timezone=True),
nullable=True,
)
"""Absolute expiry of :attr:`access_token`, or ``NULL`` when the
provider didn't state a lifetime."""
def _utcnow() -> datetime.datetime:
"""Return the current aware UTC time."""
return datetime.datetime.now(tz=datetime.UTC)
[docs]
def store_tokens(
connection: OAuthConnectionMixin,
tokens: TokenResponse,
*,
now: datetime.datetime | None = None,
) -> None:
"""Write *tokens* onto *connection* in place (encrypted on flush).
Use after a code exchange or a refresh. A ``None``
``refresh_token`` or ``scope`` in *tokens* is treated as
"unchanged" -- the existing value is kept, since refresh
responses routinely omit both (RFC 6749: an omitted scope means
the granted scope is identical). The caller commits.
Args:
connection: The row to update (mixes in
:class:`OAuthConnectionMixin`).
tokens: The freshly obtained
:class:`~fsh_lib.oauth.client.TokenResponse`.
now: Reference time for computing ``expires_at``; defaults to
the current UTC time.
"""
moment = now or _utcnow()
connection.access_token = tokens.access_token
connection.token_type = tokens.token_type
connection.expires_at = tokens.expires_at(moment)
if tokens.scope is not None:
connection.scope = tokens.scope
if tokens.refresh_token is not None:
connection.refresh_token = tokens.refresh_token
[docs]
async def ensure_fresh(
connection: OAuthConnectionMixin,
*,
client: OAuthClient,
now: datetime.datetime | None = None,
leeway: datetime.timedelta = DEFAULT_REFRESH_LEEWAY,
) -> str:
"""Return a valid access token, refreshing *connection* if needed.
If the stored access token is still valid (its ``expires_at`` is
more than *leeway* in the future, or unknown), it is returned
as-is. Otherwise the stored refresh token is spent via
:meth:`~fsh_lib.oauth.client.OAuthClient.refresh`, the new tokens
are written back onto *connection* in place (re-encrypted on
flush), and the new access token is returned. The caller commits
so the rotated tokens persist.
.. warning::
Serialize refreshes per connection row. Some providers
(Intuit/QuickBooks among them) *rotate* the refresh token on
every refresh; two concurrent ``ensure_fresh`` calls would
both spend the same stored token, and the loser can persist a
stale one -- bricking the connection until the user
re-consents. Lock the row before calling (e.g. load it with
``SELECT ... FOR UPDATE`` --
``select(Model).where(...).with_for_update()`` -- or take a
pg advisory lock keyed on the row id) and commit promptly so
the rotated tokens land before the lock is released.
Args:
connection: The stored connection (mixes in
:class:`OAuthConnectionMixin`).
client: An :class:`~fsh_lib.oauth.client.OAuthClient` for the
same provider the connection belongs to.
now: Reference time for the expiry check; defaults to the
current UTC time.
leeway: Refresh this far ahead of the stated expiry so a
token can't lapse mid-request.
Returns:
A non-expired access token.
Raises:
OAuthError: The token has expired and there is no refresh
token to renew it (or the refresh itself failed).
"""
moment = now or _utcnow()
expires_at = connection.expires_at
# DateTime(timezone=True) is a no-op on backends without a
# timestamptz type (SQLite in tests, most commonly), which hand
# back naive datetimes -- treat those as UTC rather than letting
# the comparison raise TypeError.
if expires_at is not None and expires_at.tzinfo is None:
expires_at = expires_at.replace(tzinfo=datetime.UTC)
if expires_at is None or moment + leeway < expires_at:
return connection.access_token
if connection.refresh_token is None:
msg = "access token expired and no refresh token is stored"
raise OAuthError(msg)
tokens = await client.refresh(connection.refresh_token)
store_tokens(connection, tokens, now=moment)
return connection.access_token