Repository: apache/airflow
File: task-sdk/src/airflow/sdk/crypto.py (Python)

Code before:
        raise ValueError(f"Expected bytes or str, got {type(msg)}")

    def encrypt(self, msg: bytes) -> bytes:
        """Encrypt with Fernet."""
        return msg


class _RealFernet:
    """
    A wrapper around the real Fernet to set is_encrypted to True.

    This class is only used internally to avoid changing the interface of
    the get_fernet function.
    """

    from cryptography.fernet import Fernet, MultiFernet

    is_encrypted = True

    def __init__(self, fernet: MultiFernet):
        self._fernet = fernet

    def decrypt(self, msg: bytes | str, ttl: int | None = None) -> bytes:
        """Decrypt with Fernet."""
        return self._fernet.decrypt(msg, ttl)

    def encrypt(self, msg: bytes) -> bytes:
        """Encrypt with Fernet."""
        return self._fernet.encrypt(msg)

    def rotate(self, msg: bytes | str) -> bytes:
        """Rotate the Fernet key for the given message."""
        return self._fernet.rotate(msg)


@cache
def get_fernet() -> FernetProtocol:
    """
    Deferred load of Fernet key from SDK configuration.

    This function could fail either because Cryptography is not installed
    or because the Fernet key is invalid.

    :return: Fernet object
    :raises: airflow.sdk.exceptions.AirflowException if there's a problem trying to load Fernet
    """
    from cryptography.fernet import Fernet, MultiFernet

    from airflow.sdk.configuration import conf
    from airflow.sdk.exceptions import AirflowException


Review comment (suggestion):
These imports at class body level execute at class definition time (not lazy) and are redundant since `get_fernet()` already imports them. Should be removed.

Code after:
    `airflow db migrate` is run.
    """

    is_encrypted = False

    def decrypt(self, msg: bytes | str, ttl: int | None = None) -> bytes:
        """Decrypt with Fernet."""
        if isinstance(msg, bytes):
            return msg
        if isinstance(msg, str):
            return msg.encode("utf-8")
        raise ValueError(f"Expected bytes or str, got {type(msg)}")

    def encrypt(self, msg: bytes) -> bytes:
        """Encrypt with Fernet."""
        return msg


class _RealFernet:
    """
    A wrapper around the real Fernet to set is_encrypted to True.

    This class is only used internally to avoid changing the interface of
    the get_fernet function.
    """

    is_encrypted = True

    def __init__(self, fernet: MultiFernet):
        self._fernet = fernet

    def decrypt(self, msg: bytes | str, ttl: int | None = None) -> bytes:
        """Decrypt with Fernet."""
        return self._fernet.decrypt(msg, ttl)

    def encrypt(self, msg: bytes) -> bytes:
        """Encrypt with Fernet."""
        return self._fernet.encrypt(msg)

    def rotate(self, msg: bytes | str) -> bytes:
        """Rotate the Fernet key for the given message."""
        return self._fernet.rotate(msg)


@cache
def get_fernet() -> FernetProtocol:
    """
    Deferred load of Fernet key from SDK configuration.

    This function could fail either because Cryptography is not installed
    or because the Fernet key is invalid.