From efdf4b5cab08d8607900b302ce4286473cdc1d44 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Tue, 12 Nov 2024 02:17:22 -0800 Subject: [PATCH] introduce pyotp and support generating totp within skyvern (#1176) --- poetry.lock | 16 +++++++++++++- pyproject.toml | 1 + skyvern/forge/sdk/services/bitwarden.py | 22 +------------------ skyvern/forge/sdk/workflow/context_manager.py | 5 +++++ skyvern/webeye/actions/handler.py | 6 +++-- 5 files changed, 26 insertions(+), 24 deletions(-) diff --git a/poetry.lock b/poetry.lock index 9cf80b4e..7a1c3345 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3957,6 +3957,20 @@ files = [ [package.extras] windows-terminal = ["colorama (>=0.4.6)"] +[[package]] +name = "pyotp" +version = "2.9.0" +description = "Python One Time Password Library" +optional = false +python-versions = ">=3.7" +files = [ + {file = "pyotp-2.9.0-py3-none-any.whl", hash = "sha256:81c2e5865b8ac55e825b0358e496e1d9387c811e85bb40e71a3b29b288963612"}, + {file = "pyotp-2.9.0.tar.gz", hash = "sha256:346b6642e0dbdde3b4ff5a930b664ca82abfa116356ed48cc42c7d6590d36f63"}, +] + +[package.extras] +test = ["coverage", "mypy", "ruff", "wheel"] + [[package]] name = "pypdf" version = "5.1.0" @@ -6092,4 +6106,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = "^3.11,<3.12" -content-hash = "2735a607e8ed37460c94e447b9cb24848046e7bb29b573a7b47a76c85250b08a" +content-hash = "ccde71d44563bdca3436ecf2ca26c73a8b6ca3f20d17687177f4a59aad28064a" diff --git a/pyproject.toml b/pyproject.toml index 7a53d38b..575f28a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ temporalio = "^1.6.0" requests-toolbelt = "^1.0.0" posthog = "^3.7.0" aiofiles = "^24.1.0" +pyotp = "^2.9.0" [tool.poetry.group.dev.dependencies] isort = "^5.13.2" diff --git a/skyvern/forge/sdk/services/bitwarden.py b/skyvern/forge/sdk/services/bitwarden.py index fe05e61c..ba2df00f 100644 --- a/skyvern/forge/sdk/services/bitwarden.py +++ b/skyvern/forge/sdk/services/bitwarden.py @@ -16,7 +16,6 @@ from skyvern.exceptions import ( BitwardenLoginError, BitwardenLogoutError, BitwardenSyncError, - BitwardenTOTPError, BitwardenUnlockError, ) @@ -211,27 +210,12 @@ class BitwardenService: 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", domain, "--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: - LOG.warning( - "Bitwarden TOTP Error", - error=totp_result.stderr, - e=BitwardenTOTPError(totp_result.stderr), - ) - totp_code = totp_result.stdout bitwarden_result: list[BitwardenQueryResult] = [ BitwardenQueryResult( credential={ BitwardenConstants.USERNAME: item.get("login", {}).get("username", ""), BitwardenConstants.PASSWORD: item.get("login", {}).get("password", ""), + BitwardenConstants.TOTP: item.get("login", {}).get("totp", ""), }, uris=[uri.get("uri") for uri in item.get("login", {}).get("uris", []) if "uri" in uri], ) @@ -239,10 +223,6 @@ class BitwardenService: if "login" in item ] - if totp_code: - for single_result in bitwarden_result: - single_result.credential[BitwardenConstants.TOTP] = totp_code - if len(bitwarden_result) == 0: return {} diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index d39ce3df..74431d7e 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -199,6 +199,8 @@ class WorkflowRunContext: if BitwardenConstants.TOTP in secret_credentials and secret_credentials[BitwardenConstants.TOTP]: 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] = secret_credentials[BitwardenConstants.TOTP] self.values[parameter.key]["totp"] = totp_secret_id except BitwardenBaseError as e: @@ -451,6 +453,9 @@ class WorkflowRunContext: self.parameters[parameter.key] = parameter await self.register_parameter_value(aws_client, parameter, organization) + def totp_secret_value_key(self, totp_secret_id: str) -> str: + return f"{totp_secret_id}_value" + class WorkflowContextManager: aws_client: AsyncAWSClient diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 07ca6ecb..246dc9ab 100644 --- a/skyvern/webeye/actions/handler.py +++ b/skyvern/webeye/actions/handler.py @@ -7,6 +7,7 @@ import uuid from datetime import datetime, timedelta, timezone from typing import Any, Awaitable, Callable, List +import pyotp import structlog from playwright.async_api import FileChooser, Frame, Locator, Page, TimeoutError from pydantic import BaseModel @@ -1082,8 +1083,9 @@ async def get_actual_value_of_parameter_if_secret(task: Task, parameter: str) -> secret_value = workflow_run_context.get_original_secret_value_or_none(parameter) if secret_value == BitwardenConstants.TOTP: - secrets = await workflow_run_context.get_secrets_from_password_manager() - secret_value = secrets[BitwardenConstants.TOTP] + totp_secret_key = workflow_run_context.totp_secret_value_key(parameter) + totp_secret = workflow_run_context.get_original_secret_value_or_none(totp_secret_key) + secret_value = pyotp.TOTP(totp_secret).now() return secret_value if secret_value is not None else parameter