Source code for codegen_database.types.encrypted

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]:
    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]):
    impl = Text
    cache_ok = True

    def __init__(self, key: str | Callable[[], str]) -> None:
        self.key = key
        super().__init__()

    def _fernet(self) -> Fernet:
        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:
        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:
        if value is None:
            return None

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