fix totp call sign (#3907)

skyvern/core/totp.py

@@ -1,125 +0,0 @@
-import asyncio
-from datetime import datetime, timedelta
-
-import structlog
-
-from skyvern.config import settings
-from skyvern.exceptions import FailedToGetTOTPVerificationCode, NoTOTPVerificationCodeFound
-from skyvern.forge import app
-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.schemas.totp_codes import OTPType
-
-LOG = structlog.get_logger()
-
-
-async def poll_verification_code(
-    organization_id: str,
-    task_id: str | None = None,
-    workflow_id: str | None = None,
-    workflow_run_id: str | None = None,
-    workflow_permanent_id: str | None = None,
-    totp_verification_url: str | None = None,
-    totp_identifier: str | None = None,
-) -> str | None:
-    timeout = timedelta(minutes=settings.VERIFICATION_CODE_POLLING_TIMEOUT_MINS)
-    start_datetime = datetime.utcnow()
-    timeout_datetime = start_datetime + timeout
-    org_token = await app.DATABASE.get_valid_org_auth_token(organization_id, OrganizationAuthTokenType.api.value)
-    if not org_token:
-        LOG.error("Failed to get organization token when trying to get verification code")
-        return None
-    LOG.info(
-        "Polling verification code",
-        task_id=task_id,
-        workflow_run_id=workflow_run_id,
-        workflow_permanent_id=workflow_permanent_id,
-        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 verification code 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,
-            )
-        verification_code = None
-        if totp_verification_url:
-            verification_code = await _get_verification_code_from_url(
-                totp_verification_url,
-                org_token.token,
-                task_id=task_id,
-                workflow_run_id=workflow_run_id,
-            )
-        elif totp_identifier:
-            verification_code = await _get_verification_code_from_db(
-                organization_id,
-                totp_identifier,
-                task_id=task_id,
-                workflow_id=workflow_permanent_id,
-                workflow_run_id=workflow_run_id,
-            )
-        if verification_code:
-            LOG.info("Got verification code", verification_code=verification_code)
-            return verification_code
-
-
-async def _get_verification_code_from_url(
-    url: str,
-    api_key: str,
-    task_id: str | None = None,
-    workflow_run_id: str | None = None,
-    workflow_permanent_id: str | None = None,
-) -> str | None:
-    request_data = {}
-    if task_id:
-        request_data["task_id"] = task_id
-    if workflow_run_id:
-        request_data["workflow_run_id"] = workflow_run_id
-    if workflow_permanent_id:
-        request_data["workflow_permanent_id"] = workflow_permanent_id
-    signed_data = generate_skyvern_webhook_signature(payload=request_data, api_key=api_key)
-    try:
-        json_resp = await aiohttp_post(url=url, data=request_data, headers=signed_data.headers, raise_exception=False)
-    except Exception as e:
-        LOG.error("Failed to get verification code from url", exc_info=True)
-        raise FailedToGetTOTPVerificationCode(
-            task_id=task_id,
-            workflow_run_id=workflow_run_id,
-            workflow_id=workflow_permanent_id,
-            totp_verification_url=url,
-            reason=str(e),
-        )
-    if not json_resp:
-        return None
-    return json_resp.get("verification_code", None)
-
-
-async def _get_verification_code_from_db(
-    organization_id: str,
-    totp_identifier: str,
-    task_id: str | None = None,
-    workflow_id: str | None = None,
-    workflow_run_id: str | None = None,
-) -> str | None:
-    totp_codes = await app.DATABASE.get_otp_codes(
-        organization_id=organization_id, totp_identifier=totp_identifier, otp_type=OTPType.TOTP
-    )
-    for totp_code in totp_codes:
-        if totp_code.workflow_run_id and workflow_run_id and totp_code.workflow_run_id != workflow_run_id:
-            continue
-        if totp_code.workflow_id and workflow_id and totp_code.workflow_id != workflow_id:
-            continue
-        if totp_code.task_id and totp_code.task_id != task_id:
-            continue
-        if totp_code.expired_at and totp_code.expired_at < datetime.utcnow():
-            continue
-        return totp_code.code
-    return None

skyvern/forge/sdk/core/aiohttp_helper.py

@@ -130,6 +130,7 @@ async def aiohttp_get_text(
 async def aiohttp_post(
     url: str,
     data: dict[str, Any] | None = None,
+    str_data: str | None = None,
     headers: dict[str, str] | None = None,
     cookies: dict[str, str] | None = None,
     retry: int = 0,
@@ -145,6 +146,7 @@ async def aiohttp_post(
                 async with session.post(
                     url,
                     # TODO: make sure to test this out
+                    data=str_data,
                     json=data,
                     headers=headers,
                     cookies=cookies,

skyvern/services/otp_service.py

@@ -126,7 +126,9 @@ async def _get_otp_value_from_url(
         api_key=api_key,
     )
     try:
-        json_resp = await aiohttp_post(url=url, data=request_data, headers=signed_data.headers, raise_exception=False)
+        json_resp = await aiohttp_post(
+            url=url, str_data=signed_data.signed_payload, headers=signed_data.headers, raise_exception=False
+        )
     except Exception as e:
         LOG.error("Failed to get otp value from url", exc_info=True)
         raise FailedToGetTOTPVerificationCode(