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