Source code for codegen_database.types.encrypted

"""Encrypted-at-rest SQLAlchemy column type for codegen_database.

:class:`EncryptedText` stores a Fernet-encrypted ``TEXT`` column: the
mapped attribute reads and writes plaintext, ciphertext is what hits
the database.  Fernet ciphertext is non-deterministic, so an
encrypted column can't drive a ``WHERE`` clause or an index -- right
for secrets (OAuth tokens, API keys), wrong for anything queried.

The encryption key is supplied either as a string or -- the usual
case -- as a zero-argument callable resolved on every encrypt /
decrypt, so a key read from the environment is picked up lazily and
never captured at import time.  Any sufficiently-random secret
works: the raw key material is run through SHA-256 to derive the
actual Fernet key (the same derivation ``sqlalchemy-utils``'s
``FernetEngine`` uses, so columns written by
``StringEncryptedType(..., engine=FernetEngine)`` decrypt
unchanged).

Requires the ``encrypted`` extra
(``pip install 'codegen-database[encrypted]'``) for ``cryptography``.
"""

from __future__ import annotations

import base64
import hashlib
from typing import TYPE_CHECKING, Any

from sqlalchemy import Text
from sqlalchemy import types as sa_types

if TYPE_CHECKING:
    from collections.abc import Callable

    from cryptography.fernet import Fernet


def _load_fernet() -> type[Fernet]:
    """Import :class:`~cryptography.fernet.Fernet`, or fail honestly.

    ``cryptography`` lives behind the ``encrypted`` extra so the core
    package stays light; the import is deferred to first use because
    ``codegen_database.types`` is imported by the package root.
    """
    try:
        # Deferred on purpose (not a circular import): cryptography is
        # the optional ``encrypted`` extra, and this module is imported
        # by the package root -- a top-level import would make plain
        # ``import codegen_database`` require it.
        from cryptography.fernet import Fernet  # noqa: PLC0415

    except ImportError as exc:  # pragma: no cover -- env-dependent
        msg = (
            "EncryptedText requires the 'cryptography' package -- "
            "install the extra: pip install 'codegen-database[encrypted]'"
        )
        raise ImportError(msg) from exc

    return Fernet




class EncryptedText(sa_types.TypeDecorator[str]):
    """Store a string Fernet-encrypted in a ``TEXT`` column.

    Args:
        key: The encryption key material, or a zero-argument callable
            returning it.  A callable is invoked on every encrypt /
            decrypt, so keys read from the environment resolve
            lazily -- pass e.g. ``lambda: os.environ["TOKEN_KEY"]``.

    Example::

        import os
        from sqlalchemy.orm import Mapped, mapped_column
        from codegen_database.types import EncryptedText

        class Connection(Base):
            ...
            access_token: Mapped[str] = mapped_column(
                EncryptedText(key=lambda: os.environ["TOKEN_KEY"]),
            )

    Reading a value encrypted under a different key raises
    :class:`cryptography.fernet.InvalidToken` -- rotating the key
    makes previously stored values undecryptable, so treat the key
    like a database credential.
    """

    impl = Text
    cache_ok = True

    def __init__(self, key: str | Callable[[], str]) -> None:
        """Store the key material (or its lazy resolver)."""
        self.key = key
        super().__init__()

    def _fernet(self) -> Fernet:
        """Build the Fernet engine from the (lazily resolved) key.

        The raw key is hashed through SHA-256 to derive the 256-bit
        urlsafe-base64 key Fernet wants -- so callers can supply any
        high-entropy string rather than a pre-formatted Fernet key.
        """
        fernet_cls = _load_fernet()
        material = self.key() if callable(self.key) else self.key
        digest = hashlib.sha256(material.encode("utf-8")).digest()

        return fernet_cls(base64.urlsafe_b64encode(digest))



    def process_bind_param(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> str | None:
        """Encrypt *value* on its way into the database."""
        if value is None:
            return None

        return (
            self._fernet()
            .encrypt(str(value).encode("utf-8"))
            .decode(
                "ascii",
            )
        )




    def process_result_value(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> str | None:
        """Decrypt a stored ciphertext back to plaintext."""
        if value is None:
            return None

        return self._fernet().decrypt(value.encode("ascii")).decode("utf-8")