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_DOWNLOAD_DIRECTORY_PARAMETER_KEY: str = "SKYVERN_DOWNLOAD_DIRECTORY"
#####################
# Bitwarden Configs #
#####################
BITWARDEN_TIMEOUT_SECONDS: int = 60
BITWARDEN_MAX_RETRIES: int = 1
#####################
# LLM Configuration #
#####################

View file

@ -278,6 +278,11 @@ class BitwardenLogoutError(BitwardenBaseError):
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):
def __init__(self, fmt: str) -> None:
super().__init__(f"Unknown element tree format {fmt}")

View file

@ -1,3 +1,4 @@
import asyncio
import json
import os
import re
@ -7,10 +8,12 @@ from enum import StrEnum
import structlog
import tldextract
from skyvern.config import settings
from skyvern.exceptions import (
BitwardenListItemsError,
BitwardenLoginError,
BitwardenLogoutError,
BitwardenSyncError,
BitwardenTOTPError,
BitwardenUnlockError,
)
@ -68,7 +71,47 @@ class BitwardenService:
return None
@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_secret: str,
master_password: str,
@ -80,6 +123,7 @@ class BitwardenService:
"""
try:
BitwardenService.login(client_id, client_secret)
BitwardenService.sync()
session_key = BitwardenService.unlock(master_password)
# Extract the domain from the URL and search for items in Bitwarden with that domain
@ -151,7 +195,50 @@ class BitwardenService:
BitwardenService.logout()
@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_secret: str,
master_password: str,
@ -164,6 +251,7 @@ class BitwardenService:
"""
try:
BitwardenService.login(client_id, client_secret)
BitwardenService.sync()
session_key = BitwardenService.unlock(master_password)
# Step 3: Retrieve the items
@ -267,6 +355,18 @@ class BitwardenService:
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
def logout() -> None:
"""
@ -274,5 +374,5 @@ class BitwardenService:
"""
logout_command = ["bw", "logout"]
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)

View file

@ -97,11 +97,11 @@ class WorkflowRunContext:
return self.secrets.get(secret_id_or_value)
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.
"""
secret_credentials = BitwardenService.get_secret_value_from_url(
secret_credentials = await BitwardenService.get_secret_value_from_url(
url=self.secrets[BitwardenConstants.URL],
client_secret=self.secrets[BitwardenConstants.CLIENT_SECRET],
client_id=self.secrets[BitwardenConstants.CLIENT_ID],
@ -161,7 +161,7 @@ class WorkflowRunContext:
collection_id = parameter.bitwarden_collection_id
try:
secret_credentials = BitwardenService.get_secret_value_from_url(
secret_credentials = await BitwardenService.get_secret_value_from_url(
client_id,
client_secret,
master_password,
@ -218,7 +218,7 @@ class WorkflowRunContext:
collection_id = self.values[parameter.bitwarden_collection_id]
try:
sensitive_values = BitwardenService.get_sensitive_information_from_identity(
sensitive_values = await BitwardenService.get_sensitive_information_from_identity(
client_id,
client_secret,
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)
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]
return secret_value if secret_value is not None else parameter