diff --git a/alembic/versions/2025_09_12_1622-f78486c3f895_create_azure_vault_credential_.py b/alembic/versions/2025_09_12_1622-f78486c3f895_create_azure_vault_credential_.py new file mode 100644 index 00000000..9e494484 --- /dev/null +++ b/alembic/versions/2025_09_12_1622-f78486c3f895_create_azure_vault_credential_.py @@ -0,0 +1,54 @@ +"""create azure_vault_credential_parameters table + +Revision ID: f78486c3f895 +Revises: d15d85a702a4 +Create Date: 2025-09-12 16:22:11.988864+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "f78486c3f895" +down_revision: Union[str, None] = "d15d85a702a4" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "azure_vault_credential_parameters", + sa.Column("azure_vault_credential_parameter_id", sa.String(), nullable=False), + sa.Column("workflow_id", sa.String(), nullable=False), + sa.Column("key", sa.String(), nullable=False), + sa.Column("description", sa.String(), nullable=True), + sa.Column("vault_name", sa.String(), nullable=False), + sa.Column("username_key", sa.String(), nullable=False), + sa.Column("password_key", sa.String(), nullable=False), + sa.Column("totp_secret_key", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("modified_at", sa.DateTime(), nullable=False), + sa.Column("deleted_at", sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint("azure_vault_credential_parameter_id"), + ) + op.create_index( + op.f("ix_azure_vault_credential_parameters_workflow_id"), + "azure_vault_credential_parameters", + ["workflow_id"], + unique=False, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index( + op.f("ix_azure_vault_credential_parameters_workflow_id"), table_name="azure_vault_credential_parameters" + ) + op.drop_table("azure_vault_credential_parameters") + # ### end Alembic commands ### diff --git a/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx b/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx index 067acbd7..5bfb3a12 100644 --- a/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx @@ -50,6 +50,7 @@ import { ContextParameterYAML, CredentialParameterYAML, OnePasswordCredentialParameterYAML, + AzureVaultCredentialParameterYAML, ParameterYAML, WorkflowParameterYAML, } from "../types/workflowYamlTypes"; @@ -65,6 +66,7 @@ import { parameterIsSkyvernCredential, parameterIsOnePasswordCredential, parameterIsBitwardenCredential, + parameterIsAzureVaultCredential, } from "./types"; import "./reactFlowOverrideStyles.css"; import { @@ -90,6 +92,7 @@ function convertToParametersYAML( | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | CredentialParameterYAML > { return parameters @@ -103,6 +106,7 @@ function convertToParametersYAML( | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | CredentialParameterYAML | undefined => { if (parameter.parameterType === WorkflowEditorParameterTypes.Workflow) { @@ -191,6 +195,16 @@ function convertToParametersYAML( vault_id: parameter.vaultId, item_id: parameter.itemId, }; + } else if (parameterIsAzureVaultCredential(parameter)) { + return { + parameter_type: WorkflowParameterTypes.Azure_Vault_Credential, + key: parameter.key, + description: parameter.description || null, + vault_name: parameter.vaultName, + username_key: parameter.usernameKey, + password_key: parameter.passwordKey, + totp_secret_key: parameter.totpSecretKey, + }; } } return undefined; @@ -205,6 +219,7 @@ function convertToParametersYAML( | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | CredentialParameterYAML | undefined, ): param is @@ -214,6 +229,7 @@ function convertToParametersYAML( | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | CredentialParameterYAML => param !== undefined, ); } diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx index 2acf2e01..ad686a61 100644 --- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx @@ -1,4 +1,3 @@ -import { SwitchBar } from "@/components/SwitchBar"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; @@ -81,7 +80,7 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) { >(undefined); const [credentialType, setCredentialType] = useState< - "bitwarden" | "skyvern" | "onepassword" + "bitwarden" | "skyvern" | "onepassword" | "azurevault" >("skyvern"); const [vaultId, setVaultId] = useState(""); const [itemId, setItemId] = useState(""); @@ -93,6 +92,11 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) { const [credentialId, setCredentialId] = useState(""); + const [azureVaultName, setAzureVaultName] = useState(""); + const [azureUsernameKey, setAzureUsernameKey] = useState(""); + const [azurePasswordKey, setAzurePasswordKey] = useState(""); + const [azureTotpSecretKey, setAzureTotpKey] = useState(""); + return ( @@ -199,19 +203,37 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) { )} {type === "credential" && ( - { - setCredentialType( - value as "bitwarden" | "skyvern" | "onepassword", - ); - }} - options={[ - { label: "Skyvern", value: "skyvern" }, - { label: "Bitwarden", value: "bitwarden" }, - { label: "1Password", value: "onepassword" }, - ]} - /> + <> +
+ + +
+ )} {type === "credential" && credentialType === "bitwarden" && ( <> @@ -260,6 +282,41 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) { )} + {type === "credential" && credentialType === "azurevault" && ( + <> +
+ + setAzureVaultName(e.target.value)} + /> +
+
+ + setAzureUsernameKey(e.target.value)} + /> +
+
+ + setAzurePasswordKey(e.target.value)} + /> +
+
+ + setAzureTotpKey(e.target.value)} + /> +
+ + )} {type === "context" && (
@@ -427,6 +484,31 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) { description, }); } + if (type === "credential" && credentialType === "azurevault") { + if ( + azureVaultName.trim() === "" || + azureUsernameKey.trim() === "" || + azurePasswordKey.trim() === "" + ) { + toast({ + variant: "destructive", + title: "Failed to add parameter", + description: + "Azure Vault Name, Username Key and Password Key are required", + }); + return; + } + onSave({ + key, + parameterType: "credential", + vaultName: azureVaultName, + usernameKey: azureUsernameKey, + passwordKey: azurePasswordKey, + totpSecretKey: + azureTotpSecretKey === "" ? null : azureTotpSecretKey, + description: description, + }); + } if (type === "secret" || type === "creditCardData") { if (!bitwardenCollectionId) { toast({ diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx index d4b2f993..0034cc20 100644 --- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx @@ -1,4 +1,3 @@ -import { SwitchBar } from "@/components/SwitchBar"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; @@ -28,6 +27,7 @@ import { parameterIsSkyvernCredential, parameterIsOnePasswordCredential, ParametersState, + parameterIsAzureVaultCredential, } from "../types"; import { getDefaultValueForParameterType } from "../workflowEditorUtils"; import { validateBitwardenLoginCredential } from "./util"; @@ -82,14 +82,19 @@ function WorkflowParameterEditPanel({ const isOnePasswordCredential = initialValues.parameterType === "onepassword" && parameterIsOnePasswordCredential(initialValues); + const isAzureVaultCredential = + initialValues.parameterType === "credential" && + parameterIsAzureVaultCredential(initialValues); const [credentialType, setCredentialType] = useState< - "bitwarden" | "skyvern" | "onepassword" + "bitwarden" | "skyvern" | "onepassword" | "azurevault" >( isBitwardenCredential ? "bitwarden" : isOnePasswordCredential ? "onepassword" - : "skyvern", + : isAzureVaultCredential + ? "azurevault" + : "skyvern", ); const [urlParameterKey, setUrlParameterKey] = useState( isBitwardenCredential ? initialValues.urlParameterKey ?? "" : "", @@ -163,6 +168,19 @@ function WorkflowParameterEditPanel({ const [bitwardenLoginCredentialItemId, setBitwardenLoginCredentialItemId] = useState(isBitwardenCredential ? initialValues.itemId ?? "" : ""); + const [azureVaultName, setAzureVaultName] = useState( + isAzureVaultCredential ? initialValues.vaultName : "", + ); + const [azureUsernameKey, setAzureUsernameKey] = useState( + isAzureVaultCredential ? initialValues.usernameKey : "", + ); + const [azurePasswordKey, setAzurePasswordKey] = useState( + isAzureVaultCredential ? initialValues.passwordKey : "", + ); + const [azureTotpSecretKey, setAzureTotpKey] = useState( + isAzureVaultCredential ? initialValues.totpSecretKey ?? "" : "", + ); + return ( @@ -269,19 +287,37 @@ function WorkflowParameterEditPanel({ )} {type === "credential" && ( - { - setCredentialType( - value as "bitwarden" | "skyvern" | "onepassword", - ); - }} - options={[ - { label: "Skyvern", value: "skyvern" }, - { label: "Bitwarden", value: "bitwarden" }, - { label: "1Password", value: "onepassword" }, - ]} - /> + <> +
+ + +
+ )} {type === "credential" && credentialType === "bitwarden" && ( <> @@ -330,6 +366,41 @@ function WorkflowParameterEditPanel({
)} + {type === "credential" && credentialType === "azurevault" && ( + <> +
+ + setAzureVaultName(e.target.value)} + /> +
+
+ + setAzureUsernameKey(e.target.value)} + /> +
+
+ + setAzurePasswordKey(e.target.value)} + /> +
+
+ + setAzureTotpKey(e.target.value)} + /> +
+ + )} {type === "context" && (
@@ -484,6 +555,31 @@ function WorkflowParameterEditPanel({ description, }); } + if (type === "credential" && credentialType === "azurevault") { + if ( + azureVaultName.trim() === "" || + azureUsernameKey.trim() === "" || + azurePasswordKey.trim() === "" + ) { + toast({ + variant: "destructive", + title: "Failed to add parameter", + description: + "Azure Vault Name, Username Key and Password Key are required", + }); + return; + } + onSave({ + key, + parameterType: "credential", + vaultName: azureVaultName, + usernameKey: azureUsernameKey, + passwordKey: azurePasswordKey, + totpSecretKey: + azureTotpSecretKey === "" ? null : azureTotpSecretKey, + description: description, + }); + } if (type === "secret" || type === "creditCardData") { if (!collectionId) { toast({ diff --git a/skyvern-frontend/src/routes/workflows/editor/types.ts b/skyvern-frontend/src/routes/workflows/editor/types.ts index 2b8c6d92..96a2d17a 100644 --- a/skyvern-frontend/src/routes/workflows/editor/types.ts +++ b/skyvern-frontend/src/routes/workflows/editor/types.ts @@ -24,6 +24,16 @@ export type OnePasswordCredential = { itemId: string; }; +export type AzureVaultCredential = { + key: string; + description?: string | null; + parameterType: "credential"; + vaultName: string; + usernameKey: string; + passwordKey: string; + totpSecretKey: string | null; +}; + export function parameterIsBitwardenCredential( parameter: CredentialParameterState, ): parameter is BitwardenLoginCredential { @@ -42,10 +52,21 @@ export function parameterIsOnePasswordCredential( return "vaultId" in parameter && "itemId" in parameter; } +export function parameterIsAzureVaultCredential( + parameter: CredentialParameterState, +): parameter is AzureVaultCredential { + return ( + "vaultName" in parameter && + "usernameKey" in parameter && + "passwordKey" in parameter + ); +} + export type CredentialParameterState = | BitwardenLoginCredential | SkyvernCredential - | OnePasswordCredential; + | OnePasswordCredential + | AzureVaultCredential; export type ParametersState = Array< | { diff --git a/skyvern-frontend/src/routes/workflows/editor/utils.ts b/skyvern-frontend/src/routes/workflows/editor/utils.ts index ed78925a..3ed1dbfb 100644 --- a/skyvern-frontend/src/routes/workflows/editor/utils.ts +++ b/skyvern-frontend/src/routes/workflows/editor/utils.ts @@ -79,6 +79,19 @@ const getInitialParameters = (workflow: WorkflowApiResponse) => { itemId: parameter.item_id, description: parameter.description, }; + } else if ( + parameter.parameter_type === + WorkflowParameterTypes.Azure_Vault_Credential + ) { + return { + key: parameter.key, + parameterType: WorkflowEditorParameterTypes.Credential, + vaultName: parameter.vault_name, + usernameKey: parameter.username_key, + passwordKey: parameter.password_key, + totpSecretKey: parameter.totp_secret_key, + description: parameter.description, + }; } else if ( parameter.parameter_type === WorkflowParameterTypes.Bitwarden_Login_Credential diff --git a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts index c1bbf866..2ca36e87 100644 --- a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts +++ b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts @@ -1805,6 +1805,16 @@ function convertParametersToParameterYAML( item_id: parameter.item_id, }; } + case WorkflowParameterTypes.Azure_Vault_Credential: { + return { + ...base, + parameter_type: WorkflowParameterTypes.Azure_Vault_Credential, + vault_name: parameter.vault_name, + username_key: parameter.username_key, + password_key: parameter.password_key, + totp_secret_key: parameter.totp_secret_key, + }; + } } }) .filter(Boolean); diff --git a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts index 822c2d2d..c2be2de4 100644 --- a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts +++ b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts @@ -72,6 +72,19 @@ export type OnePasswordCredentialParameter = WorkflowParameterBase & { deleted_at: string | null; }; +export type AzureVaultCredentialParameter = WorkflowParameterBase & { + parameter_type: "azure_vault_credential"; + workflow_id: string; + azure_vault_credential_parameter_id: string; + vault_name: string; + username_key: string; + password_key: string; + totp_secret_key: string | null; + created_at: string; + modified_at: string; + deleted_at: string | null; +}; + export type CredentialParameter = WorkflowParameterBase & { parameter_type: "credential"; workflow_id: string; @@ -130,6 +143,7 @@ export const WorkflowParameterTypes = { Bitwarden_Sensitive_Information: "bitwarden_sensitive_information", Bitwarden_Credit_Card_Data: "bitwarden_credit_card_data", OnePassword: "onepassword", + Azure_Vault_Credential: "azure_vault_credential", Credential: "credential", } as const; @@ -145,6 +159,7 @@ export function isDisplayedInWorkflowEditor( | BitwardenLoginCredentialParameter | BitwardenSensitiveInformationParameter | OnePasswordCredentialParameter + | AzureVaultCredentialParameter | CredentialParameter { return ( parameter.parameter_type === WorkflowParameterTypes.Workflow || @@ -156,6 +171,8 @@ export function isDisplayedInWorkflowEditor( parameter.parameter_type === WorkflowParameterTypes.Bitwarden_Credit_Card_Data || parameter.parameter_type === WorkflowParameterTypes.OnePassword || + parameter.parameter_type === + WorkflowParameterTypes.Azure_Vault_Credential || parameter.parameter_type === WorkflowParameterTypes.Credential ); } @@ -168,6 +185,7 @@ export type Parameter = | BitwardenSensitiveInformationParameter | BitwardenCreditCardDataParameter | OnePasswordCredentialParameter + | AzureVaultCredentialParameter | AWSSecretParameter | CredentialParameter; diff --git a/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts b/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts index 33426b5c..72e0db9f 100644 --- a/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts +++ b/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts @@ -31,6 +31,7 @@ export type ParameterYAML = | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | ContextParameterYAML | OutputParameterYAML | CredentialParameterYAML; @@ -90,6 +91,14 @@ export type OnePasswordCredentialParameterYAML = ParameterYAMLBase & { item_id: string; }; +export type AzureVaultCredentialParameterYAML = ParameterYAMLBase & { + parameter_type: "azure_vault_credential"; + vault_name: string; + username_key: string; + password_key: string; + totp_secret_key: string | null; +}; + export type ContextParameterYAML = ParameterYAMLBase & { parameter_type: "context"; source_parameter_key: string; diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index 714ab0d4..a3711a23 100644 --- a/skyvern/exceptions.py +++ b/skyvern/exceptions.py @@ -757,3 +757,13 @@ class NoElementFound(SkyvernException): class OutputParameterNotFound(SkyvernException): def __init__(self, block_label: str, workflow_permanent_id: str) -> None: super().__init__(f"Output parameter for {block_label} not found in workflow {workflow_permanent_id}") + + +class AzureBaseError(SkyvernException): + def __init__(self, message: str) -> None: + super().__init__(f"Azure error: {message}") + + +class AzureConfigurationError(AzureBaseError): + def __init__(self, message: str) -> None: + super().__init__(f"Error in Azure configuration: {message}") diff --git a/skyvern/forge/sdk/api/azure.py b/skyvern/forge/sdk/api/azure.py index 40eb0a25..d4073523 100644 --- a/skyvern/forge/sdk/api/azure.py +++ b/skyvern/forge/sdk/api/azure.py @@ -3,26 +3,37 @@ from azure.identity.aio import DefaultAzureCredential from azure.keyvault.secrets.aio import SecretClient from azure.storage.blob.aio import BlobServiceClient +from skyvern.exceptions import AzureConfigurationError + LOG = structlog.get_logger() class AsyncAzureClient: - def __init__(self, account_name: str, account_key: str): - self.account_name = account_name - self.account_key = account_key - self.blob_service_client = BlobServiceClient( - account_url=f"https://{account_name}.blob.core.windows.net", - credential=account_key, - ) + def __init__(self, storage_account_name: str | None, storage_account_key: str | None): + self.storage_account_name = storage_account_name + self.storage_account_key = storage_account_key + + if storage_account_name and storage_account_key: + self.blob_service_client = BlobServiceClient( + account_url=f"https://{storage_account_name}.blob.core.windows.net", + credential=storage_account_key, + ) + else: + self.blob_service_client = None + self.credential = DefaultAzureCredential() - async def get_secret(self, secret_name: str) -> str | None: + async def get_secret(self, secret_name: str, vault_name: str | None = None) -> str | None: + vault_subdomain = vault_name or self.storage_account_name + if not vault_subdomain: + raise AzureConfigurationError("Missing vault") + try: # Azure Key Vault URL format: https://.vault.azure.net # Assuming the secret_name is actually the Key Vault URL and the secret name # This needs to be clarified or passed as separate parameters # For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings. - key_vault_url = f"https://{self.account_name}.vault.azure.net" # Placeholder, adjust as needed + key_vault_url = f"https://{vault_subdomain}.vault.azure.net" # Placeholder, adjust as needed secret_client = SecretClient(vault_url=key_vault_url, credential=self.credential) secret = await secret_client.get_secret(secret_name) return secret.value @@ -33,6 +44,9 @@ class AsyncAzureClient: await self.credential.close() async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None: + if not self.blob_service_client: + raise AzureConfigurationError("Storage is not configured") + try: container_client = self.blob_service_client.get_container_client(container_name) # Create the container if it doesn't exist diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 95df1724..67ac3054 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -17,6 +17,7 @@ from skyvern.forge.sdk.db.models import ( AISuggestionModel, ArtifactModel, AWSSecretParameterModel, + AzureVaultCredentialParameterModel, BitwardenCreditCardDataParameterModel, BitwardenLoginCredentialParameterModel, BitwardenSensitiveInformationParameterModel, @@ -88,6 +89,7 @@ from skyvern.forge.sdk.schemas.totp_codes import TOTPCode from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock from skyvern.forge.sdk.workflow.models.parameter import ( AWSSecretParameter, + AzureVaultCredentialParameter, BitwardenCreditCardDataParameter, BitwardenLoginCredentialParameter, BitwardenSensitiveInformationParameter, @@ -2107,6 +2109,43 @@ class AgentDB: deleted_at=parameter.deleted_at, ) + 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, + ) -> AzureVaultCredentialParameter: + async with self.Session() as session: + parameter = AzureVaultCredentialParameterModel( + workflow_id=workflow_id, + key=key, + description=description, + vault_name=vault_name, + username_key=username_key, + password_key=password_key, + totp_secret_key=totp_secret_key, + ) + session.add(parameter) + await session.commit() + await session.refresh(parameter) + return AzureVaultCredentialParameter( + azure_vault_credential_parameter_id=parameter.azure_vault_credential_parameter_id, + workflow_id=parameter.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, + created_at=parameter.created_at, + modified_at=parameter.modified_at, + deleted_at=parameter.deleted_at, + ) + async def get_workflow_run_output_parameters(self, workflow_run_id: str) -> list[WorkflowRunOutputParameter]: try: async with self.Session() as session: diff --git a/skyvern/forge/sdk/db/id.py b/skyvern/forge/sdk/db/id.py index 1bbabbdb..48a92f75 100644 --- a/skyvern/forge/sdk/db/id.py +++ b/skyvern/forge/sdk/db/id.py @@ -35,6 +35,7 @@ BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd" BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc" BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi" CREDENTIAL_ONEPASSWORD_PARAMETER_PREFIX = "opp" +CREDENTIAL_AZURE_VAULT_PARAMETER_PREFIX = "azcp" CREDENTIAL_PARAMETER_PREFIX = "cp" CREDENTIAL_PREFIX = "cred" DEBUG_SESSION_PREFIX = "ds" @@ -123,6 +124,11 @@ def generate_onepassword_credential_parameter_id() -> str: return f"{CREDENTIAL_ONEPASSWORD_PARAMETER_PREFIX}_{int_id}" +def generate_azure_vault_credential_parameter_id() -> str: + int_id = generate_id() + return f"{CREDENTIAL_AZURE_VAULT_PARAMETER_PREFIX}_{int_id}" + + def generate_organization_auth_token_id() -> str: int_id = generate_id() return f"{ORGANIZATION_AUTH_TOKEN_PREFIX}_{int_id}" diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 1b1bbba1..6409935e 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -24,6 +24,7 @@ from skyvern.forge.sdk.db.id import ( generate_ai_suggestion_id, generate_artifact_id, generate_aws_secret_parameter_id, + generate_azure_vault_credential_parameter_id, generate_bitwarden_credit_card_data_parameter_id, generate_bitwarden_login_credential_parameter_id, generate_bitwarden_sensitive_information_parameter_id, @@ -467,6 +468,30 @@ class OnePasswordCredentialParameterModel(Base): deleted_at = Column(DateTime, nullable=True) +class AzureVaultCredentialParameterModel(Base): + __tablename__ = "azure_vault_credential_parameters" + + azure_vault_credential_parameter_id = Column( + String, primary_key=True, default=generate_azure_vault_credential_parameter_id + ) + workflow_id = Column(String, index=True, nullable=False) + key = Column(String, nullable=False) + description = Column(String, nullable=True) + vault_name = Column(String, nullable=False) + username_key = Column(String, nullable=False) + password_key = Column(String, nullable=False) + totp_secret_key = Column(String, nullable=True) + + created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) + modified_at = Column( + DateTime, + default=datetime.datetime.utcnow, + onupdate=datetime.datetime.utcnow, + nullable=False, + ) + deleted_at = Column(DateTime, nullable=True) + + class WorkflowRunParameterModel(Base): __tablename__ = "workflow_run_parameters" diff --git a/skyvern/forge/sdk/routes/run_blocks.py b/skyvern/forge/sdk/routes/run_blocks.py index e1eafe3b..7b68b0a4 100644 --- a/skyvern/forge/sdk/routes/run_blocks.py +++ b/skyvern/forge/sdk/routes/run_blocks.py @@ -20,6 +20,7 @@ from skyvern.forge.sdk.workflow.models.workflow import WorkflowRequestBody from skyvern.schemas.run_blocks import CredentialType, LoginRequest from skyvern.schemas.runs import ProxyLocation, RunType, WorkflowRunRequest, WorkflowRunResponse from skyvern.schemas.workflows import ( + AzureVaultCredentialParameterYAML, BitwardenLoginCredentialParameterYAML, LoginBlockYAML, OnePasswordCredentialParameterYAML, @@ -63,7 +64,7 @@ If login is completed, you're successful.""" } ], }, - description="Log in to a website using either credential stored in Skyvern, Bitwarden or 1Password", + description="Log in to a website using either credential stored in Skyvern, Bitwarden, 1Password, or Azure Vault", summary="Login Task", ) async def login( @@ -130,6 +131,28 @@ async def login( item_id=login_request.onepassword_item_id, ) ] + elif login_request.credential_type == CredentialType.azure_vault: + if not login_request.azure_vault_name: + raise HTTPException( + status_code=400, detail="azure_vault_name is required to login with Azure Vault credential" + ) + if not login_request.azure_vault_username_key: + raise HTTPException( + status_code=400, detail="azure_vault_username_key is required to login with Azure Vault credential" + ) + if not login_request.azure_vault_password_key: + raise HTTPException( + status_code=400, detail="azure_vault_password_key is required to login with Azure Vault credential" + ) + yaml_parameters = [ + AzureVaultCredentialParameterYAML( + key=parameter_key, + vault_name=login_request.azure_vault_name, + username_key=login_request.azure_vault_username_key, + password_key=login_request.azure_vault_password_key, + totp_secret_key=login_request.azure_vault_totp_secret_key, + ) + ] login_block_yaml = LoginBlockYAML( label=label, diff --git a/skyvern/forge/sdk/services/credentials.py b/skyvern/forge/sdk/services/credentials.py index 330afb65..7b4382a5 100644 --- a/skyvern/forge/sdk/services/credentials.py +++ b/skyvern/forge/sdk/services/credentials.py @@ -14,6 +14,12 @@ class OnePasswordConstants(StrEnum): TOTP = "OP_TOTP" # Special value to indicate a TOTP code +class AzureVaultConstants(StrEnum): + """Constants for Azure Vault integration.""" + + TOTP = "AZ_TOTP" # Special value to indicate a TOTP code + + def parse_totp_secret(totp_secret: str) -> str: if not totp_secret: return "" diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 5d376948..290ce809 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -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_`. # 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: diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 54231ded..6193c6b4 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -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 diff --git a/skyvern/forge/sdk/workflow/models/parameter.py b/skyvern/forge/sdk/workflow/models/parameter.py index 5ce7be2c..db071cbc 100644 --- a/skyvern/forge/sdk/workflow/models/parameter.py +++ b/skyvern/forge/sdk/workflow/models/parameter.py @@ -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, ] diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 03791a14..8699cef8 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -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( diff --git a/skyvern/schemas/run_blocks.py b/skyvern/schemas/run_blocks.py index 391f3989..0db9315b 100644 --- a/skyvern/schemas/run_blocks.py +++ b/skyvern/schemas/run_blocks.py @@ -9,6 +9,7 @@ class CredentialType(StrEnum): skyvern = "skyvern" bitwarden = "bitwarden" onepassword = "1password" + azure_vault = "azure_vault" class LoginRequest(BaseModel): @@ -51,3 +52,9 @@ class LoginRequest(BaseModel): # 1Password credential onepassword_vault_id: str | None = Field(default=None, description="1Password vault ID") onepassword_item_id: str | None = Field(default=None, description="1Password item ID") + + # Azure Vault credential + azure_vault_name: str | None = Field(default=None, description="Azure Vault Name") + azure_vault_username_key: str | None = Field(default=None, description="Azure Vault username key") + azure_vault_password_key: str | None = Field(default=None, description="Azure Vault password key") + azure_vault_totp_secret_key: str | None = Field(default=None, description="Azure Vault TOTP secret key") diff --git a/skyvern/schemas/workflows.py b/skyvern/schemas/workflows.py index 8721218b..cab3535d 100644 --- a/skyvern/schemas/workflows.py +++ b/skyvern/schemas/workflows.py @@ -150,6 +150,14 @@ class OnePasswordCredentialParameterYAML(ParameterYAML): item_id: str +class AzureVaultCredentialParameterYAML(ParameterYAML): + parameter_type: Literal[ParameterType.AZURE_VAULT_CREDENTIAL] = ParameterType.AZURE_VAULT_CREDENTIAL # type: ignore + vault_name: str + username_key: str + password_key: str + totp_secret_key: str | None = None + + class WorkflowParameterYAML(ParameterYAML): # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error: # Parameter 1 of Literal[...] cannot be of type "Any" @@ -454,6 +462,7 @@ PARAMETER_YAML_SUBCLASSES = ( | BitwardenSensitiveInformationParameterYAML | BitwardenCreditCardDataParameterYAML | OnePasswordCredentialParameterYAML + | AzureVaultCredentialParameterYAML | WorkflowParameterYAML | ContextParameterYAML | OutputParameterYAML diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 01ffe0ed..3dc131d4 100644 --- a/skyvern/webeye/actions/handler.py +++ b/skyvern/webeye/actions/handler.py @@ -67,7 +67,7 @@ from skyvern.forge.sdk.core.skyvern_context import ensure_context from skyvern.forge.sdk.models import Step from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.services.bitwarden import BitwardenConstants -from skyvern.forge.sdk.services.credentials import OnePasswordConstants +from skyvern.forge.sdk.services.credentials import AzureVaultConstants, OnePasswordConstants from skyvern.forge.sdk.trace import TraceManager from skyvern.services.task_v1_service import is_cua_task from skyvern.utils.prompt_engine import ( @@ -961,7 +961,9 @@ async def handle_input_text_action( if text is None: return [ActionFailure(FailedToFetchSecret())] - is_totp_value = text == BitwardenConstants.TOTP or text == OnePasswordConstants.TOTP + is_totp_value = ( + text == BitwardenConstants.TOTP or text == OnePasswordConstants.TOTP or text == AzureVaultConstants.TOTP + ) is_secret_value = text != action.text # dynamically validate the attr, since it could change into enabled after the previous actions