From 1421bc10c6e70256d8c97507e4e0d6916065c2d1 Mon Sep 17 00:00:00 2001 From: Stanislav Novosad Date: Thu, 9 Oct 2025 11:39:01 -0600 Subject: [PATCH] Extract CredentialVaultService (#3669) --- skyvern/forge/app.py | 4 + skyvern/forge/sdk/routes/credentials.py | 124 +-------------- .../forge/sdk/services/credential/__init__.py | 0 .../bitwarden_credential_service.py | 150 ++++++++++++++++++ .../credential/credential_vault_service.py | 36 +++++ skyvern/forge/sdk/workflow/context_manager.py | 22 ++- 6 files changed, 208 insertions(+), 128 deletions(-) create mode 100644 skyvern/forge/sdk/services/credential/__init__.py create mode 100644 skyvern/forge/sdk/services/credential/bitwarden_credential_service.py create mode 100644 skyvern/forge/sdk/services/credential/credential_vault_service.py diff --git a/skyvern/forge/app.py b/skyvern/forge/app.py index ae6308a4..27c4f8b8 100644 --- a/skyvern/forge/app.py +++ b/skyvern/forge/app.py @@ -14,6 +14,8 @@ from skyvern.forge.sdk.cache.factory import CacheFactory from skyvern.forge.sdk.db.client import AgentDB from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider from skyvern.forge.sdk.schemas.organizations import Organization +from skyvern.forge.sdk.services.credential.bitwarden_credential_service import BitwardenCredentialVaultService +from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.forge.sdk.trace import TraceManager from skyvern.forge.sdk.trace.lmnr import LaminarTrace @@ -94,6 +96,8 @@ WORKFLOW_CONTEXT_MANAGER = WorkflowContextManager() WORKFLOW_SERVICE = WorkflowService() AGENT_FUNCTION = AgentFunction() PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=DATABASE) +CREDENTIAL_VAULT_SERVICE: CredentialVaultService = BitwardenCredentialVaultService() + scrape_exclude: ScrapeExcludeFunc | None = None authentication_function: Callable[[str], Awaitable[Organization]] | None = None authenticate_user_function: Callable[[str], Awaitable[str | None]] | None = None diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py index 1d29288a..30719349 100644 --- a/skyvern/forge/sdk/routes/credentials.py +++ b/skyvern/forge/sdk/routes/credentials.py @@ -155,37 +155,12 @@ async def create_credential( ), current_org: Organization = Depends(org_auth_service.get_current_org), ) -> CredentialResponse: - org_collection = await app.DATABASE.get_organization_bitwarden_collection(current_org.organization_id) - - if not org_collection: - LOG.info( - "There is no collection for the organization. Creating new collection.", - organization_id=current_org.organization_id, - ) - collection_id = await BitwardenService.create_collection( - name=current_org.organization_id, - ) - org_collection = await app.DATABASE.create_organization_bitwarden_collection( - current_org.organization_id, - collection_id, - ) - - item_id = await BitwardenService.create_credential_item( - collection_id=org_collection.collection_id, - name=data.name, - credential=data.credential, - ) - - credential = await app.DATABASE.create_credential( - organization_id=current_org.organization_id, - item_id=item_id, - name=data.name, - credential_type=data.credential_type, - totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none", + credential = await app.CREDENTIAL_VAULT_SERVICE.create_credential( + organization_id=current_org.organization_id, data=data ) # Early resyncing the Bitwarden vault - background_tasks.add_task(fetch_credential_item_background, item_id) + background_tasks.add_task(fetch_credential_item_background, credential.item_id) if data.credential_type == CredentialType.PASSWORD: credential_response = PasswordCredentialResponse( @@ -209,6 +184,8 @@ async def create_credential( credential_type=data.credential_type, name=data.name, ) + else: + raise HTTPException(status_code=400, detail=f"Unsupported credential type: {data.credential_type}") @legacy_base_router.delete("/credentials/{credential_id}") @@ -238,20 +215,13 @@ async def delete_credential( ), current_org: Organization = Depends(org_auth_service.get_current_org), ) -> None: - organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( - current_org.organization_id - ) - if not organization_bitwarden_collection: - raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.") - credential = await app.DATABASE.get_credential( credential_id=credential_id, organization_id=current_org.organization_id ) if not credential: raise HTTPException(status_code=404, detail=f"Credential not found, credential_id={credential_id}") - await app.DATABASE.delete_credential(credential.credential_id, current_org.organization_id) - await BitwardenService.delete_credential_item(credential.item_id) + await app.CREDENTIAL_VAULT_SERVICE.delete_credential(credential) return None @@ -283,45 +253,7 @@ async def get_credential( ), current_org: Organization = Depends(org_auth_service.get_current_org), ) -> CredentialResponse: - organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( - current_org.organization_id - ) - if not organization_bitwarden_collection: - raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.") - - credential = await app.DATABASE.get_credential( - credential_id=credential_id, organization_id=current_org.organization_id - ) - if not credential: - raise HTTPException(status_code=404, detail="Credential not found") - - credential_item = await BitwardenService.get_credential_item(credential.item_id) - if not credential_item: - raise HTTPException(status_code=404, detail="Credential not found") - - if credential_item.credential_type == CredentialType.PASSWORD: - credential_response = PasswordCredentialResponse( - username=credential_item.credential.username, - totp_type=credential.totp_type, - ) - return CredentialResponse( - credential=credential_response, - credential_id=credential.credential_id, - credential_type=credential_item.credential_type, - name=credential_item.name, - ) - if credential_item.credential_type == CredentialType.CREDIT_CARD: - credential_response = CreditCardCredentialResponse( - last_four=credential_item.credential.card_number[-4:], - brand=credential_item.credential.card_brand, - ) - return CredentialResponse( - credential=credential_response, - credential_id=credential.credential_id, - credential_type=credential_item.credential_type, - name=credential_item.name, - ) - raise HTTPException(status_code=400, detail="Invalid credential type") + return await app.CREDENTIAL_VAULT_SERVICE.get_credential(current_org.organization_id, credential_id) @legacy_base_router.get("/credentials") @@ -359,47 +291,7 @@ async def get_credentials( openapi_extra={"x-fern-sdk-parameter-name": "page_size"}, ), ) -> list[CredentialResponse]: - organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( - current_org.organization_id - ) - if not organization_bitwarden_collection: - return [] - - credentials = await app.DATABASE.get_credentials(current_org.organization_id, page=page, page_size=page_size) - items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id) - - response_items = [] - for credential in credentials: - item = next((item for item in items if item.item_id == credential.item_id), None) - if not item: - continue - if item.credential_type == CredentialType.PASSWORD: - credential_response = PasswordCredentialResponse( - username=item.credential.username, - totp_type=credential.totp_type, - ) - response_items.append( - CredentialResponse( - credential=credential_response, - credential_id=credential.credential_id, - credential_type=item.credential_type, - name=item.name, - ) - ) - elif item.credential_type == CredentialType.CREDIT_CARD: - credential_response = CreditCardCredentialResponse( - last_four=item.credential.card_number[-4:], - brand=item.credential.card_brand, - ) - response_items.append( - CredentialResponse( - credential=credential_response, - credential_id=credential.credential_id, - credential_type=item.credential_type, - name=item.name, - ) - ) - return response_items + return await app.CREDENTIAL_VAULT_SERVICE.get_credentials(current_org.organization_id, page, page_size) @base_router.get( diff --git a/skyvern/forge/sdk/services/credential/__init__.py b/skyvern/forge/sdk/services/credential/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py b/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py new file mode 100644 index 00000000..f41fd975 --- /dev/null +++ b/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py @@ -0,0 +1,150 @@ +import structlog +from fastapi import HTTPException + +from skyvern.forge import app +from skyvern.forge.sdk.schemas.credentials import ( + CreateCredentialRequest, + Credential, + CredentialItem, + CredentialResponse, + CredentialType, + CreditCardCredentialResponse, + PasswordCredentialResponse, +) +from skyvern.forge.sdk.services.bitwarden import BitwardenService +from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService + +LOG = structlog.get_logger() + + +class BitwardenCredentialVaultService(CredentialVaultService): + async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential: + org_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id) + + if not org_collection: + LOG.info( + "There is no collection for the organization. Creating new collection.", + organization_id=organization_id, + ) + collection_id = await BitwardenService.create_collection( + name=organization_id, + ) + org_collection = await app.DATABASE.create_organization_bitwarden_collection( + organization_id, + collection_id, + ) + + item_id = await BitwardenService.create_credential_item( + collection_id=org_collection.collection_id, + name=data.name, + credential=data.credential, + ) + + credential = await app.DATABASE.create_credential( + organization_id=organization_id, + item_id=item_id, + name=data.name, + credential_type=data.credential_type, + totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none", + ) + + return credential + + async def delete_credential( + self, + credential: Credential, + ) -> None: + organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( + credential.organization_id + ) + if not organization_bitwarden_collection: + raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.") + + await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id) + await BitwardenService.delete_credential_item(credential.item_id) + + async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse: + organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id) + if not organization_bitwarden_collection: + raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.") + + credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id) + if not credential: + raise HTTPException(status_code=404, detail="Credential not found") + + credential_item = await BitwardenService.get_credential_item(credential.item_id) + if not credential_item: + raise HTTPException(status_code=404, detail="Credential not found") + + if credential_item.credential_type == CredentialType.PASSWORD: + credential_response = PasswordCredentialResponse( + username=credential_item.credential.username, + totp_type=credential.totp_type, + ) + return CredentialResponse( + credential=credential_response, + credential_id=credential.credential_id, + credential_type=credential_item.credential_type, + name=credential_item.name, + ) + if credential_item.credential_type == CredentialType.CREDIT_CARD: + credential_response = CreditCardCredentialResponse( + last_four=credential_item.credential.card_number[-4:], + brand=credential_item.credential.card_brand, + ) + return CredentialResponse( + credential=credential_response, + credential_id=credential.credential_id, + credential_type=credential_item.credential_type, + name=credential_item.name, + ) + raise HTTPException(status_code=400, detail="Invalid credential type") + + async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]: + organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id) + if not organization_bitwarden_collection: + return [] + + credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size) + items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id) + + response_items = [] + for credential in credentials: + item = next((item for item in items if item.item_id == credential.item_id), None) + if not item: + LOG.warning( + "Credential item not found in vault", + credential_id=credential.credential_id, + item_id=credential.item_id, + ) + continue + if item.credential_type == CredentialType.PASSWORD: + credential_response = PasswordCredentialResponse( + username=item.credential.username, + totp_type=credential.totp_type, + ) + response_items.append( + CredentialResponse( + credential=credential_response, + credential_id=credential.credential_id, + credential_type=item.credential_type, + name=item.name, + ) + ) + elif item.credential_type == CredentialType.CREDIT_CARD: + credential_response = CreditCardCredentialResponse( + last_four=item.credential.card_number[-4:], + brand=item.credential.card_brand, + ) + response_items.append( + CredentialResponse( + credential=credential_response, + credential_id=credential.credential_id, + credential_type=item.credential_type, + name=item.name, + ) + ) + return response_items + + async def get_credential_item(self, db_credential: Credential) -> CredentialItem: + return await BitwardenService.get_credential_item(db_credential.item_id) diff --git a/skyvern/forge/sdk/services/credential/credential_vault_service.py b/skyvern/forge/sdk/services/credential/credential_vault_service.py new file mode 100644 index 00000000..2161a673 --- /dev/null +++ b/skyvern/forge/sdk/services/credential/credential_vault_service.py @@ -0,0 +1,36 @@ +from abc import ABC, abstractmethod + +from skyvern.forge.sdk.schemas.credentials import ( + CreateCredentialRequest, + Credential, + CredentialItem, + CredentialResponse, +) + + +class CredentialVaultService(ABC): + """Abstract interface for credential vault services. + + This interface defines the contract for storing and retrieving credentials + from different vault providers (e.g., Bitwarden, OnePassword, AWS Secrets Manager). + """ + + @abstractmethod + async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential: + """Create a new credential in the vault and database.""" + + @abstractmethod + async def delete_credential(self, credential: Credential) -> None: + """Delete a credential from the vault and database.""" + + @abstractmethod + async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse: + """Retrieve a credential with masked sensitive data.""" + + @abstractmethod + async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]: + """Retrieve all credentials for an organization with pagination.""" + + @abstractmethod + async def get_credential_item(self, db_credential: Credential) -> CredentialItem: + """Retrieve the full credential data from the vault.""" diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 62164338..eadc2e2a 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -295,15 +295,14 @@ class WorkflowRunContext: if db_credential is None: raise CredentialParameterNotFoundError(credential_id) - bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id) - - credential_item = bitwarden_credential.credential + credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential) + credential = credential_item.credential self.parameters[parameter.key] = parameter self.values[parameter.key] = { "context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)", } - credential_dict = credential_item.model_dump() + credential_dict = credential.model_dump() for key, value in credential_dict.items(): if value is None: continue @@ -312,12 +311,12 @@ class WorkflowRunContext: self.secrets[secret_id] = value self.values[parameter.key][key] = secret_id - if isinstance(credential_item, PasswordCredential) and credential_item.totp is not None: + if isinstance(credential, PasswordCredential) and credential.totp is not None: random_secret_id = self.generate_random_secret_id() totp_secret_id = f"{random_secret_id}_totp" self.secrets[totp_secret_id] = BitwardenConstants.TOTP totp_secret_value = self.totp_secret_value_key(totp_secret_id) - self.secrets[totp_secret_value] = parse_totp_secret(credential_item.totp) + self.secrets[totp_secret_value] = parse_totp_secret(credential.totp) self.values[parameter.key]["totp"] = totp_secret_id except Exception as e: LOG.error(f"Failed to get credential from database: {credential_id}. Error: {e}") @@ -345,27 +344,26 @@ class WorkflowRunContext: if db_credential is None: raise CredentialParameterNotFoundError(credential_id) - bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id) - - credential_item = bitwarden_credential.credential + credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential) + credential = credential_item.credential self.parameters[parameter.key] = parameter self.values[parameter.key] = { "context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)", } - credential_dict = credential_item.model_dump() + credential_dict = credential.model_dump() for key, value in credential_dict.items(): random_secret_id = self.generate_random_secret_id() secret_id = f"{random_secret_id}_{key}" self.secrets[secret_id] = value self.values[parameter.key][key] = secret_id - if isinstance(credential_item, PasswordCredential) and credential_item.totp is not None: + if isinstance(credential, PasswordCredential) and credential.totp is not None: random_secret_id = self.generate_random_secret_id() totp_secret_id = f"{random_secret_id}_totp" self.secrets[totp_secret_id] = BitwardenConstants.TOTP totp_secret_value = self.totp_secret_value_key(totp_secret_id) - self.secrets[totp_secret_value] = parse_totp_secret(credential_item.totp) + self.secrets[totp_secret_value] = parse_totp_secret(credential.totp) self.values[parameter.key]["totp"] = totp_secret_id async def register_aws_secret_parameter_value(