airbyte.secrets.base

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2"""Base classes and methods for working with secrets in PyAirbyte."""
  3
  4from __future__ import annotations
  5
  6import json
  7from abc import ABC, abstractmethod
  8from enum import Enum
  9from typing import TYPE_CHECKING, Any, cast
 10
 11from pydantic_core import CoreSchema, core_schema
 12
 13from airbyte import exceptions as exc
 14
 15
 16if TYPE_CHECKING:
 17    from pathlib import Path
 18
 19    from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, ValidationInfo
 20    from pydantic.json_schema import JsonSchemaValue
 21
 22
 23class SecretSourceEnum(str, Enum):
 24    """Enumeration of secret sources supported by PyAirbyte."""
 25
 26    ENV = "env"
 27    DOTENV = "dotenv"
 28    GOOGLE_COLAB = "google_colab"
 29    GOOGLE_GSM = "google_gsm"  # Not enabled by default
 30
 31    PROMPT = "prompt"
 32
 33
 34class SecretString(str):
 35    """A string that represents a secret.
 36
 37    This class is used to mark a string as a secret. When a secret is printed, it
 38    will be masked to prevent accidental exposure of sensitive information when debugging
 39    or when printing containing objects like dictionaries.
 40
 41    To create a secret string, simply instantiate the class with any string value:
 42
 43        ```python
 44        secret = SecretString("my_secret_password")
 45        ```
 46
 47    """
 48
 49    __slots__ = ()
 50
 51    def __repr__(self) -> str:
 52        """Override the representation of the secret string to return a masked value.
 53
 54        The secret string is always masked with `****` to prevent accidental exposure, unless
 55        explicitly converted to a string. For instance, printing a config dictionary that contains
 56        a secret will automatically mask the secret value instead of printing it in plain text.
 57
 58        However, if you explicitly convert the cast the secret as a string, such as when used
 59        in an f-string, the secret will be exposed. This is the desired behavior to allow
 60        secrets to be used in a controlled manner.
 61        """
 62        return "<SecretString: ****>"
 63
 64    def is_empty(self) -> bool:
 65        """Check if the secret is an empty string."""
 66        return len(self) == 0
 67
 68    def is_json(self) -> bool:
 69        """Check if the secret string is a valid JSON string."""
 70        try:
 71            json.loads(self)
 72        except (json.JSONDecodeError, Exception):
 73            return False
 74
 75        return True
 76
 77    def __bool__(self) -> bool:
 78        """Override the boolean value of the secret string.
 79
 80        Always returns `True` without inspecting contents.
 81        """
 82        return True
 83
 84    def parse_json(self) -> dict:
 85        """Parse the secret string as JSON."""
 86        try:
 87            return json.loads(self)
 88        except json.JSONDecodeError as ex:
 89            raise exc.PyAirbyteInputError(
 90                message="Failed to parse secret as JSON.",
 91                context={
 92                    "Message": ex.msg,
 93                    "Position": ex.pos,
 94                    "SecretString_Length": len(self),  # Debug secret blank or an unexpected format.
 95                },
 96            ) from None
 97
 98    # Pydantic compatibility
 99
100    @classmethod
101    def validate(
102        cls,
103        v: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
104        info: ValidationInfo,
105    ) -> SecretString:
106        """Validate the input value is valid as a secret string."""
107        _ = info  # Unused
108        if not isinstance(v, str):
109            raise exc.PyAirbyteInputError(
110                message="A valid `str` or `SecretString` object is required.",
111            )
112        return cls(v)
113
114    @classmethod
115    def __get_pydantic_core_schema__(  # noqa: PLW3201  # Pydantic dunder
116        cls,
117        source_type: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
118        handler: GetCoreSchemaHandler,
119    ) -> CoreSchema:
120        """Return a modified core schema for the secret string."""
121        return core_schema.with_info_after_validator_function(
122            function=cls.validate, schema=handler(str), field_name=handler.field_name
123        )
124
125    @classmethod
126    def __get_pydantic_json_schema__(  # noqa: PLW3201  # Pydantic dunder method
127        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
128    ) -> JsonSchemaValue:
129        """Return a modified JSON schema for the secret string.
130
131        - `writeOnly=True` is the official way to prevent secrets from being exposed inadvertently.
132        - `Format=password` is a popular and readable convention to indicate the field is sensitive.
133        """
134        _ = _core_schema, handler  # Unused
135        return {
136            "type": "string",
137            "format": "password",
138            "writeOnly": True,
139        }
140
141
142class SecretManager(ABC):
143    """Abstract base class for secret managers.
144
145    Secret managers are used to retrieve secrets from a secret store.
146
147    By registering a secret manager, PyAirbyte can automatically locate and
148    retrieve secrets from the secret store when needed. This allows you to
149    securely store and access sensitive information such as API keys, passwords,
150    and other credentials without hardcoding them in your code.
151
152    To create a custom secret manager, subclass this class and implement the
153    `get_secret` method. By default, the secret manager will be automatically
154    registered as a global secret source, but will not replace any existing
155    secret sources. To customize this behavior, override the `auto_register` and
156    `replace_existing` attributes in your subclass as needed.
157
158    Note: Registered secrets managers always have priority over the default
159    secret sources such as environment variables, dotenv files, and Google Colab
160    secrets. If multiple secret managers are registered, the last one registered
161    will take priority.
162    """
163
164    replace_existing = False
165    as_backup = False
166
167    def __init__(self) -> None:
168        """Instantiate the new secret manager."""
169        if not hasattr(self, "name"):
170            # Default to the class name if no name is provided
171            self.name: str = self.__class__.__name__
172
173    @abstractmethod
174    def get_secret(self, secret_name: str) -> SecretString | None:
175        """Get a named secret from the secret manager.
176
177        This method should be implemented by subclasses to retrieve secrets from
178        the secret store. If the secret is not found, the method should return `None`.
179        """
180        ...
181
182    def __str__(self) -> str:
183        """Return the name of the secret manager."""
184        return self.name
185
186    def __eq__(self, value: object) -> bool:
187        """Check if the secret manager is equal to another secret manager."""
188        if isinstance(value, SecretManager):
189            return self.name == value.name
190
191        if isinstance(value, str):
192            return self.name == value
193
194        if isinstance(value, SecretSourceEnum):
195            return self.name == str(value)
196
197        return super().__eq__(value)
198
199    def __hash__(self) -> int:
200        """Return a hash of the secret manager name.
201
202        This allows the secret manager to be used in sets and dictionaries.
203        """
204        return hash(self.name)
205
206
207class SecretHandle:
208    """A handle for a secret in a secret manager.
209
210    This class is used to store a reference to a secret in a secret manager.
211    The secret is not retrieved until the `get_value()` or `parse_json()` methods are
212    called.
213    """
214
215    def __init__(
216        self,
217        parent: SecretManager,
218        secret_name: str,
219    ) -> None:
220        """Instantiate a new secret handle."""
221        self.parent = parent
222        self.secret_name = secret_name
223
224    def get_value(self) -> SecretString:
225        """Get the secret from the secret manager.
226
227        Subclasses can optionally override this method to provide a more optimized code path.
228        """
229        return cast(SecretString, self.parent.get_secret(self.secret_name))
230
231    def parse_json(self) -> dict:
232        """Parse the secret as JSON.
233
234        This method is a convenience method to parse the secret as JSON without
235        needing to call `get_value()` first. If the secret is not a valid JSON
236        string, a `PyAirbyteInputError` will be raised.
237        """
238        return self.get_value().parse_json()
239
240    def write_to_file(
241        self,
242        file_path: Path,
243        /,
244        *,
245        silent: bool = False,
246    ) -> None:
247        """Write the secret to a file.
248
249        If `silent` is True, the method will not print any output to the console. Otherwise,
250        the method will print a message to the console indicating the file path to which the secret
251        is being written.
252
253        This method is a convenience method that writes the secret to a file as text.
254        """
255        if not silent:
256            print(
257                f"Writing secret '{self.secret_name.split('/')[-1]}' to '{file_path.absolute()!s}'"
258            )
259
260        file_path.write_text(
261            str(self.get_value()),
262            encoding="utf-8",
263        )