Timeout and retry mechanism for Bitwarden CLI (#670)

This commit is contained in:
Kerem Yilmaz 2024-08-02 12:53:05 -07:00 committed by GitHub
parent 7fd07ece2f
commit e4fd825497
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 119 additions and 8 deletions

View file

@ -62,6 +62,12 @@ class Settings(BaseSettings):
# Workflow constant parameters # Workflow constant parameters
WORKFLOW_DOWNLOAD_DIRECTORY_PARAMETER_KEY: str = "SKYVERN_DOWNLOAD_DIRECTORY" WORKFLOW_DOWNLOAD_DIRECTORY_PARAMETER_KEY: str = "SKYVERN_DOWNLOAD_DIRECTORY"
#####################
# Bitwarden Configs #
#####################
BITWARDEN_TIMEOUT_SECONDS: int = 60
BITWARDEN_MAX_RETRIES: int = 1
##################### #####################
# LLM Configuration # # LLM Configuration #
##################### #####################

View file

@ -278,6 +278,11 @@ class BitwardenLogoutError(BitwardenBaseError):
super().__init__(f"Error logging out of Bitwarden: {message}") super().__init__(f"Error logging out of Bitwarden: {message}")
class BitwardenSyncError(BitwardenBaseError):
def __init__(self, message: str) -> None:
super().__init__(f"Error syncing Bitwarden: {message}")
class UnknownElementTreeFormat(SkyvernException): class UnknownElementTreeFormat(SkyvernException):
def __init__(self, fmt: str) -> None: def __init__(self, fmt: str) -> None:
super().__init__(f"Unknown element tree format {fmt}") super().__init__(f"Unknown element tree format {fmt}")

View file

@ -1,3 +1,4 @@
import asyncio
import json import json
import os import os
import re import re
@ -7,10 +8,12 @@ from enum import StrEnum
import structlog import structlog
import tldextract import tldextract
from skyvern.config import settings
from skyvern.exceptions import ( from skyvern.exceptions import (
BitwardenListItemsError, BitwardenListItemsError,
BitwardenLoginError, BitwardenLoginError,
BitwardenLogoutError, BitwardenLogoutError,
BitwardenSyncError,
BitwardenTOTPError, BitwardenTOTPError,
BitwardenUnlockError, BitwardenUnlockError,
) )
@ -68,7 +71,47 @@ class BitwardenService:
return None return None
@staticmethod @staticmethod
def get_secret_value_from_url( async def get_secret_value_from_url(
client_id: str,
client_secret: str,
master_password: str,
url: str,
collection_id: str | None = None,
remaining_retries: int = settings.BITWARDEN_MAX_RETRIES,
fail_reasons: list[str] = [],
) -> dict[str, str]:
"""
Get the secret value from the Bitwarden CLI.
"""
try:
async with asyncio.timeout(settings.BITWARDEN_TIMEOUT_SECONDS):
return await BitwardenService._get_secret_value_from_url(
client_id=client_id,
client_secret=client_secret,
master_password=master_password,
url=url,
collection_id=collection_id,
)
except Exception as e:
if remaining_retries <= 0:
raise BitwardenListItemsError(
f"Bitwarden CLI failed after all retry attempts. Fail reasons: {fail_reasons}"
)
remaining_retries -= 1
LOG.info("Retrying to get secret value from Bitwarden", remaining_retries=remaining_retries)
return await BitwardenService.get_secret_value_from_url(
client_id=client_id,
client_secret=client_secret,
master_password=master_password,
url=url,
collection_id=collection_id,
remaining_retries=remaining_retries,
fail_reasons=fail_reasons + [f"{type(e).__name__}: {str(e)}"],
)
@staticmethod
async def _get_secret_value_from_url(
client_id: str, client_id: str,
client_secret: str, client_secret: str,
master_password: str, master_password: str,
@ -80,6 +123,7 @@ class BitwardenService:
""" """
try: try:
BitwardenService.login(client_id, client_secret) BitwardenService.login(client_id, client_secret)
BitwardenService.sync()
session_key = BitwardenService.unlock(master_password) session_key = BitwardenService.unlock(master_password)
# Extract the domain from the URL and search for items in Bitwarden with that domain # Extract the domain from the URL and search for items in Bitwarden with that domain
@ -151,7 +195,50 @@ class BitwardenService:
BitwardenService.logout() BitwardenService.logout()
@staticmethod @staticmethod
def get_sensitive_information_from_identity( async def get_sensitive_information_from_identity(
client_id: str,
client_secret: str,
master_password: str,
collection_id: str,
identity_key: str,
identity_fields: list[str],
remaining_retries: int = settings.BITWARDEN_MAX_RETRIES,
fail_reasons: list[str] = [],
) -> dict[str, str]:
"""
Get the secret value from the Bitwarden CLI.
"""
try:
async with asyncio.timeout(settings.BITWARDEN_TIMEOUT_SECONDS):
return await BitwardenService._get_sensitive_information_from_identity(
client_id=client_id,
client_secret=client_secret,
master_password=master_password,
collection_id=collection_id,
identity_key=identity_key,
identity_fields=identity_fields,
)
except Exception as e:
if remaining_retries <= 0:
raise BitwardenListItemsError(
f"Bitwarden CLI failed after all retry attempts. Fail reasons: {fail_reasons}"
)
remaining_retries -= 1
LOG.info("Retrying to get sensitive information from Bitwarden", remaining_retries=remaining_retries)
return await BitwardenService.get_sensitive_information_from_identity(
client_id=client_id,
client_secret=client_secret,
master_password=master_password,
collection_id=collection_id,
identity_key=identity_key,
identity_fields=identity_fields,
remaining_retries=remaining_retries,
fail_reasons=fail_reasons + [f"{type(e).__name__}: {str(e)}"],
)
@staticmethod
async def _get_sensitive_information_from_identity(
client_id: str, client_id: str,
client_secret: str, client_secret: str,
master_password: str, master_password: str,
@ -164,6 +251,7 @@ class BitwardenService:
""" """
try: try:
BitwardenService.login(client_id, client_secret) BitwardenService.login(client_id, client_secret)
BitwardenService.sync()
session_key = BitwardenService.unlock(master_password) session_key = BitwardenService.unlock(master_password)
# Step 3: Retrieve the items # Step 3: Retrieve the items
@ -267,6 +355,18 @@ class BitwardenService:
return session_key return session_key
@staticmethod
def sync() -> None:
"""
Sync the Bitwarden CLI.
"""
sync_command = ["bw", "sync"]
LOG.info("Bitwarden CLI sync started")
sync_result = BitwardenService.run_command(sync_command)
LOG.info("Bitwarden CLI sync completed")
if sync_result.stderr:
raise BitwardenSyncError(sync_result.stderr)
@staticmethod @staticmethod
def logout() -> None: def logout() -> None:
""" """
@ -274,5 +374,5 @@ class BitwardenService:
""" """
logout_command = ["bw", "logout"] logout_command = ["bw", "logout"]
logout_result = BitwardenService.run_command(logout_command) logout_result = BitwardenService.run_command(logout_command)
if logout_result.stderr: if logout_result.stderr and "You are not logged in." not in logout_result.stderr:
raise BitwardenLogoutError(logout_result.stderr) raise BitwardenLogoutError(logout_result.stderr)

View file

@ -97,11 +97,11 @@ class WorkflowRunContext:
return self.secrets.get(secret_id_or_value) return self.secrets.get(secret_id_or_value)
return None return None
def get_secrets_from_password_manager(self) -> dict[str, Any]: async def get_secrets_from_password_manager(self) -> dict[str, Any]:
""" """
Get the secrets from the password manager. The secrets dict will contain the actual secret values. Get the secrets from the password manager. The secrets dict will contain the actual secret values.
""" """
secret_credentials = BitwardenService.get_secret_value_from_url( secret_credentials = await BitwardenService.get_secret_value_from_url(
url=self.secrets[BitwardenConstants.URL], url=self.secrets[BitwardenConstants.URL],
client_secret=self.secrets[BitwardenConstants.CLIENT_SECRET], client_secret=self.secrets[BitwardenConstants.CLIENT_SECRET],
client_id=self.secrets[BitwardenConstants.CLIENT_ID], client_id=self.secrets[BitwardenConstants.CLIENT_ID],
@ -161,7 +161,7 @@ class WorkflowRunContext:
collection_id = parameter.bitwarden_collection_id collection_id = parameter.bitwarden_collection_id
try: try:
secret_credentials = BitwardenService.get_secret_value_from_url( secret_credentials = await BitwardenService.get_secret_value_from_url(
client_id, client_id,
client_secret, client_secret,
master_password, master_password,
@ -218,7 +218,7 @@ class WorkflowRunContext:
collection_id = self.values[parameter.bitwarden_collection_id] collection_id = self.values[parameter.bitwarden_collection_id]
try: try:
sensitive_values = BitwardenService.get_sensitive_information_from_identity( sensitive_values = await BitwardenService.get_sensitive_information_from_identity(
client_id, client_id,
client_secret, client_secret,
master_password, master_password,

View file

@ -706,7 +706,7 @@ async def get_actual_value_of_parameter_if_secret(task: Task, parameter: str) ->
secret_value = workflow_run_context.get_original_secret_value_or_none(parameter) secret_value = workflow_run_context.get_original_secret_value_or_none(parameter)
if secret_value == BitwardenConstants.TOTP: if secret_value == BitwardenConstants.TOTP:
secrets = workflow_run_context.get_secrets_from_password_manager() secrets = await workflow_run_context.get_secrets_from_password_manager()
secret_value = secrets[BitwardenConstants.TOTP] secret_value = secrets[BitwardenConstants.TOTP]
return secret_value if secret_value is not None else parameter return secret_value if secret_value is not None else parameter