bb-integrations-library 3.0.11__py3-none-any.whl

bb_integrations_lib/secrets/__init__.py

@@ -0,0 +1,4 @@
+from .credential_models import SDCredential, PECredential, RITACredential, IMAPCredential, AnyCredential, \
+    BadSecretException, AbstractCredential, GoogleCredential, MongoDBCredential
+from .credential_models import allowed_onepassword_models, onepassword_category_map
+from .providers import SecretProvider, IntegrationSecretProvider

bb_integrations_lib/secrets/adapters.py

@@ -0,0 +1,98 @@
+from copy import deepcopy
+from typing import TypeVar, Any, Union
+
+from onepasswordconnectsdk.models import Item, Field, ItemUrls
+from pydantic import TypeAdapter
+
+from bb_integrations_lib.secrets import AnyCredential, AbstractCredential, onepassword_category_map
+
+ConcreteOrTaggedCredential = TypeVar("ConcreteOrTaggedCredential", bound=Union[AbstractCredential, AnyCredential])
+
+
+class OPSecretAdapter:
+    """
+    Helper class for translating between 1Password Connect SDK models and our own credential models.
+    Defines a fixed set of transformation rules that support round-trip conversions.
+
+    Note that ``credential_to_opc`` does NOT output a ready-to-upload Item object, because credential models don't
+    contain fields like title or category, but it does provide the minimum set of field, URL, and tag information needed
+    to validate the model. See ``apply_to_opc`` for a reference implementation of updating a 1P item model with a
+    credential model.
+    """
+
+    @staticmethod
+    def credential_to_opc(credential: AbstractCredential) -> Item:
+        """
+        Convert a credential model to a barebones 1P Connet Item model. The item returned by this method is not ready
+        for upload, but can be used as the basis for merging into an existing model. This function exists primarily to
+        outline a standard set of transformations for round-trip model ser/deser to 1P.
+        """
+        url = getattr(credential, "host", None)
+        return Item(
+            fields=[
+                Field(
+                    label=name,
+                    value=getattr(credential, name)
+                )
+                for name, f in type(credential).model_fields.items() if name != "host"
+            ],
+            tags=[
+                f"type/{credential.type_tag}"
+            ],
+            urls=[
+                ItemUrls(
+                    primary=True,
+                    href=url
+                )
+            ]
+        )
+
+    @staticmethod
+    def apply_to_opc(credential: AbstractCredential, item: Item) -> Item:
+        """
+        "Apply" a credential model to an existing 1P item model. Useful if updating a 1P Item with changed secret field
+        values.
+
+        :param credential: The credential model to convert/apply to item.
+        :param item: The 1P item model, returned by the OPC SDK's ``get_item`` or ``create_item``, for example.
+        :return: A copy of item with the relevant pieces of data replaced with converted items from the credential.
+        """
+        new_item = deepcopy(item)
+        converted = OPSecretAdapter.credential_to_opc(credential)
+        new_item.category = onepassword_category_map[type(credential)]
+        new_item.fields = converted.fields
+        new_item.tags = list(set(new_item.tags) | set(converted.tags))
+        new_item.urls = converted.urls
+        return new_item
+
+    @staticmethod
+    def opc_to_credential_dict(item: Item) -> dict[str, Any]:
+        """
+        Convert a 1P Item to a dictionary of field values, with the standard transformations applied.
+
+        :param item: The 1P Item to convert.
+        :return: A dictionary of fields from the credential that can be used to construct credential models.
+        """
+        fields = {field.label.replace(" ", "_"): field.value for field in item.fields}
+        if item.urls:
+            [primary_url] = [x for x in item.urls if x.primary]
+            fields["host"] = primary_url.href
+        return fields
+
+    @staticmethod
+    def opc_to_credential(
+            item: Item,
+            credential_type: type[ConcreteOrTaggedCredential] = AnyCredential
+    ) -> ConcreteOrTaggedCredential:
+        """
+        Helper method to convert a 1P Item model to a credential model, validating it. This will either throw a
+        ValidationError or return an instance of a concrete credential type. Equivalent to validating a model on the
+        output of ``opc_to_credential_dict``.
+
+        :param item: The 1P Item model to convert.
+        :param credential_type: The type of credential to return. Supports and defaults to AnyCredential, which will
+          automatically infer the type from the type_tag discriminator field, if desired.
+        """
+        if type(credential_type) == TypeAdapter:
+            return credential_type.validate_python(OPSecretAdapter.opc_to_credential_dict(item))
+        return credential_type(**OPSecretAdapter.opc_to_credential_dict(item))

bb_integrations_lib/secrets/credential_models.py

@@ -0,0 +1,222 @@
+from abc import ABC
+from typing import Optional, Literal, Annotated, Union, Self
+
+from pydantic import BaseModel, TypeAdapter, Field, field_validator
+
+from bb_integrations_lib.provider.ftp.model import FTPAuthType, FTPType
+
+
+class BadSecretException(Exception):
+    pass
+
+
+"""
+A generic base class for any and all credentials. These must be storable in 1Password, so they need to be
+tagged (set ``type_tag=cls.__name__``) so they can be discriminated from each other in a union (see ``AnyCredential``).
+"""
+class AbstractCredential(ABC, BaseModel):
+    type_tag: str
+
+
+
+class SDCredential(AbstractCredential):
+    type_tag: Literal["SDCredential"] = "SDCredential"
+
+    host: str
+    username: str | None = None
+    password: str | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
+    mongodb_conn_str: str | None = None
+
+
+class PECredential(AbstractCredential):
+    type_tag: Literal["PECredential"] = "PECredential"
+
+    host: str
+    username: str
+    password: str
+    client_id: str
+    client_secret: str
+
+class QTCredential(AbstractCredential):
+    type_tag: Literal["QTCredential"] = "QTCredential"
+    base_url: str
+    qt_id: str
+    carrier_id: str
+    authorization: str
+
+
+class RITACredential(AbstractCredential):
+    type_tag: Literal["RITACredential"] = "RITACredential"
+
+    base_url: str = "https://rita.gravitate.energy/api/"
+    username: str | None = None
+    password: str | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
+    tenant: str
+
+class FTPCredential(AbstractCredential):
+    """
+    Model representing the credentials required to connect to an FTP server.
+
+    Attributes:
+        host (str): Host address of the FTP server.
+        username (str): Username for FTP authentication.
+        password (Optional[str]): Password for FTP authentication, if applicable.
+        passphrase (Optional[str]): Passphrase for RSA authentication, if applicable.
+        auth_type (Optional[FTPAuthType]): Type of authentication. Defaults to `FTPAuthType.basic`.
+        port (int): Port number for FTP connection. Defaults to 22.
+        ftp_type (FTPType): Type of FTP protocol.
+        private_key (Optional[str]): Private key for RSA authentication, if applicable. Defaults to None.
+    """
+    type_tag: Literal["FTPCredential"] = "FTPCredential"
+
+    host: str
+    username: str
+    password: Optional[str] = None
+    passphrase: Optional[str] = None
+    auth_type: Optional[FTPAuthType] = FTPAuthType.basic
+    port: Optional[int] = None
+    ftp_type: FTPType
+    private_key: Optional[str] = None
+
+
+class AWSCredential(AbstractCredential):
+    """
+    Model representing AWS credentials required for accessing an S3 bucket.
+
+    Attributes:
+        bucket_name (str): The name of the S3 bucket.
+        access_key_id (str): AWS access key ID used for authentication.
+        secret_access_key (str): AWS secret access key used for authentication.
+    """
+    type_tag: Literal["AWSCredential"] = "AWSCredential"
+
+    bucket_name: str
+    access_key_id: str
+    secret_access_key: str
+
+
+class GoogleCredential(AbstractCredential):
+    """
+    Model for Google Cloud credentials.
+
+    Attributes:
+        type (str): Credential type.
+        project_id (str): Project ID.
+        private_key_id (str): Private key ID.
+        private_key (str): Private key string.
+        client_email (str): Client email address.
+        client_id (str): Client ID.
+        auth_uri (str): Authentication URI.
+        token_uri (str): Token URI.
+        auth_provider_x509_cert_url (str): URL to the x509 certificate of the auth provider.
+        client_x509_cert_url (str): URL to the x509 certificate of the client.
+        universe_domain (str): Universe domain.
+    """
+    type_tag: Literal["GoogleCredential"] = "GoogleCredential"
+
+    type: str
+    project_id: str
+    private_key_id: str
+    private_key: str
+    client_email: str
+    client_id: str
+    auth_uri: str
+    token_uri: str
+    auth_provider_x509_cert_url: str
+    client_x509_cert_url: str
+    universe_domain: str
+
+    @field_validator("private_key", mode="after")
+    @classmethod
+    def unescape_newlines(cls, v):
+        return v.replace("\\n", "\n")
+
+
+# TODO: These are not compatible with the flat field layout that 1P secrets require
+class IMAPAuthOAuth(BaseModel):
+    """OAuth IMAP authentication flow."""
+    client_id: str
+    client_secret: str
+    refresh_token: str
+
+
+class IMAPAuthSimple(BaseModel):
+    """Simple (password) IMAP authentication flow."""
+    password: str
+
+
+class IMAPCredential(AbstractCredential):
+    """
+    Model representing required to connect to an IMAP server and inbox.
+    """
+    type_tag: Literal["IMAPCredential"] = "IMAPCredential"
+
+    host: str
+    port: int
+    email_address: str
+    auth: IMAPAuthOAuth | IMAPAuthSimple
+
+
+class MongoDBCredential(AbstractCredential):
+    type_tag: Literal["MongoDBCredential"] = "MongoDBCredential"
+
+    mongo_conn_str: str
+    mongo_db_name: str
+
+class PlatformScienceCredential(AbstractCredential):
+    type_tag: Literal["PlatformScienceCredential"] = "PlatformScienceCredential"
+
+    base_url: str
+    client_id: str
+    client_secret: str
+
+class CargasCredential(AbstractCredential):
+    type_tag: Literal["CargasCredential"] = "CargasCredential"
+
+    base_url: str
+    api_key: str
+
+class KeyVuCredential(AbstractCredential):
+    type_tag: Literal["KeyVuCredential"] = "KeyVuCredential"
+
+    credential: str
+
+
+"""Any credential that can be stored and retrieved in a secrets manager."""
+AnyCredential = TypeAdapter(Annotated[Union[
+    SDCredential,
+    PECredential,
+    QTCredential,
+    RITACredential,
+    FTPCredential,
+    AWSCredential,
+    GoogleCredential,
+    IMAPCredential,
+    MongoDBCredential,
+    PlatformScienceCredential,
+    CargasCredential,
+    KeyVuCredential
+], Field(discriminator="type_tag")])
+
+allowed_onepassword_models: dict[str, AbstractCredential] = {
+    k: v["cls"]
+    for k, v in AnyCredential.core_schema["choices"].items()
+}
+onepassword_category_map = {
+    SDCredential: "LOGIN",
+    PECredential: "LOGIN",
+    QTCredential: "API",
+    RITACredential: "LOGIN",
+    FTPCredential: "SERVER",
+    AWSCredential: "API",
+    GoogleCredential: "API",
+    IMAPCredential: "EMAIL",
+    MongoDBCredential: "DATABASE",
+    PlatformScienceCredential: "API",
+    CargasCredential: "API",
+    KeyVuCredential: "API"
+}

bb_integrations_lib/secrets/factory.py

@@ -0,0 +1,85 @@
+import functools
+
+from io import StringIO
+from loguru import logger
+from typing import TypeVar
+
+from gcloud.aio.auth import Token
+from gcloud.aio.storage import Storage
+from pymongo import MongoClient, AsyncMongoClient
+from pymongo.asynchronous.database import AsyncDatabase
+
+from bb_integrations_lib.gravitate.base_api import BaseAPI
+from bb_integrations_lib.gravitate.pe_api import GravitatePEAPI
+from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
+from bb_integrations_lib.gravitate.sd_api import GravitateSDAPI
+from bb_integrations_lib.provider.api.keyvu.client import KeyVuClient
+from bb_integrations_lib.provider.api.quicktrip.client import QTApiClient
+from bb_integrations_lib.provider.ftp.client import FTPIntegrationClient
+from bb_integrations_lib.provider.imap.client import IMAPClient
+from bb_integrations_lib.secrets import SecretProvider, AnyCredential, IMAPCredential
+from bb_integrations_lib.secrets.credential_models import FTPCredential, GoogleCredential, MongoDBCredential, \
+    KeyVuCredential
+
+T = TypeVar("T", bound=BaseAPI)
+
+
+def _log_creation(func):
+    async def log(self, secret_name: str, *args, **kwargs):
+        ret = await func(self, secret_name, *args, **kwargs)
+        if self.log_creations:
+            logger.debug(f"Made {type(ret).__name__} from '{secret_name}'")
+        return ret
+    return log
+
+
+class APIFactory:
+    """Builds API or client instances from named secrets using a SecretProvider."""
+
+    def __init__(self, provider: SecretProvider, log_creations: bool = True):
+        self.provider = provider
+        self.log_creations = log_creations
+
+    async def _make(self, secret_name: str, api: type[T]) -> T:
+        return api.from_credential(await self.provider.get_secret(secret_name, AnyCredential))
+
+    @_log_creation
+    async def sd(self, secret_name: str) -> GravitateSDAPI:
+        return await self._make(secret_name, GravitateSDAPI)
+
+    @_log_creation
+    async def rita(self, secret_name: str) -> GravitateRitaAPI:
+        return await self._make(secret_name, GravitateRitaAPI)
+
+    @_log_creation
+    async def pe(self, secret_name: str) -> GravitatePEAPI:
+        return await self._make(secret_name, GravitatePEAPI)
+
+    @_log_creation
+    async def ftp(self, secret_name: str) -> FTPIntegrationClient:
+        return FTPIntegrationClient(await self.provider.get_secret(secret_name, FTPCredential))
+
+    @_log_creation
+    async def gcloud_storage(self, secret_name: str) -> Storage:
+        secret = await self.provider.get_secret(secret_name, GoogleCredential)
+        as_json = secret.model_dump_json()
+        return Storage(
+            token=Token(
+                service_file=StringIO(as_json),
+                scopes=["https://www.googleapis.com/auth/cloud-platform"]
+            )
+        )
+
+    @_log_creation
+    async def mongo_db(self, secret_name: str) -> AsyncDatabase:
+        secret = await self.provider.get_secret(secret_name, MongoDBCredential)
+        return AsyncMongoClient(secret.mongo_conn_str)[secret.mongo_db_name]
+
+    @_log_creation
+    async def imap(self, secret_name: str) -> IMAPClient:
+        return IMAPClient(await self.provider.get_secret(secret_name, IMAPCredential))
+
+    @_log_creation
+    async def qt(self, secret_name: str) -> QTApiClient:
+        return await self._make(secret_name, QTApiClient)
+

bb_integrations_lib/secrets/providers.py

@@ -0,0 +1,160 @@
+import time
+import urllib.parse
+from os import getenv
+from typing import TypeVar
+
+import onepasswordconnectsdk as opc_sdk
+from loguru import logger
+from onepasswordconnectsdk.client import AsyncClient as OpConnectAsyncClient
+from pydantic.v1 import ValidationError
+
+from bb_integrations_lib.secrets import BadSecretException, AnyCredential
+from bb_integrations_lib.secrets.adapters import OPSecretAdapter
+
+T = TypeVar("T", bound=AnyCredential)
+
+
+class SecretProvider:
+    """
+    Provides secrets from a 1Password Connect vault. Getting an index of vault items is cached for cache_ttl, but
+    individual vault items are retrieved fresh in every call to get_secret.
+    """
+
+    def __init__(self, op_connect_host: str, op_token: str, vault_name: str, cache_ttl: int = 60):
+        self._op_host = op_connect_host
+        self._op_token = op_token
+        self._op_client: OpConnectAsyncClient | None = None
+        self._vault = None
+        self._item_overview = []
+        self._item_overview_refreshed = 0
+        self.vault_name = vault_name
+        self.cache_ttl = cache_ttl
+
+    async def _get_client(self) -> OpConnectAsyncClient:
+        if self._op_client:
+            return self._op_client
+        else:
+            self._op_client = opc_sdk.new_client(
+                url=self._op_host,
+                token=self._op_token,
+                is_async=True,
+            )
+            await self.use_vault(self.vault_name)
+            return self._op_client
+
+    async def _refresh_item_overviews(self) -> None:
+        self._item_overview = await (await self._get_client()).get_items(self._vault.id)
+        self._item_overview_refreshed = time.monotonic()
+
+    async def use_vault(self, vault_name: str):
+        self.vault_name = vault_name
+        self._vault = await self._op_client.get_vault_by_title(self.vault_name)
+        await self._refresh_item_overviews()
+
+    async def get_item_overviews(self) -> list[opc_sdk.models.SummaryItem]:
+        """
+        Return a list of overview items in the vault. Refreshes the list from 1P if the cached version has expired.
+        """
+        if time.monotonic() - self._item_overview_refreshed > self.cache_ttl:
+            await self._refresh_item_overviews()
+        return self._item_overview
+
+    async def get_secret_plain(self, item_name: str) -> dict:
+        """
+        Retrieves a secret from 1Password and returns it as a dict of fields, with the standard transformations applied.
+        See documentation for ``OPSecretAdapter.opc_to_credential_dict`` for details on the transformations.
+
+        :param item_name: The name of the secret to load.
+        :return: A dictionary of field_name:value for all fields in the secret.
+        """
+        c = await self._get_client()
+        item = await c.get_item_by_title(urllib.parse.quote(item_name), self._vault.id)
+
+        return OPSecretAdapter.opc_to_credential_dict(item)
+
+    async def get_secret(self, item_name: str, secret_type: type[T]) -> T:
+        """
+        Retrieves a secret from 1Password and tries to load it as a Pydantic model of type ``T``.
+
+        Model initialization is handled by Pydantic; it's recommended that the model 'ignore' ``extra`` fields
+        (the default) to avoid notes or other extra fields in the 1Password item from causing validation errors. The
+        model is initialized by a call roughly equivalent to ``T({field_name:value,...})`` for all fields present in the
+        1Password secret.
+
+        Note: If there are spaces in the 1Password field names, they will be replaced with underscores.
+
+        :param item_name: The name of the secret to load.
+        :param secret_type: The type of secret to load, or a TypeAdapter of a discriminated union. Must be derived from
+          Pydantic BaseModel.
+        :raises BadSecretException: If the secret loads successfully but cannot be parsed as the expected type.
+        :return: An instance of ``T`` with the fields from the secret loaded.
+        """
+        try:
+            c = await self._get_client()
+            item = await c.get_item_by_title(urllib.parse.quote(item_name), self._vault.id)
+            return OPSecretAdapter.opc_to_credential(item, secret_type)
+        except ValidationError:
+            raise BadSecretException(f"Failed to parse secret {item_name} as type({secret_type})")
+
+
+class IntegrationSecretProvider(SecretProvider):
+    """A SecretProvider with defaults appropriate for integration pipelines."""
+
+    def __init__(
+            self,
+            op_connect_host: str,
+            vault_name: str,
+            op_token_env_var: str = "OP_SERVICE_ACCOUNT_TOKEN",
+            op_token_override: str | None = None,
+    ):
+        self.op_connect_host = op_connect_host
+        self.vault_name = vault_name
+
+        # 3 ways to get the 1P token:
+        # 1. If op_token_override is provided, use that.
+        # 2. If op_token_env_var is set, use that.
+        # 3. If a file named .1ptoken, is next to the script or in the parent dir, use that.
+
+        if op_token_override:
+            op_token = op_token_override
+            logger.warning(
+                f"Initializing IntegrationSecretProvider with provided op_token_override on vault '{vault_name}'"
+            )
+        else:
+            # Try to load the token from the environment variable
+            op_token = getenv(op_token_env_var)
+
+            if op_token:
+                logger.info(
+                    f"Initializing IntegrationSecretProvider with token from env variable '{op_token_env_var}' "
+                    f"on vault '{vault_name}'"
+                )
+            else:
+                op_token = self._find_token_in_file()
+                if op_token:
+                    logger.info(
+                        f"Initializing IntegrationSecretProvider with token from .1ptoken file "
+                        f"on vault '{vault_name}'"
+                    )
+                if not op_token:
+                    raise ValueError(
+                        f"No op_token_override provided, env variable '{op_token_env_var}' not set,and a .1ptoken file "
+                        "is not found; cannot load secrets"
+                    )
+
+        super().__init__(
+            op_connect_host=op_connect_host,
+            op_token=op_token,
+            vault_name=self.vault_name
+        )
+
+    @staticmethod
+    def _find_token_in_file() -> str | None:
+        paths_to_try = ["../.1ptoken", ".1ptoken"]
+        for path in paths_to_try:
+            try:
+                with open(path, "r") as f:
+                    return f.read().strip()
+            except FileNotFoundError:
+                pass
+        return None

bb_integrations_lib/shared/exceptions.py

@@ -0,0 +1,25 @@
+class MappingNotFoundException(Exception):
+    pass
+
+class StepInitializationError(Exception):
+    pass
+
+class FileParsingError(Exception):
+    """Raised when file parsing fails."""
+    pass
+
+class MapperLoadError(Exception):
+    """Raised when loading the RITA mapper fails."""
+    pass
+
+class ConfigNotFoundError(Exception):
+    """Raised when a configuration is not found."""
+    pass
+
+class ConfigValidationError(Exception):
+    """Raised when configuration validation fails."""
+    pass
+
+class StepConfigValidationError(Exception):
+    """Raised when a step configuration validation fails."""
+    pass