diff --git a/alembic/versions/2024_10_02_2212-6c90d565076b_add_bitwarden_details_to_organizations.py b/alembic/versions/2024_10_02_2212-6c90d565076b_add_bitwarden_details_to_organizations.py new file mode 100644 index 00000000..ae607f60 --- /dev/null +++ b/alembic/versions/2024_10_02_2212-6c90d565076b_add_bitwarden_details_to_organizations.py @@ -0,0 +1,33 @@ +"""Add bitwarden details to organizations + +Revision ID: 6c90d565076b +Revises: c5848cc524b1 +Create Date: 2024-10-02 22:12:34.959165+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "6c90d565076b" +down_revision: Union[str, None] = "c5848cc524b1" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("organizations", sa.Column("bw_organization_id", sa.String(), nullable=True)) + op.add_column("organizations", sa.Column("bw_collection_ids", sa.JSON(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("organizations", "bw_collection_ids") + op.drop_column("organizations", "bw_organization_id") + # ### end Alembic commands ### diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index 58e08286..df5a9c85 100644 --- a/skyvern/exceptions.py +++ b/skyvern/exceptions.py @@ -278,6 +278,15 @@ class BitwardenSyncError(BitwardenBaseError): super().__init__(f"Error syncing Bitwarden: {message}") +class BitwardenAccessDeniedError(BitwardenBaseError): + def __init__(self) -> None: + super().__init__( + "Current organization does not have access to the specified Bitwarden collection. \ + Contact Skyvern support to enable access. This is a security layer on top of Bitwarden, \ + Skyvern team needs to let your Skyvern account access the Bitwarden collection." + ) + + class UnknownElementTreeFormat(SkyvernException): def __init__(self, fmt: str) -> None: super().__init__(f"Unknown element tree format {fmt}") diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 04c4d771..e4ac168c 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -109,6 +109,8 @@ class OrganizationModel(Base): max_steps_per_run = Column(Integer, nullable=True) max_retries_per_step = Column(Integer, nullable=True) domain = Column(String, nullable=True, index=True) + bw_organization_id = Column(String, nullable=True, default=None) + bw_collection_ids = Column(JSON, nullable=True, default=None) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column( DateTime, diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index ce662df2..00383194 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -111,6 +111,8 @@ def convert_to_organization(org_model: OrganizationModel) -> Organization: max_steps_per_run=org_model.max_steps_per_run, max_retries_per_step=org_model.max_retries_per_step, domain=org_model.domain, + bw_organization_id=org_model.bw_organization_id, + bw_collection_ids=org_model.bw_collection_ids, created_at=org_model.created_at, modified_at=org_model.modified_at, ) diff --git a/skyvern/forge/sdk/executor/async_executor.py b/skyvern/forge/sdk/executor/async_executor.py index 805a3956..84b51af8 100644 --- a/skyvern/forge/sdk/executor/async_executor.py +++ b/skyvern/forge/sdk/executor/async_executor.py @@ -99,8 +99,14 @@ class BackgroundTaskExecutor(AsyncExecutor): "Executing workflow using background task executor", workflow_run_id=workflow_run_id, ) + + organization = await app.DATABASE.get_organization(organization_id) + if organization is None: + raise OrganizationNotFound(organization_id) + background_tasks.add_task( app.WORKFLOW_SERVICE.execute_workflow, workflow_run_id=workflow_run_id, api_key=api_key, + organization=organization, ) diff --git a/skyvern/forge/sdk/models.py b/skyvern/forge/sdk/models.py index b4912e73..5c28886a 100644 --- a/skyvern/forge/sdk/models.py +++ b/skyvern/forge/sdk/models.py @@ -131,6 +131,8 @@ class Organization(BaseModel): max_steps_per_run: int | None = None max_retries_per_step: int | None = None domain: str | None = None + bw_organization_id: str | None = None + bw_collection_ids: list[str] | None = None created_at: datetime modified_at: datetime diff --git a/skyvern/forge/sdk/services/bitwarden.py b/skyvern/forge/sdk/services/bitwarden.py index 9573faeb..884cbbc5 100644 --- a/skyvern/forge/sdk/services/bitwarden.py +++ b/skyvern/forge/sdk/services/bitwarden.py @@ -10,6 +10,7 @@ import tldextract from skyvern.config import settings from skyvern.exceptions import ( + BitwardenAccessDeniedError, BitwardenListItemsError, BitwardenLoginError, BitwardenLogoutError, @@ -29,6 +30,9 @@ def is_valid_email(email: str | None) -> bool: class BitwardenConstants(StrEnum): + BW_ORGANIZATION_ID = "BW_ORGANIZATION_ID" + BW_COLLECTION_IDS = "BW_COLLECTION_IDS" + CLIENT_ID = "BW_CLIENT_ID" CLIENT_SECRET = "BW_CLIENT_SECRET" MASTER_PASSWORD = "BW_MASTER_PASSWORD" @@ -79,6 +83,8 @@ class BitwardenService: client_id: str, client_secret: str, master_password: str, + bw_organization_id: str | None, + bw_collection_ids: list[str] | None, url: str, collection_id: str | None = None, remaining_retries: int = settings.BITWARDEN_MAX_RETRIES, @@ -94,6 +100,8 @@ class BitwardenService: client_id=client_id, client_secret=client_secret, master_password=master_password, + bw_organization_id=bw_organization_id, + bw_collection_ids=bw_collection_ids, url=url, collection_id=collection_id, ) @@ -109,6 +117,8 @@ class BitwardenService: client_id=client_id, client_secret=client_secret, master_password=master_password, + bw_organization_id=bw_organization_id, + bw_collection_ids=bw_collection_ids, url=url, collection_id=collection_id, remaining_retries=remaining_retries, @@ -122,12 +132,16 @@ class BitwardenService: client_id: str, client_secret: str, master_password: str, + bw_organization_id: str | None, + bw_collection_ids: list[str] | None, url: str, collection_id: str | None = None, ) -> dict[str, str]: """ Get the secret value from the Bitwarden CLI. """ + if not bw_organization_id and bw_collection_ids and collection_id not in bw_collection_ids: + raise BitwardenAccessDeniedError() try: BitwardenService.login(client_id, client_secret) BitwardenService.sync() @@ -144,7 +158,13 @@ class BitwardenService: "--session", session_key, ] - if collection_id: + if bw_organization_id: + LOG.info( + "Organization ID is provided, filtering items by organization ID", + bw_organization_id=bw_organization_id, + ) + list_command.extend(["--organizationid", bw_organization_id]) + elif collection_id: LOG.info("Collection ID is provided, filtering items by collection ID", collection_id=collection_id) list_command.extend(["--collectionid", collection_id]) items_result = BitwardenService.run_command(list_command) @@ -158,11 +178,26 @@ class BitwardenService: except json.JSONDecodeError: raise BitwardenListItemsError("Failed to parse items JSON. Output: " + items_result.stdout) + # Since Bitwarden can't AND multiple filters, we only use organization id in the list command + # but we still need to filter the items by collection id here + if bw_organization_id and collection_id: + filtered_items = [] + for item in items: + if "collectionIds" in item and collection_id in item["collectionIds"]: + filtered_items.append(item) + items = filtered_items + if not items: collection_id_str = f" in collection with ID: {collection_id}" if collection_id else "" raise BitwardenListItemsError(f"No items found in Bitwarden for URL: {url}{collection_id_str}") + # TODO (kerem): To make this more robust, we need to store the item id of the totp login item + # and use it here to get the TOTP code for that specific item totp_command = ["bw", "get", "totp", url, "--session", session_key] + if bw_organization_id: + # We need to add this filter because the TOTP command fails if there are multiple results + # For now, we require that the bitwarden organization id has only one totp login item for the domain + totp_command.extend(["--organizationid", bw_organization_id]) totp_result = BitwardenService.run_command(totp_command) if totp_result.stderr and "Event post failed" not in totp_result.stderr: @@ -208,6 +243,8 @@ class BitwardenService: client_id: str, client_secret: str, master_password: str, + bw_organization_id: str | None, + bw_collection_ids: list[str] | None, collection_id: str, identity_key: str, identity_fields: list[str], @@ -224,6 +261,8 @@ class BitwardenService: client_id=client_id, client_secret=client_secret, master_password=master_password, + bw_organization_id=bw_organization_id, + bw_collection_ids=bw_collection_ids, collection_id=collection_id, identity_key=identity_key, identity_fields=identity_fields, @@ -240,6 +279,8 @@ class BitwardenService: client_id=client_id, client_secret=client_secret, master_password=master_password, + bw_organization_id=bw_organization_id, + bw_collection_ids=bw_collection_ids, collection_id=collection_id, identity_key=identity_key, identity_fields=identity_fields, @@ -257,10 +298,14 @@ class BitwardenService: collection_id: str, identity_key: str, identity_fields: list[str], + bw_organization_id: str | None, + bw_collection_ids: list[str] | None, ) -> dict[str, str]: """ Get the sensitive information from the Bitwarden CLI. """ + if not bw_organization_id and bw_collection_ids and collection_id not in bw_collection_ids: + raise BitwardenAccessDeniedError() try: BitwardenService.login(client_id, client_secret) BitwardenService.sync() @@ -278,6 +323,8 @@ class BitwardenService: "--collectionid", collection_id, ] + if bw_organization_id: + list_command.extend(["--organizationid", bw_organization_id]) items_result = BitwardenService.run_command(list_command) # Parse the items and extract sensitive information diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index ded0367e..05e5f99d 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -5,6 +5,7 @@ import structlog from skyvern.exceptions import BitwardenBaseError, WorkflowRunContextNotInitialized from skyvern.forge.sdk.api.aws import AsyncAWSClient +from skyvern.forge.sdk.models import Organization from skyvern.forge.sdk.services.bitwarden import BitwardenConstants, BitwardenService from skyvern.forge.sdk.workflow.exceptions import OutputParameterKeyCollisionError from skyvern.forge.sdk.workflow.models.parameter import ( @@ -106,6 +107,8 @@ class WorkflowRunContext: client_secret=self.secrets[BitwardenConstants.CLIENT_SECRET], client_id=self.secrets[BitwardenConstants.CLIENT_ID], master_password=self.secrets[BitwardenConstants.MASTER_PASSWORD], + bw_organization_id=self.secrets[BitwardenConstants.BW_ORGANIZATION_ID], + bw_collection_ids=self.secrets[BitwardenConstants.BW_COLLECTION_IDS], ) return secret_credentials @@ -117,6 +120,7 @@ class WorkflowRunContext: self, aws_client: AsyncAWSClient, parameter: PARAMETER_TYPE, + organization: Organization, ) -> None: if parameter.parameter_type == ParameterType.WORKFLOW: LOG.error(f"Workflow parameters are set while initializing context manager. Parameter key: {parameter.key}") @@ -165,10 +169,14 @@ class WorkflowRunContext: client_id, client_secret, master_password, + organization.bw_organization_id, + organization.bw_collection_ids, url, collection_id=collection_id, ) if secret_credentials: + self.secrets[BitwardenConstants.BW_ORGANIZATION_ID] = organization.bw_organization_id + self.secrets[BitwardenConstants.BW_COLLECTION_IDS] = organization.bw_collection_ids self.secrets[BitwardenConstants.URL] = url self.secrets[BitwardenConstants.CLIENT_SECRET] = client_secret self.secrets[BitwardenConstants.CLIENT_ID] = client_id @@ -223,11 +231,15 @@ class WorkflowRunContext: client_id, client_secret, master_password, + organization.bw_organization_id, + organization.bw_collection_ids, collection_id, bitwarden_identity_key, parameter.bitwarden_identity_fields, ) if sensitive_values: + self.secrets[BitwardenConstants.BW_ORGANIZATION_ID] = organization.bw_organization_id + self.secrets[BitwardenConstants.BW_COLLECTION_IDS] = organization.bw_collection_ids self.secrets[BitwardenConstants.IDENTITY_KEY] = bitwarden_identity_key self.secrets[BitwardenConstants.CLIENT_SECRET] = client_secret self.secrets[BitwardenConstants.CLIENT_ID] = client_id @@ -333,6 +345,7 @@ class WorkflowRunContext: self, aws_client: AsyncAWSClient, parameters: list[PARAMETER_TYPE], + organization: Organization, ) -> None: # Sort the parameters so that ContextParameter and BitwardenLoginCredentialParameter are processed last # ContextParameter should be processed at the end since it requires the source parameter to be set @@ -369,7 +382,7 @@ class WorkflowRunContext: ) self.parameters[parameter.key] = parameter - await self.register_parameter_value(aws_client, parameter) + await self.register_parameter_value(aws_client, parameter, organization) class WorkflowContextManager: @@ -410,6 +423,9 @@ class WorkflowContextManager: self, workflow_run_id: str, parameters: list[PARAMETER_TYPE], + organization: Organization, ) -> None: self._validate_workflow_run_context(workflow_run_id) - await self.workflow_run_contexts[workflow_run_id].register_block_parameters(self.aws_client, parameters) + await self.workflow_run_contexts[workflow_run_id].register_block_parameters( + self.aws_client, parameters, organization + ) diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index ceeb040b..918566ae 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -12,7 +12,7 @@ from skyvern.forge.sdk.artifact.models import ArtifactType from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.core.security import generate_skyvern_signature from skyvern.forge.sdk.core.skyvern_context import SkyvernContext -from skyvern.forge.sdk.models import Step +from skyvern.forge.sdk.models import Organization, Step from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus from skyvern.forge.sdk.workflow.exceptions import ( ContextParameterSourceNotDefined, @@ -150,9 +150,10 @@ class WorkflowService: self, workflow_run_id: str, api_key: str, - organization_id: str | None = None, + organization: Organization, ) -> WorkflowRun: """Execute a workflow.""" + organization_id = organization.organization_id workflow_run = await self.get_workflow_run(workflow_run_id=workflow_run_id) workflow = await self.get_workflow(workflow_id=workflow_run.workflow_id, organization_id=organization_id) @@ -181,7 +182,7 @@ class WorkflowService: try: parameters = block.get_all_parameters(workflow_run_id) await app.WORKFLOW_CONTEXT_MANAGER.register_block_parameters_for_workflow_run( - workflow_run_id, parameters + workflow_run_id, parameters, organization ) LOG.info( f"Executing root block {block.block_type} at index {block_idx} for workflow run {workflow_run_id}",