Extract CredentialVaultService (#3669)

This commit is contained in:
Stanislav Novosad
2025-10-09 11:39:01 -06:00
committed by GitHub
parent e5e812ff67
commit 1421bc10c6
6 changed files with 208 additions and 128 deletions

View File

@@ -14,6 +14,8 @@ from skyvern.forge.sdk.cache.factory import CacheFactory
from skyvern.forge.sdk.db.client import AgentDB from skyvern.forge.sdk.db.client import AgentDB
from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
from skyvern.forge.sdk.schemas.organizations import Organization 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.settings_manager import SettingsManager
from skyvern.forge.sdk.trace import TraceManager from skyvern.forge.sdk.trace import TraceManager
from skyvern.forge.sdk.trace.lmnr import LaminarTrace from skyvern.forge.sdk.trace.lmnr import LaminarTrace
@@ -94,6 +96,8 @@ WORKFLOW_CONTEXT_MANAGER = WorkflowContextManager()
WORKFLOW_SERVICE = WorkflowService() WORKFLOW_SERVICE = WorkflowService()
AGENT_FUNCTION = AgentFunction() AGENT_FUNCTION = AgentFunction()
PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=DATABASE) PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=DATABASE)
CREDENTIAL_VAULT_SERVICE: CredentialVaultService = BitwardenCredentialVaultService()
scrape_exclude: ScrapeExcludeFunc | None = None scrape_exclude: ScrapeExcludeFunc | None = None
authentication_function: Callable[[str], Awaitable[Organization]] | None = None authentication_function: Callable[[str], Awaitable[Organization]] | None = None
authenticate_user_function: Callable[[str], Awaitable[str | None]] | None = None authenticate_user_function: Callable[[str], Awaitable[str | None]] | None = None

View File

@@ -155,37 +155,12 @@ async def create_credential(
), ),
current_org: Organization = Depends(org_auth_service.get_current_org), current_org: Organization = Depends(org_auth_service.get_current_org),
) -> CredentialResponse: ) -> CredentialResponse:
org_collection = await app.DATABASE.get_organization_bitwarden_collection(current_org.organization_id) credential = await app.CREDENTIAL_VAULT_SERVICE.create_credential(
organization_id=current_org.organization_id, data=data
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",
) )
# Early resyncing the Bitwarden vault # 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: if data.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse( credential_response = PasswordCredentialResponse(
@@ -209,6 +184,8 @@ async def create_credential(
credential_type=data.credential_type, credential_type=data.credential_type,
name=data.name, name=data.name,
) )
else:
raise HTTPException(status_code=400, detail=f"Unsupported credential type: {data.credential_type}")
@legacy_base_router.delete("/credentials/{credential_id}") @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), current_org: Organization = Depends(org_auth_service.get_current_org),
) -> None: ) -> 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 = await app.DATABASE.get_credential(
credential_id=credential_id, organization_id=current_org.organization_id credential_id=credential_id, organization_id=current_org.organization_id
) )
if not credential: if not credential:
raise HTTPException(status_code=404, detail=f"Credential not found, credential_id={credential_id}") 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 app.CREDENTIAL_VAULT_SERVICE.delete_credential(credential)
await BitwardenService.delete_credential_item(credential.item_id)
return None return None
@@ -283,45 +253,7 @@ async def get_credential(
), ),
current_org: Organization = Depends(org_auth_service.get_current_org), current_org: Organization = Depends(org_auth_service.get_current_org),
) -> CredentialResponse: ) -> CredentialResponse:
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( return await app.CREDENTIAL_VAULT_SERVICE.get_credential(current_org.organization_id, credential_id)
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")
@legacy_base_router.get("/credentials") @legacy_base_router.get("/credentials")
@@ -359,47 +291,7 @@ async def get_credentials(
openapi_extra={"x-fern-sdk-parameter-name": "page_size"}, openapi_extra={"x-fern-sdk-parameter-name": "page_size"},
), ),
) -> list[CredentialResponse]: ) -> list[CredentialResponse]:
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection( return await app.CREDENTIAL_VAULT_SERVICE.get_credentials(current_org.organization_id, page, page_size)
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
@base_router.get( @base_router.get(

View File

@@ -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)

View File

@@ -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."""

View File

@@ -295,15 +295,14 @@ class WorkflowRunContext:
if db_credential is None: if db_credential is None:
raise CredentialParameterNotFoundError(credential_id) raise CredentialParameterNotFoundError(credential_id)
bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id) credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential)
credential = credential_item.credential
credential_item = bitwarden_credential.credential
self.parameters[parameter.key] = parameter self.parameters[parameter.key] = parameter
self.values[parameter.key] = { self.values[parameter.key] = {
"context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)", "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(): for key, value in credential_dict.items():
if value is None: if value is None:
continue continue
@@ -312,12 +311,12 @@ class WorkflowRunContext:
self.secrets[secret_id] = value self.secrets[secret_id] = value
self.values[parameter.key][key] = secret_id 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() random_secret_id = self.generate_random_secret_id()
totp_secret_id = f"{random_secret_id}_totp" totp_secret_id = f"{random_secret_id}_totp"
self.secrets[totp_secret_id] = BitwardenConstants.TOTP self.secrets[totp_secret_id] = BitwardenConstants.TOTP
totp_secret_value = self.totp_secret_value_key(totp_secret_id) 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 self.values[parameter.key]["totp"] = totp_secret_id
except Exception as e: except Exception as e:
LOG.error(f"Failed to get credential from database: {credential_id}. Error: {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: if db_credential is None:
raise CredentialParameterNotFoundError(credential_id) raise CredentialParameterNotFoundError(credential_id)
bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id) credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential)
credential = credential_item.credential
credential_item = bitwarden_credential.credential
self.parameters[parameter.key] = parameter self.parameters[parameter.key] = parameter
self.values[parameter.key] = { self.values[parameter.key] = {
"context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)", "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(): for key, value in credential_dict.items():
random_secret_id = self.generate_random_secret_id() random_secret_id = self.generate_random_secret_id()
secret_id = f"{random_secret_id}_{key}" secret_id = f"{random_secret_id}_{key}"
self.secrets[secret_id] = value self.secrets[secret_id] = value
self.values[parameter.key][key] = secret_id 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() random_secret_id = self.generate_random_secret_id()
totp_secret_id = f"{random_secret_id}_totp" totp_secret_id = f"{random_secret_id}_totp"
self.secrets[totp_secret_id] = BitwardenConstants.TOTP self.secrets[totp_secret_id] = BitwardenConstants.TOTP
totp_secret_value = self.totp_secret_value_key(totp_secret_id) 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 self.values[parameter.key]["totp"] = totp_secret_id
async def register_aws_secret_parameter_value( async def register_aws_secret_parameter_value(