support uri format otp secrete (#3126)

skyvern/exceptions.py

@@ -744,3 +744,8 @@ class ElementOutOfCurrentViewport(SkyvernException):
 class ScriptNotFound(SkyvernHTTPException):
     def __init__(self, script_id: str) -> None:
         super().__init__(f"Script {script_id} not found")
+
+
+class NoTOTPSecretFound(SkyvernException):
+    def __init__(self) -> None:
+        super().__init__("No TOTP secret found")

skyvern/forge/sdk/services/bitwarden.py

@@ -2,7 +2,6 @@ import asyncio
 import json
 import os
 import re
-import urllib.parse
 from enum import IntEnum, StrEnum
 from typing import Tuple
 
@@ -32,6 +31,7 @@ from skyvern.forge.sdk.schemas.credentials import (
     CreditCardCredential,
     PasswordCredential,
 )
+from skyvern.forge.sdk.services.credentials import parse_totp_secret
 
 LOG = structlog.get_logger()
 BITWARDEN_SERVER_BASE_URL = f"{settings.BITWARDEN_SERVER}:{settings.BITWARDEN_SERVER_PORT or 8002}"
@@ -245,25 +245,7 @@ class BitwardenService:
             >>> BitwardenService.extract_totp_secret("otpauth://totp/user@domain.com?secret=AAAAAABBBBBBB")
             "AAAAAABBBBBBB"
         """
-        if not totp_value:
-            return ""
-
-        # Handle TOTP URI format
-        if totp_value.startswith("otpauth://"):
-            try:
-                # Parse the URI to extract the secret
-                query = urllib.parse.urlparse(totp_value).query
-                params = dict(urllib.parse.parse_qsl(query))
-                return params.get("secret", "")
-            except Exception:
-                LOG.error(
-                    "Failed to parse TOTP URI",
-                    totp_value=totp_value,
-                    exc_info=True,
-                )
-                return ""
-
-        return totp_value
+        return parse_totp_secret(totp_value)
 
     @staticmethod
     async def _get_secret_value_from_url(

skyvern/forge/sdk/services/credentials.py

@@ -1,10 +1,39 @@
-import logging
+import re
 from enum import StrEnum
+from urllib.parse import unquote
 
-LOG = logging.getLogger(__name__)
+import pyotp
+import structlog
+
+from skyvern.exceptions import NoTOTPSecretFound
+
+LOG = structlog.get_logger()
 
 
 class OnePasswordConstants(StrEnum):
     """Constants for 1Password integration."""
 
     TOTP = "OP_TOTP"  # Special value to indicate a TOTP code
+
+
+def parse_totp_secret(totp_secret: str) -> str:
+    if not totp_secret:
+        return ""
+
+    totp_secret_no_whitespace = "".join(totp_secret.split())
+    if len(totp_secret_no_whitespace) == 32:
+        return totp_secret_no_whitespace
+
+    LOG.info("TOTP secret key is not 32 characters, try to parse it from URI format")
+    try:
+        totp_secret = pyotp.parse_uri(totp_secret_no_whitespace).secret
+        totp_secret_no_whitespace = "".join(totp_secret.split())
+        return totp_secret_no_whitespace
+    except Exception:
+        LOG.warning("Failed to parse TOTP secret key from URI format, going to extract secret by regex", exc_info=True)
+        m = re.search(r"(?i)(?:^|[?&])secret=([^&#]+)", unquote(totp_secret_no_whitespace))
+        if m is None:
+            raise NoTOTPSecretFound()
+        totp_secret = m.group(1)
+        totp_secret_no_whitespace = "".join(totp_secret.split())
+        return totp_secret_no_whitespace

skyvern/forge/sdk/workflow/context_manager.py

@@ -19,7 +19,7 @@ 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
-from skyvern.forge.sdk.services.credentials import OnePasswordConstants
+from skyvern.forge.sdk.services.credentials import OnePasswordConstants, parse_totp_secret
 from skyvern.forge.sdk.workflow.exceptions import OutputParameterKeyCollisionError
 from skyvern.forge.sdk.workflow.models.parameter import (
     PARAMETER_TYPE,
@@ -362,7 +362,7 @@ class WorkflowRunContext:
                 totp_secret_id = f"{random_secret_id}_totp"
                 self.secrets[totp_secret_id] = OnePasswordConstants.TOTP
                 totp_secret_value = self.totp_secret_value_key(totp_secret_id)
-                self.secrets[totp_secret_value] = field.value
+                self.secrets[totp_secret_value] = parse_totp_secret(field.value)
                 self.values[parameter.key]["totp"] = totp_secret_id
             else:
                 # this will be the username or password or other field

skyvern/webeye/actions/handler.py

@@ -1964,8 +1964,10 @@ def generate_totp_value(task: Task, parameter: str) -> str:
     workflow_run_context = app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context(task.workflow_run_id)
     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)
-    totp_secret_no_whitespace = "".join(totp_secret.split())
-    return pyotp.TOTP(totp_secret_no_whitespace).now()
+    if not totp_secret:
+        LOG.warning("No TOTP secret found, returning the parameter value as is", parameter=parameter)
+        return parameter
+    return pyotp.TOTP(totp_secret).now()
 
 
 async def chain_click(