vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Azure Vault credential support (#3394)

Browse Source
This commit is contained in:
stenn930
2025-09-12 11:01:57 -06:00
committed by GitHub
parent c876566c57
commit 8df506660e
23 changed files with 624 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
skyvern/forge/sdk/workflow/context_manager.py
View File

@@ -20,12 +20,13 @@ from skyvern.forge.sdk.schemas.credentials import PasswordCredential
from skyvern.forge.sdk.schemas.organizations import Organization
from skyvern.forge.sdk.schemas.tasks import TaskStatus
from skyvern.forge.sdk.services.bitwarden import BitwardenConstants, BitwardenService
from skyvern.forge.sdk.services.credentials import OnePasswordConstants, parse_totp_secret
from skyvern.forge.sdk.services.credentials import AzureVaultConstants, OnePasswordConstants, parse_totp_secret
from skyvern.forge.sdk.workflow.exceptions import OutputParameterKeyCollisionError
from skyvern.forge.sdk.workflow.models.parameter import (
PARAMETER_TYPE,
AWSSecretParameter,
AzureSecretParameter,
AzureVaultCredentialParameter,
BitwardenCreditCardDataParameter,
BitwardenLoginCredentialParameter,
BitwardenSensitiveInformationParameter,
@@ -55,7 +56,7 @@ class WorkflowRunContext:
async def init(
cls,
aws_client: AsyncAWSClient,
azure_client: AsyncAzureClient | None,
azure_client: AsyncAzureClient,
organization: Organization,
workflow_parameter_tuples: list[tuple[WorkflowParameter, "WorkflowRunParameter"]],
workflow_output_parameters: list[OutputParameter],
@@ -96,28 +97,30 @@ class WorkflowRunContext:
for label, value in block_outputs.items():
workflow_run_context.values[f"{label}_output"] = value
for secrete_parameter in secret_parameters:
if isinstance(secrete_parameter, AWSSecretParameter):
await workflow_run_context.register_aws_secret_parameter_value(secrete_parameter)
elif isinstance(secrete_parameter, AzureSecretParameter):
await workflow_run_context.register_azure_secret_parameter_value(secrete_parameter)
elif isinstance(secrete_parameter, CredentialParameter):
await workflow_run_context.register_credential_parameter_value(secrete_parameter, organization)
elif isinstance(secrete_parameter, OnePasswordCredentialParameter):
for secret_parameter in secret_parameters:
if isinstance(secret_parameter, AWSSecretParameter):
await workflow_run_context.register_aws_secret_parameter_value(secret_parameter)
elif isinstance(secret_parameter, AzureSecretParameter):
await workflow_run_context.register_azure_secret_parameter_value(secret_parameter)
elif isinstance(secret_parameter, CredentialParameter):
await workflow_run_context.register_credential_parameter_value(secret_parameter, organization)
elif isinstance(secret_parameter, OnePasswordCredentialParameter):
await workflow_run_context.register_onepassword_credential_parameter_value(
secrete_parameter, organization
secret_parameter, organization
)
elif isinstance(secrete_parameter, BitwardenLoginCredentialParameter):
elif isinstance(secret_parameter, AzureVaultCredentialParameter):
await workflow_run_context.register_azure_vault_credential_parameter_value(secret_parameter)
elif isinstance(secret_parameter, BitwardenLoginCredentialParameter):
await workflow_run_context.register_bitwarden_login_credential_parameter_value(
secrete_parameter, organization
secret_parameter, organization
)
elif isinstance(secrete_parameter, BitwardenCreditCardDataParameter):
elif isinstance(secret_parameter, BitwardenCreditCardDataParameter):
await workflow_run_context.register_bitwarden_credit_card_data_parameter_value(
secrete_parameter, organization
secret_parameter, organization
)
elif isinstance(secrete_parameter, BitwardenSensitiveInformationParameter):
elif isinstance(secret_parameter, BitwardenSensitiveInformationParameter):
await workflow_run_context.register_bitwarden_sensitive_information_parameter_value(
secrete_parameter, organization
secret_parameter, organization
)
for context_parameter in context_parameters:
@@ -128,7 +131,7 @@ class WorkflowRunContext:
return workflow_run_context
def __init__(self, aws_client: AsyncAWSClient, azure_client: AsyncAzureClient | None) -> None:
def __init__(self, aws_client: AsyncAWSClient, azure_client: AsyncAzureClient) -> None:
self.blocks_metadata: dict[str, BlockMetadata] = {}
self.parameters: dict[str, PARAMETER_TYPE] = {}
self.values: dict[str, Any] = {}
@@ -343,9 +346,6 @@ class WorkflowRunContext:
# If the parameter is an Azure secret, fetch the secret value and store it in the secrets dict
# The value of the parameter will be the random secret id with format `secret_<uuid>`.
# We'll replace the random secret id with the actual secret value when we need to use it.
if self._azure_client is None:
LOG.error("Azure client not initialized, cannot register Azure secret parameter value")
raise ValueError("Azure client not initialized")
secret_value = await self._azure_client.get_secret(parameter.azure_key)
if secret_value is not None:
random_secret_id = self.generate_random_secret_id()
@@ -491,6 +491,48 @@ class WorkflowRunContext:
LOG.error(f"Failed to get secret from Bitwarden. Error: {e}")
raise e
async def register_azure_vault_credential_parameter_value(self, parameter: AzureVaultCredentialParameter) -> None:
vault_name = self._resolve_parameter_value(parameter.vault_name)
if not vault_name:
raise ValueError("Azure Vault Name is missing")
username_key = self._resolve_parameter_value(parameter.username_key)
if not username_key:
raise ValueError("Azure Username Key is missing")
password_key = self._resolve_parameter_value(parameter.password_key)
if not password_key:
raise ValueError("Azure Password Key is missing")
totp_secret_key = self._resolve_parameter_value(parameter.totp_secret_key)
secret_login = await self._azure_client.get_secret(username_key, vault_name)
secret_password = await self._azure_client.get_secret(password_key, vault_name)
if totp_secret_key:
totp_secret = await self._azure_client.get_secret(totp_secret_key, vault_name)
else:
totp_secret = None
if secret_login is not None and secret_password is not None:
random_secret_id = self.generate_random_secret_id()
# login secret
username_secret_id = f"{random_secret_id}_username"
self.secrets[username_secret_id] = secret_login
# password secret
password_secret_id = f"{random_secret_id}_password"
self.secrets[password_secret_id] = secret_password
self.values[parameter.key] = {
"context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)",
"username": username_secret_id,
"password": password_secret_id,
}
self.parameters[parameter.key] = parameter
if totp_secret:
totp_secret_id = f"{random_secret_id}_totp"
self.secrets[totp_secret_id] = AzureVaultConstants.TOTP
totp_secret_value = self.totp_secret_value_key(totp_secret_id)
self.secrets[totp_secret_value] = parse_totp_secret(totp_secret)
self.values[parameter.key]["totp"] = totp_secret_id
async def register_bitwarden_sensitive_information_parameter_value(
self,
parameter: BitwardenSensitiveInformationParameter,
@@ -856,7 +898,7 @@ class WorkflowRunContext:
class WorkflowContextManager:
aws_client: AsyncAWSClient
azure_client: AsyncAzureClient | None
azure_client: AsyncAzureClient
workflow_run_contexts: dict[str, WorkflowRunContext]
parameters: dict[str, PARAMETER_TYPE]
@@ -865,12 +907,10 @@ class WorkflowContextManager:
def __init__(self) -> None:
self.aws_client = AsyncAWSClient()
self.azure_client = None
if settings.AZURE_STORAGE_ACCOUNT_NAME and settings.AZURE_STORAGE_ACCOUNT_KEY:
self.azure_client = AsyncAzureClient(
account_name=settings.AZURE_STORAGE_ACCOUNT_NAME,
account_key=settings.AZURE_STORAGE_ACCOUNT_KEY,
)
self.azure_client = AsyncAzureClient(
storage_account_name=settings.AZURE_STORAGE_ACCOUNT_NAME,
storage_account_key=settings.AZURE_STORAGE_ACCOUNT_KEY,
)
self.workflow_run_contexts = {}
def _validate_workflow_run_context(self, workflow_run_id: str) -> None:

4
skyvern/forge/sdk/workflow/models/block.py
View File

@@ -2062,8 +2062,8 @@ class FileUploadBlock(Block):
or self.azure_storage_account_key
)
azure_client = AsyncAzureClient(
account_name=actual_azure_storage_account_name or "",
account_key=actual_azure_storage_account_key or "",
storage_account_name=actual_azure_storage_account_name,
storage_account_key=actual_azure_storage_account_key,
)
for file_path in files_to_upload:
blob_name = Path(file_path).name

17
skyvern/forge/sdk/workflow/models/parameter.py
View File

@@ -19,6 +19,7 @@ class ParameterType(StrEnum):
BITWARDEN_SENSITIVE_INFORMATION = "bitwarden_sensitive_information"
BITWARDEN_CREDIT_CARD_DATA = "bitwarden_credit_card_data"
ONEPASSWORD = "onepassword"
AZURE_VAULT_CREDENTIAL = "azure_vault_credential"
OUTPUT = "output"
CREDENTIAL = "credential"
AZURE_SECRET = "azure_secret"
@@ -154,6 +155,21 @@ class OnePasswordCredentialParameter(Parameter):
deleted_at: datetime | None = None
class AzureVaultCredentialParameter(Parameter):
parameter_type: Literal[ParameterType.AZURE_VAULT_CREDENTIAL] = ParameterType.AZURE_VAULT_CREDENTIAL
azure_vault_credential_parameter_id: str
workflow_id: str
vault_name: str
username_key: str
password_key: str
totp_secret_key: str | None = None
created_at: datetime
modified_at: datetime
deleted_at: datetime | None = None
class WorkflowParameterType(StrEnum):
STRING = "string"
INTEGER = "integer"
@@ -232,6 +248,7 @@ ParameterSubclasses = Union[
BitwardenSensitiveInformationParameter,
BitwardenCreditCardDataParameter,
OnePasswordCredentialParameter,
AzureVaultCredentialParameter,
OutputParameter,
CredentialParameter,
]

33
skyvern/forge/sdk/workflow/service.py
View File

@@ -32,6 +32,7 @@ from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.security import generate_skyvern_webhook_headers
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
from skyvern.forge.sdk.db.enums import TaskType
from skyvern.forge.sdk.db.models import AzureVaultCredentialParameterModel
from skyvern.forge.sdk.models import Step, StepStatus
from skyvern.forge.sdk.schemas.files import FileInfo
from skyvern.forge.sdk.schemas.organizations import Organization
@@ -74,6 +75,7 @@ from skyvern.forge.sdk.workflow.models.parameter import (
PARAMETER_TYPE,
RESERVED_PARAMETER_KEYS,
AWSSecretParameter,
AzureVaultCredentialParameter,
BitwardenCreditCardDataParameter,
BitwardenLoginCredentialParameter,
BitwardenSensitiveInformationParameter,
@@ -308,6 +310,7 @@ class WorkflowService:
BitwardenCreditCardDataParameter,
BitwardenSensitiveInformationParameter,
OnePasswordCredentialParameter,
AzureVaultCredentialParameter,
CredentialParameter,
),
)
@@ -1119,6 +1122,26 @@ class WorkflowService:
description=description,
)
async def create_azure_vault_credential_parameter(
self,
workflow_id: str,
key: str,
vault_name: str,
username_key: str,
password_key: str,
totp_secret_key: str | None = None,
description: str | None = None,
) -> AzureVaultCredentialParameterModel:
return await app.DATABASE.create_azure_vault_credential_parameter(
workflow_id=workflow_id,
key=key,
vault_name=vault_name,
username_key=username_key,
password_key=password_key,
totp_secret_key=totp_secret_key,
description=description,
)
async def create_bitwarden_sensitive_information_parameter(
self,
workflow_id: str,
@@ -1771,6 +1794,16 @@ class WorkflowService:
vault_id=parameter.vault_id,
item_id=parameter.item_id,
)
elif parameter.parameter_type == ParameterType.AZURE_VAULT_CREDENTIAL:
parameters[parameter.key] = await self.create_azure_vault_credential_parameter(
workflow_id=workflow.workflow_id,
key=parameter.key,
description=parameter.description,
vault_name=parameter.vault_name,
username_key=parameter.username_key,
password_key=parameter.password_key,
totp_secret_key=parameter.totp_secret_key,
)
elif parameter.parameter_type == ParameterType.BITWARDEN_LOGIN_CREDENTIAL:
if not parameter.bitwarden_collection_id and not parameter.bitwarden_item_id:
raise WorkflowParameterMissingRequiredValue(