introduce pyotp and support generating totp within skyvern (#1176)

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"

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"

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 {}
 

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

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()
+        totp_secret_key = workflow_run_context.totp_secret_value_key(parameter)
-        secret_value = secrets[BitwardenConstants.TOTP]
+        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