diff --git a/skyvern/core/totp.py b/skyvern/core/totp.py deleted file mode 100644 index 273bc1c3..00000000 --- a/skyvern/core/totp.py +++ /dev/null @@ -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 diff --git a/skyvern/forge/sdk/core/aiohttp_helper.py b/skyvern/forge/sdk/core/aiohttp_helper.py index ce5d3ed8..1ce597e6 100644 --- a/skyvern/forge/sdk/core/aiohttp_helper.py +++ b/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, diff --git a/skyvern/services/otp_service.py b/skyvern/services/otp_service.py index b2bff38f..c7a83e30 100644 --- a/skyvern/services/otp_service.py +++ b/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(