From cc0d9653f0af910241c52a42f506fe39cac37958 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 28 Feb 2025 09:07:50 -0500 Subject: [PATCH] Add totp field in bitwarden password credential (#1855) Co-authored-by: Muhammed Salih Altun --- skyvern/forge/sdk/schemas/credentials.py | 1 + skyvern/forge/sdk/services/bitwarden.py | 4 ++++ skyvern/forge/sdk/workflow/context_manager.py | 20 +++++++++++++++---- 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/skyvern/forge/sdk/schemas/credentials.py b/skyvern/forge/sdk/schemas/credentials.py index b9262f1b..074df332 100644 --- a/skyvern/forge/sdk/schemas/credentials.py +++ b/skyvern/forge/sdk/schemas/credentials.py @@ -21,6 +21,7 @@ class CreditCardCredentialResponse(BaseModel): class PasswordCredential(BaseModel): password: str username: str + totp: str | None = None class CreditCardCredential(BaseModel): diff --git a/skyvern/forge/sdk/services/bitwarden.py b/skyvern/forge/sdk/services/bitwarden.py index 0c631b91..5cbf5529 100644 --- a/skyvern/forge/sdk/services/bitwarden.py +++ b/skyvern/forge/sdk/services/bitwarden.py @@ -60,6 +60,7 @@ def get_list_response_item_from_bitwarden_item(item: dict) -> CredentialItem: credential=PasswordCredential( username=login["username"], password=login["password"], + totp=login["totp"], ), name=item["name"], credential_type=CredentialType.PASSWORD, @@ -726,6 +727,7 @@ class BitwardenService: return PasswordCredential( username=login["username"], password=login["password"], + totp=login["totp"], ) @staticmethod @@ -743,6 +745,7 @@ class BitwardenService: login_template["username"] = credential.username login_template["password"] = credential.password + login_template["totp"] = credential.totp item_template["type"] = get_bitwarden_item_type_code(BitwardenItemType.LOGIN) item_template["name"] = name @@ -968,6 +971,7 @@ class BitwardenService: credential=PasswordCredential( username=login_item["username"], password=login_item["password"], + totp=login_item["totp"], ), ) elif response["data"]["type"] == BitwardenItemType.CREDIT_CARD: diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index ae4dd9bc..82e529e0 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -12,6 +12,7 @@ from skyvern.exceptions import ( ) from skyvern.forge import app from skyvern.forge.sdk.api.aws import AsyncAWSClient +from skyvern.forge.sdk.schemas.credentials import PasswordCredential from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import TaskStatus from skyvern.forge.sdk.services.bitwarden import BitwardenConstants, BitwardenService @@ -177,22 +178,33 @@ class WorkflowRunContext: organization: Organization, ) -> None: LOG.info(f"Fetching credential parameter value for credential: {parameter.credential_id}") - credential_item = await app.DATABASE.get_credential( + db_credential = await app.DATABASE.get_credential( parameter.credential_id, organization_id=organization.organization_id ) - if credential_item is None: + if db_credential is None: raise CredentialParameterNotFoundError(parameter.credential_id) - credential = await BitwardenService.get_credential_item(credential_item.item_id) - credential_dict = credential.credential.model_dump() + + bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id) + + credential_item = bitwarden_credential.credential self.parameters[parameter.key] = parameter self.values[parameter.key] = {} + credential_dict = credential_item.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: + 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] = credential_item.totp + self.values[parameter.key]["totp"] = totp_secret_id + async def register_aws_secret_parameter_value( self, aws_client: AsyncAWSClient,