From 0ede4fdfa07fa0e30a545ff616ba507613eaf24a Mon Sep 17 00:00:00 2001 From: Kerem Yilmaz Date: Mon, 10 Jun 2024 22:06:58 -0700 Subject: [PATCH] Introduce collectionid filter for bitwarden parameters (#454) --- ...e606a3d_add_collection_id_to_bitwarden_.py | 33 +++++++++++++++++++ skyvern/forge/sdk/db/client.py | 2 ++ skyvern/forge/sdk/db/models.py | 1 + skyvern/forge/sdk/db/utils.py | 2 ++ skyvern/forge/sdk/services/bitwarden.py | 16 ++++----- skyvern/forge/sdk/workflow/context_manager.py | 3 +- .../forge/sdk/workflow/models/parameter.py | 3 ++ skyvern/forge/sdk/workflow/models/yaml.py | 3 ++ skyvern/forge/sdk/workflow/service.py | 3 ++ 9 files changed, 56 insertions(+), 10 deletions(-) create mode 100644 alembic/versions/2024_06_11_0502-2c163e606a3d_add_collection_id_to_bitwarden_.py diff --git a/alembic/versions/2024_06_11_0502-2c163e606a3d_add_collection_id_to_bitwarden_.py b/alembic/versions/2024_06_11_0502-2c163e606a3d_add_collection_id_to_bitwarden_.py new file mode 100644 index 00000000..76f033f0 --- /dev/null +++ b/alembic/versions/2024_06_11_0502-2c163e606a3d_add_collection_id_to_bitwarden_.py @@ -0,0 +1,33 @@ +"""Add collection id to bitwarden credential parameters + +Revision ID: 2c163e606a3d +Revises: 312d305c6b18 +Create Date: 2024-06-11 05:02:25.023252+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "2c163e606a3d" +down_revision: Union[str, None] = "312d305c6b18" +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( + "bitwarden_login_credential_parameters", sa.Column("bitwarden_collection_id", sa.String(), nullable=True) + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("bitwarden_login_credential_parameters", "bitwarden_collection_id") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 98f96950..a2463788 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -1026,6 +1026,7 @@ class AgentDB: url_parameter_key: str, key: str, description: str | None = None, + bitwarden_collection_id: str | None = None, ) -> BitwardenLoginCredentialParameter: async with self.Session() as session: bitwarden_login_credential_parameter = BitwardenLoginCredentialParameterModel( @@ -1036,6 +1037,7 @@ class AgentDB: url_parameter_key=url_parameter_key, key=key, description=description, + bitwarden_collection_id=bitwarden_collection_id, ) session.add(bitwarden_login_credential_parameter) await session.commit() diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index bd8042f8..a3b02e4a 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -278,6 +278,7 @@ class BitwardenLoginCredentialParameterModel(Base): bitwarden_client_id_aws_secret_key = Column(String, nullable=False) bitwarden_client_secret_aws_secret_key = Column(String, nullable=False) bitwarden_master_password_aws_secret_key = Column(String, nullable=False) + bitwarden_collection_id = Column(String, nullable=True, default=None) url_parameter_key = Column(String, nullable=False) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column( diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index 528fef67..582eb8a6 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -241,6 +241,7 @@ def convert_to_bitwarden_login_credential_parameter( LOG.debug( "Converting BitwardenLoginCredentialParameterModel to BitwardenLoginCredentialParameter", bitwarden_login_credential_parameter_id=bitwarden_login_credential_parameter_model.bitwarden_login_credential_parameter_id, + bitwarden_collection_id=bitwarden_login_credential_parameter_model.bitwarden_collection_id, ) return BitwardenLoginCredentialParameter( @@ -251,6 +252,7 @@ def convert_to_bitwarden_login_credential_parameter( bitwarden_client_id_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_client_id_aws_secret_key, bitwarden_client_secret_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_client_secret_aws_secret_key, bitwarden_master_password_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_master_password_aws_secret_key, + bitwarden_collection_id=bitwarden_login_credential_parameter_model.bitwarden_collection_id, url_parameter_key=bitwarden_login_credential_parameter_model.url_parameter_key, created_at=bitwarden_login_credential_parameter_model.created_at, modified_at=bitwarden_login_credential_parameter_model.modified_at, diff --git a/skyvern/forge/sdk/services/bitwarden.py b/skyvern/forge/sdk/services/bitwarden.py index 20c2946a..31a31f68 100644 --- a/skyvern/forge/sdk/services/bitwarden.py +++ b/skyvern/forge/sdk/services/bitwarden.py @@ -27,6 +27,7 @@ class BitwardenConstants(StrEnum): CLIENT_SECRET = "BW_CLIENT_SECRET" MASTER_PASSWORD = "BW_MASTER_PASSWORD" URL = "BW_URL" + BW_COLLECTION_ID = "BW_COLLECTION_ID" USERNAME = "BW_USERNAME" PASSWORD = "BW_PASSWORD" @@ -68,6 +69,7 @@ class BitwardenService: client_secret: str, master_password: str, url: str, + collection_id: str | None = None, ) -> dict[str, str]: """ Get the secret value from the Bitwarden CLI. @@ -105,14 +107,6 @@ class BitwardenService: f"Failed to unlock vault. stdout: {unlock_result.stdout} stderr: {unlock_result.stderr}" ) - # This is a part of Bitwarden's client-side telemetry - # TODO -- figure out how to disable this telemetry so we never get this error - # https://github.com/bitwarden/clients/blob/9d10825dbd891c0f41fe1b4c4dd3ca4171f63be5/libs/common/src/services/api.service.ts#L1473 - if unlock_result.stderr and "Event post failed" not in unlock_result.stderr: - raise BitwardenUnlockError( - f"Failed to unlock vault. stdout: {unlock_result.stdout} stderr: {unlock_result.stderr}" - ) - # Extract session key try: session_key = BitwardenService._extract_session_key(unlock_result.stdout) @@ -132,6 +126,9 @@ class BitwardenService: "--session", session_key, ] + if 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) if items_result.stderr and "Event post failed" not in items_result.stderr: @@ -144,7 +141,8 @@ class BitwardenService: raise BitwardenListItemsError("Failed to parse items JSON. Output: " + items_result.stdout) if not items: - raise BitwardenListItemsError("No items found in Bitwarden.") + 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}") totp_command = ["bw", "get", "totp", url, "--session", session_key] totp_result = BitwardenService.run_command(totp_command) diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 1b37fdc0..03f12e8b 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -157,12 +157,14 @@ class WorkflowRunContext: client_secret, master_password, url, + collection_id=parameter.bitwarden_collection_id, ) if secret_credentials: self.secrets[BitwardenConstants.URL] = url self.secrets[BitwardenConstants.CLIENT_SECRET] = client_secret self.secrets[BitwardenConstants.CLIENT_ID] = client_id self.secrets[BitwardenConstants.MASTER_PASSWORD] = master_password + self.secrets[BitwardenConstants.BW_COLLECTION_ID] = parameter.bitwarden_collection_id random_secret_id = self.generate_random_secret_id() # username secret @@ -181,7 +183,6 @@ class WorkflowRunContext: "totp": totp_secret_id, } except BitwardenBaseError as e: - BitwardenService.logout() LOG.error(f"Failed to get secret from Bitwarden. Error: {e}") raise e elif isinstance(parameter, ContextParameter): diff --git a/skyvern/forge/sdk/workflow/models/parameter.py b/skyvern/forge/sdk/workflow/models/parameter.py index 8d50469a..b0f3fe07 100644 --- a/skyvern/forge/sdk/workflow/models/parameter.py +++ b/skyvern/forge/sdk/workflow/models/parameter.py @@ -52,6 +52,9 @@ class BitwardenLoginCredentialParameter(Parameter): bitwarden_master_password_aws_secret_key: str # url to request the login credentials from bitwarden url_parameter_key: str + # bitwarden collection id to filter the login credentials from, + # if not provided, no filtering will be done + bitwarden_collection_id: str | None = None created_at: datetime modified_at: datetime diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py index 97dca26e..9bd0a113 100644 --- a/skyvern/forge/sdk/workflow/models/yaml.py +++ b/skyvern/forge/sdk/workflow/models/yaml.py @@ -36,6 +36,9 @@ class BitwardenLoginCredentialParameterYAML(ParameterYAML): bitwarden_master_password_aws_secret_key: str # parameter key for the url to request the login credentials from bitwarden url_parameter_key: str + # bitwarden collection id to filter the login credentials from, + # if not provided, no filtering will be done + bitwarden_collection_id: str | None = None class WorkflowParameterYAML(ParameterYAML): diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 479f399f..aaebe99b 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -450,6 +450,7 @@ class WorkflowService: url_parameter_key: str, key: str, description: str | None = None, + bitwarden_collection_id: str | None = None, ) -> Parameter: return await app.DATABASE.create_bitwarden_login_credential_parameter( workflow_id=workflow_id, @@ -459,6 +460,7 @@ class WorkflowService: url_parameter_key=url_parameter_key, key=key, description=description, + bitwarden_collection_id=bitwarden_collection_id, ) async def create_output_parameter( @@ -833,6 +835,7 @@ class WorkflowService: url_parameter_key=parameter.url_parameter_key, key=parameter.key, description=parameter.description, + bitwarden_collection_id=parameter.bitwarden_collection_id, ) elif parameter.parameter_type == ParameterType.WORKFLOW: parameters[parameter.key] = await self.create_workflow_parameter(