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

Code before:

    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

    try:
        fernet_key = conf.get("core", "FERNET_KEY")
        if not fernet_key:
            log.warning("empty cryptography key - values will not be stored encrypted.")
            return _NullFernet()

        fernet = MultiFernet([Fernet(fernet_part.encode("utf-8")) for fernet_part in fernet_key.split(",")])
        return _RealFernet(fernet)
    except (ValueError, TypeError) as value_error:
        raise AirflowException(f"Could not create Fernet object: {value_error}")


Review comment (bug):
`_RealFernet` has `rotate` but `FernetProtocol` doesn't define it. Calling `rotate()` on a `FernetProtocol` type would fail at runtime if it's a `_NullFernet`.

Code after:
        """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.

    :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

    try:
        fernet_key = conf.get("core", "FERNET_KEY")
        if not fernet_key:
            log.warning("empty cryptography key - values will not be stored encrypted.")
            return _NullFernet()