Fix MFA resolution priority: credential TOTP over webhook (#SKY-7980) (#4800)
Co-authored-by: Suchintan Singh <suchintan@skyvern.com>
This commit is contained in:
@@ -19,7 +19,7 @@ from skyvern.forge.sdk.core import skyvern_context
|
||||
from skyvern.forge.sdk.schemas.totp_codes import OTPType
|
||||
from skyvern.schemas.workflows import BlockStatus
|
||||
from skyvern.services import script_service
|
||||
from skyvern.services.otp_service import poll_otp_value
|
||||
from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential
|
||||
from skyvern.utils.prompt_engine import load_prompt_with_elements
|
||||
from skyvern.webeye.actions import handler_utils
|
||||
from skyvern.webeye.actions.actions import (
|
||||
@@ -245,7 +245,10 @@ class RealSkyvernPageAi(SkyvernPageAi):
|
||||
if value and isinstance(data, dict) and "value" not in data:
|
||||
data["value"] = value
|
||||
|
||||
if (totp_identifier or totp_url) and context and organization_id and task_id:
|
||||
# Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier)
|
||||
otp_value = try_generate_totp_from_credential(workflow_run_id)
|
||||
# Fall back to webhook/totp_identifier
|
||||
if not otp_value and (totp_identifier or totp_url) and context and organization_id and task_id:
|
||||
if totp_identifier:
|
||||
totp_identifier = _render_template_with_label(totp_identifier, label=self.current_label)
|
||||
if totp_url:
|
||||
@@ -257,16 +260,16 @@ class RealSkyvernPageAi(SkyvernPageAi):
|
||||
totp_identifier=totp_identifier,
|
||||
totp_verification_url=totp_url,
|
||||
)
|
||||
if otp_value and otp_value.get_otp_type() == OTPType.TOTP:
|
||||
verification_code = otp_value.value
|
||||
if isinstance(data, dict) and SPECIAL_FIELD_VERIFICATION_CODE not in data:
|
||||
data[SPECIAL_FIELD_VERIFICATION_CODE] = verification_code
|
||||
elif isinstance(data, str) and SPECIAL_FIELD_VERIFICATION_CODE not in data:
|
||||
data = f"{data}\n" + str({SPECIAL_FIELD_VERIFICATION_CODE: verification_code})
|
||||
elif isinstance(data, list):
|
||||
data.append({SPECIAL_FIELD_VERIFICATION_CODE: verification_code})
|
||||
else:
|
||||
data = {SPECIAL_FIELD_VERIFICATION_CODE: verification_code}
|
||||
if otp_value and otp_value.get_otp_type() == OTPType.TOTP:
|
||||
verification_code = otp_value.value
|
||||
if isinstance(data, dict) and SPECIAL_FIELD_VERIFICATION_CODE not in data:
|
||||
data[SPECIAL_FIELD_VERIFICATION_CODE] = verification_code
|
||||
elif isinstance(data, str) and SPECIAL_FIELD_VERIFICATION_CODE not in data:
|
||||
data = f"{data}\n" + str({SPECIAL_FIELD_VERIFICATION_CODE: verification_code})
|
||||
elif isinstance(data, list):
|
||||
data.append({SPECIAL_FIELD_VERIFICATION_CODE: verification_code})
|
||||
else:
|
||||
data = {SPECIAL_FIELD_VERIFICATION_CODE: verification_code}
|
||||
|
||||
await self._refresh_scraped_page(take_screenshots=False)
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ from skyvern.forge.sdk.api.files import (
|
||||
from skyvern.forge.sdk.artifact.models import ArtifactType
|
||||
from skyvern.forge.sdk.core import skyvern_context
|
||||
from skyvern.schemas.steps import AgentStepOutput
|
||||
from skyvern.services.otp_service import poll_otp_value
|
||||
from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential
|
||||
from skyvern.utils.url_validators import prepend_scheme_and_validate_url
|
||||
from skyvern.webeye.actions.action_types import ActionType
|
||||
from skyvern.webeye.actions.actions import (
|
||||
@@ -615,16 +615,21 @@ class ScriptSkyvernPage(SkyvernPage):
|
||||
if is_totp_value:
|
||||
value = generate_totp_value(context.workflow_run_id, original_value)
|
||||
elif (totp_identifier or totp_url) and organization_id:
|
||||
totp_value = await poll_otp_value(
|
||||
organization_id=organization_id,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
totp_verification_url=totp_url,
|
||||
totp_identifier=totp_identifier,
|
||||
)
|
||||
if totp_value:
|
||||
# use the totp verification code
|
||||
value = totp_value.value
|
||||
# Try credential TOTP first (higher priority than webhook/totp_identifier)
|
||||
credential_totp = try_generate_totp_from_credential(workflow_run_id)
|
||||
if credential_totp:
|
||||
value = credential_totp.value
|
||||
else:
|
||||
totp_value = await poll_otp_value(
|
||||
organization_id=organization_id,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
totp_verification_url=totp_url,
|
||||
totp_identifier=totp_identifier,
|
||||
)
|
||||
if totp_value:
|
||||
# use the totp verification code
|
||||
value = totp_value.value
|
||||
|
||||
return value
|
||||
|
||||
|
||||
@@ -104,7 +104,7 @@ from skyvern.schemas.runs import CUA_ENGINES, RunEngine
|
||||
from skyvern.schemas.steps import AgentStepOutput
|
||||
from skyvern.services import run_service, service_utils
|
||||
from skyvern.services.action_service import get_action_history
|
||||
from skyvern.services.otp_service import poll_otp_value
|
||||
from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential
|
||||
from skyvern.utils.image_resizer import Resolution
|
||||
from skyvern.utils.prompt_engine import MaxStepsReasonResponse, load_prompt_with_elements
|
||||
from skyvern.webeye.actions.action_types import ActionType
|
||||
@@ -4526,21 +4526,25 @@ class ForgeAgent:
|
||||
should_enter_verification_code = json_response.get("should_enter_verification_code")
|
||||
if place_to_enter_verification_code and should_enter_verification_code and task.organization_id:
|
||||
LOG.info("Need verification code")
|
||||
workflow_id = workflow_permanent_id = None
|
||||
if task.workflow_run_id:
|
||||
workflow_run = await app.DATABASE.get_workflow_run(task.workflow_run_id)
|
||||
if workflow_run:
|
||||
workflow_id = workflow_run.workflow_id
|
||||
workflow_permanent_id = workflow_run.workflow_permanent_id
|
||||
otp_value = await poll_otp_value(
|
||||
organization_id=task.organization_id,
|
||||
task_id=task.task_id,
|
||||
workflow_id=workflow_id,
|
||||
workflow_run_id=task.workflow_run_id,
|
||||
workflow_permanent_id=workflow_permanent_id,
|
||||
totp_verification_url=task.totp_verification_url,
|
||||
totp_identifier=task.totp_identifier,
|
||||
)
|
||||
# Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier)
|
||||
otp_value = try_generate_totp_from_credential(task.workflow_run_id)
|
||||
# Fall back to webhook/totp_identifier
|
||||
if not otp_value and (task.totp_verification_url or task.totp_identifier):
|
||||
workflow_id = workflow_permanent_id = None
|
||||
if task.workflow_run_id:
|
||||
workflow_run = await app.DATABASE.get_workflow_run(task.workflow_run_id)
|
||||
if workflow_run:
|
||||
workflow_id = workflow_run.workflow_id
|
||||
workflow_permanent_id = workflow_run.workflow_permanent_id
|
||||
otp_value = await poll_otp_value(
|
||||
organization_id=task.organization_id,
|
||||
task_id=task.task_id,
|
||||
workflow_id=workflow_id,
|
||||
workflow_run_id=task.workflow_run_id,
|
||||
workflow_permanent_id=workflow_permanent_id,
|
||||
totp_verification_url=task.totp_verification_url,
|
||||
totp_identifier=task.totp_identifier,
|
||||
)
|
||||
if not otp_value or otp_value.get_otp_type() != OTPType.TOTP:
|
||||
return json_response
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pyotp
|
||||
import structlog
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@@ -57,6 +58,47 @@ async def parse_otp_login(
|
||||
return None
|
||||
|
||||
|
||||
def try_generate_totp_from_credential(workflow_run_id: str | None) -> OTPValue | None:
|
||||
"""Try to generate a TOTP code from a credential secret stored in the workflow run context.
|
||||
|
||||
Scans workflow_run_context.values for credential entries with a "totp" key
|
||||
(e.g. Bitwarden, 1Password, Azure Key Vault credentials) and generates a
|
||||
TOTP code using pyotp. This should be checked BEFORE poll_otp_value so that
|
||||
credential-based TOTP takes priority over webhook (totp_url) and totp_identifier.
|
||||
"""
|
||||
if not workflow_run_id:
|
||||
return None
|
||||
|
||||
workflow_run_context = app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context(workflow_run_id)
|
||||
if not workflow_run_context:
|
||||
return None
|
||||
|
||||
for key, value in workflow_run_context.values.items():
|
||||
if isinstance(value, dict) and "totp" in value:
|
||||
totp_secret_id = value.get("totp")
|
||||
if not totp_secret_id or not isinstance(totp_secret_id, str):
|
||||
continue
|
||||
totp_secret_key = workflow_run_context.totp_secret_value_key(totp_secret_id)
|
||||
totp_secret = workflow_run_context.get_original_secret_value_or_none(totp_secret_key)
|
||||
if totp_secret:
|
||||
try:
|
||||
code = pyotp.TOTP(totp_secret).now()
|
||||
LOG.info(
|
||||
"Generated TOTP from credential secret",
|
||||
workflow_run_id=workflow_run_id,
|
||||
credential_key=key,
|
||||
)
|
||||
return OTPValue(value=code, type=OTPType.TOTP)
|
||||
except Exception:
|
||||
LOG.warning(
|
||||
"Failed to generate TOTP from credential secret",
|
||||
workflow_run_id=workflow_run_id,
|
||||
credential_key=key,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def poll_otp_value(
|
||||
organization_id: str,
|
||||
task_id: str | None = None,
|
||||
|
||||
@@ -14,7 +14,7 @@ from skyvern.forge.sdk.core import skyvern_context
|
||||
from skyvern.forge.sdk.models import Step
|
||||
from skyvern.forge.sdk.schemas.tasks import Task
|
||||
from skyvern.forge.sdk.schemas.totp_codes import OTPType
|
||||
from skyvern.services.otp_service import poll_otp_value
|
||||
from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential
|
||||
from skyvern.utils.image_resizer import Resolution, scale_coordinates
|
||||
from skyvern.webeye.actions.action_types import ActionType
|
||||
from skyvern.webeye.actions.actions import (
|
||||
@@ -913,7 +913,10 @@ async def generate_cua_fallback_actions(
|
||||
)
|
||||
|
||||
elif skyvern_action_type == "get_verification_code":
|
||||
if (task.totp_verification_url or task.totp_identifier) and task.organization_id:
|
||||
# Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier)
|
||||
otp_value = try_generate_totp_from_credential(task.workflow_run_id)
|
||||
# Fall back to webhook/totp_identifier
|
||||
if not otp_value and (task.totp_verification_url or task.totp_identifier) and task.organization_id:
|
||||
LOG.info(
|
||||
"Getting verification code for CUA",
|
||||
task_id=task.task_id,
|
||||
@@ -930,29 +933,21 @@ async def generate_cua_fallback_actions(
|
||||
totp_verification_url=task.totp_verification_url,
|
||||
totp_identifier=task.totp_identifier,
|
||||
)
|
||||
if not otp_value or otp_value.get_otp_type() != OTPType.TOTP:
|
||||
raise NoTOTPVerificationCodeFound()
|
||||
verification_code = otp_value.value
|
||||
reasoning = reasoning or f"Received verification code: {verification_code}"
|
||||
action = VerificationCodeAction(
|
||||
verification_code=verification_code,
|
||||
reasoning=reasoning,
|
||||
intention=reasoning,
|
||||
)
|
||||
except NoTOTPVerificationCodeFound:
|
||||
reasoning_suffix = "No verification code found"
|
||||
reasoning = f"{reasoning}. {reasoning_suffix}" if reasoning else reasoning_suffix
|
||||
action = TerminateAction(
|
||||
reasoning=reasoning,
|
||||
intention=reasoning,
|
||||
)
|
||||
except FailedToGetTOTPVerificationCode as e:
|
||||
reasoning_suffix = f"Failed to get verification code. Reason: {e.reason}"
|
||||
reasoning = f"{reasoning}. {reasoning_suffix}" if reasoning else reasoning_suffix
|
||||
action = TerminateAction(
|
||||
reasoning=reasoning,
|
||||
intention=reasoning,
|
||||
)
|
||||
|
||||
if otp_value and otp_value.get_otp_type() == OTPType.TOTP:
|
||||
verification_code = otp_value.value
|
||||
reasoning = reasoning or f"Received verification code: {verification_code}"
|
||||
action = VerificationCodeAction(
|
||||
verification_code=verification_code,
|
||||
reasoning=reasoning,
|
||||
intention=reasoning,
|
||||
)
|
||||
else:
|
||||
action = TerminateAction(
|
||||
reasoning=reasoning,
|
||||
|
||||
Reference in New Issue
Block a user