[SKY-6] Backend: Enable 2FA code detection without TOTP credentials (#4786)
This commit is contained in:
Aaron Perez
@@ -11,6 +11,7 @@ from skyvern.forge.prompts import prompt_engine
|
||||
from skyvern.forge.sdk.core.aiohttp_helper import aiohttp_post
|
||||
from skyvern.forge.sdk.core.security import generate_skyvern_webhook_signature
|
||||
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
|
||||
from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory
|
||||
from skyvern.forge.sdk.schemas.totp_codes import OTPType
|
||||
|
||||
LOG = structlog.get_logger()
|
||||
@@ -80,38 +81,147 @@ async def poll_otp_value(
|
||||
totp_verification_url=totp_verification_url,
|
||||
totp_identifier=totp_identifier,
|
||||
)
|
||||
while True:
|
||||
await asyncio.sleep(10)
|
||||
# check timeout
|
||||
if datetime.utcnow() > timeout_datetime:
|
||||
LOG.warning("Polling otp value timed out")
|
||||
raise NoTOTPVerificationCodeFound(
|
||||
task_id=task_id,
|
||||
|
||||
# Set the waiting state in the database when polling starts
|
||||
identifier_for_ui = totp_identifier
|
||||
if workflow_run_id:
|
||||
try:
|
||||
await app.DATABASE.update_workflow_run(
|
||||
workflow_run_id=workflow_run_id,
|
||||
workflow_id=workflow_permanent_id,
|
||||
totp_verification_url=totp_verification_url,
|
||||
totp_identifier=totp_identifier,
|
||||
waiting_for_verification_code=True,
|
||||
verification_code_identifier=identifier_for_ui,
|
||||
verification_code_polling_started_at=start_datetime,
|
||||
)
|
||||
otp_value: OTPValue | None = None
|
||||
if totp_verification_url:
|
||||
otp_value = await _get_otp_value_from_url(
|
||||
organization_id,
|
||||
totp_verification_url,
|
||||
org_token.token,
|
||||
task_id=task_id,
|
||||
LOG.info(
|
||||
"Set 2FA waiting state for workflow run",
|
||||
workflow_run_id=workflow_run_id,
|
||||
verification_code_identifier=identifier_for_ui,
|
||||
)
|
||||
elif totp_identifier:
|
||||
otp_value = await _get_otp_value_from_db(
|
||||
organization_id,
|
||||
totp_identifier,
|
||||
try:
|
||||
NotificationRegistryFactory.get_registry().publish(
|
||||
organization_id,
|
||||
{
|
||||
"type": "verification_code_required",
|
||||
"workflow_run_id": workflow_run_id,
|
||||
"task_id": task_id,
|
||||
"identifier": identifier_for_ui,
|
||||
"polling_started_at": start_datetime.isoformat(),
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
LOG.warning("Failed to publish 2FA required notification for workflow run", exc_info=True)
|
||||
except Exception:
|
||||
LOG.warning("Failed to set 2FA waiting state for workflow run", exc_info=True)
|
||||
elif task_id:
|
||||
try:
|
||||
await app.DATABASE.update_task_2fa_state(
|
||||
task_id=task_id,
|
||||
workflow_id=workflow_permanent_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
waiting_for_verification_code=True,
|
||||
verification_code_identifier=identifier_for_ui,
|
||||
verification_code_polling_started_at=start_datetime,
|
||||
)
|
||||
if otp_value:
|
||||
LOG.info("Got otp value", otp_value=otp_value)
|
||||
return otp_value
|
||||
LOG.info(
|
||||
"Set 2FA waiting state for task",
|
||||
task_id=task_id,
|
||||
verification_code_identifier=identifier_for_ui,
|
||||
)
|
||||
try:
|
||||
NotificationRegistryFactory.get_registry().publish(
|
||||
organization_id,
|
||||
{
|
||||
"type": "verification_code_required",
|
||||
"task_id": task_id,
|
||||
"identifier": identifier_for_ui,
|
||||
"polling_started_at": start_datetime.isoformat(),
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
LOG.warning("Failed to publish 2FA required notification for task", exc_info=True)
|
||||
except Exception:
|
||||
LOG.warning("Failed to set 2FA waiting state for task", exc_info=True)
|
||||
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(10)
|
||||
# check timeout
|
||||
if datetime.utcnow() > timeout_datetime:
|
||||
LOG.warning("Polling otp value timed out")
|
||||
raise NoTOTPVerificationCodeFound(
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
workflow_id=workflow_permanent_id,
|
||||
totp_verification_url=totp_verification_url,
|
||||
totp_identifier=totp_identifier,
|
||||
)
|
||||
otp_value: OTPValue | None = None
|
||||
if totp_verification_url:
|
||||
otp_value = await _get_otp_value_from_url(
|
||||
organization_id,
|
||||
totp_verification_url,
|
||||
org_token.token,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
)
|
||||
elif totp_identifier:
|
||||
otp_value = await _get_otp_value_from_db(
|
||||
organization_id,
|
||||
totp_identifier,
|
||||
task_id=task_id,
|
||||
workflow_id=workflow_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
)
|
||||
if not otp_value:
|
||||
otp_value = await _get_otp_value_by_run(
|
||||
organization_id,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
)
|
||||
else:
|
||||
# No pre-configured TOTP — poll for manually submitted codes by run context
|
||||
otp_value = await _get_otp_value_by_run(
|
||||
organization_id,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
)
|
||||
if otp_value:
|
||||
LOG.info("Got otp value", otp_value=otp_value)
|
||||
return otp_value
|
||||
finally:
|
||||
# Clear the waiting state when polling completes (success, timeout, or error)
|
||||
if workflow_run_id:
|
||||
try:
|
||||
await app.DATABASE.update_workflow_run(
|
||||
workflow_run_id=workflow_run_id,
|
||||
waiting_for_verification_code=False,
|
||||
)
|
||||
LOG.info("Cleared 2FA waiting state for workflow run", workflow_run_id=workflow_run_id)
|
||||
try:
|
||||
NotificationRegistryFactory.get_registry().publish(
|
||||
organization_id,
|
||||
{"type": "verification_code_resolved", "workflow_run_id": workflow_run_id, "task_id": task_id},
|
||||
)
|
||||
except Exception:
|
||||
LOG.warning("Failed to publish 2FA resolved notification for workflow run", exc_info=True)
|
||||
except Exception:
|
||||
LOG.warning("Failed to clear 2FA waiting state for workflow run", exc_info=True)
|
||||
elif task_id:
|
||||
try:
|
||||
await app.DATABASE.update_task_2fa_state(
|
||||
task_id=task_id,
|
||||
organization_id=organization_id,
|
||||
waiting_for_verification_code=False,
|
||||
)
|
||||
LOG.info("Cleared 2FA waiting state for task", task_id=task_id)
|
||||
try:
|
||||
NotificationRegistryFactory.get_registry().publish(
|
||||
organization_id,
|
||||
{"type": "verification_code_resolved", "task_id": task_id},
|
||||
)
|
||||
except Exception:
|
||||
LOG.warning("Failed to publish 2FA resolved notification for task", exc_info=True)
|
||||
except Exception:
|
||||
LOG.warning("Failed to clear 2FA waiting state for task", exc_info=True)
|
||||
|
||||
|
||||
async def _get_otp_value_from_url(
|
||||
@@ -175,6 +285,28 @@ async def _get_otp_value_from_url(
|
||||
return otp_value
|
||||
|
||||
|
||||
async def _get_otp_value_by_run(
|
||||
organization_id: str,
|
||||
task_id: str | None = None,
|
||||
workflow_run_id: str | None = None,
|
||||
) -> OTPValue | None:
|
||||
"""Look up OTP codes by task_id/workflow_run_id when no totp_identifier is configured.
|
||||
|
||||
Used for the manual 2FA input flow where users submit codes through the UI
|
||||
without pre-configured TOTP credentials.
|
||||
"""
|
||||
codes = await app.DATABASE.get_otp_codes_by_run(
|
||||
organization_id=organization_id,
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
limit=1,
|
||||
)
|
||||
if codes:
|
||||
code = codes[0]
|
||||
return OTPValue(value=code.code, type=code.otp_type)
|
||||
return None
|
||||
|
||||
|
||||
async def _get_otp_value_from_db(
|
||||
organization_id: str,
|
||||
totp_identifier: str,
|
||||
|
||||
Reference in New Issue
Block a user