From 44c9a81156a2ffb7e3a404bfb78f12d1c7b02fd9 Mon Sep 17 00:00:00 2001 From: LawyZheng Date: Thu, 7 Aug 2025 14:59:29 +0800 Subject: [PATCH] support uri format otp secrete (#3126) --- skyvern/exceptions.py | 5 +++ skyvern/forge/sdk/services/bitwarden.py | 22 ++----------- skyvern/forge/sdk/services/credentials.py | 33 +++++++++++++++++-- skyvern/forge/sdk/workflow/context_manager.py | 4 +-- skyvern/webeye/actions/handler.py | 6 ++-- 5 files changed, 44 insertions(+), 26 deletions(-) diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index 3169b939..26603002 100644 --- a/skyvern/exceptions.py +++ b/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") diff --git a/skyvern/forge/sdk/services/bitwarden.py b/skyvern/forge/sdk/services/bitwarden.py index 43d03f05..083afb85 100644 --- a/skyvern/forge/sdk/services/bitwarden.py +++ b/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( diff --git a/skyvern/forge/sdk/services/credentials.py b/skyvern/forge/sdk/services/credentials.py index e4cca78c..206f676e 100644 --- a/skyvern/forge/sdk/services/credentials.py +++ b/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 diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index bd8865b2..29fe1359 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/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 diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 879213cf..b1853050 100644 --- a/skyvern/webeye/actions/handler.py +++ b/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(