Implement BitwardenSensitiveInformationParameter (#589)

This commit is contained in:
Kerem Yilmaz
2024-07-11 09:48:14 -07:00
committed by GitHub
parent 87d6e71768
commit 6f88ae31a0
10 changed files with 422 additions and 41 deletions

View File

@@ -14,6 +14,7 @@ from skyvern.forge.sdk.db.models import (
ArtifactModel,
AWSSecretParameterModel,
BitwardenLoginCredentialParameterModel,
BitwardenSensitiveInformationParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
@@ -31,6 +32,7 @@ from skyvern.forge.sdk.db.utils import (
convert_to_artifact,
convert_to_aws_secret_parameter,
convert_to_bitwarden_login_credential_parameter,
convert_to_bitwarden_sensitive_information_parameter,
convert_to_organization,
convert_to_organization_auth_token,
convert_to_output_parameter,
@@ -48,6 +50,7 @@ from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
BitwardenLoginCredentialParameter,
BitwardenSensitiveInformationParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
@@ -1138,6 +1141,35 @@ class AgentDB:
await session.refresh(bitwarden_login_credential_parameter)
return convert_to_bitwarden_login_credential_parameter(bitwarden_login_credential_parameter)
async def create_bitwarden_sensitive_information_parameter(
self,
workflow_id: str,
bitwarden_client_id_aws_secret_key: str,
bitwarden_client_secret_aws_secret_key: str,
bitwarden_master_password_aws_secret_key: str,
bitwarden_collection_id: str,
bitwarden_identity_key: str,
bitwarden_identity_fields: list[str],
key: str,
description: str | None = None,
) -> BitwardenSensitiveInformationParameter:
async with self.Session() as session:
bitwarden_sensitive_information_parameter = BitwardenSensitiveInformationParameterModel(
workflow_id=workflow_id,
bitwarden_client_id_aws_secret_key=bitwarden_client_id_aws_secret_key,
bitwarden_client_secret_aws_secret_key=bitwarden_client_secret_aws_secret_key,
bitwarden_master_password_aws_secret_key=bitwarden_master_password_aws_secret_key,
bitwarden_collection_id=bitwarden_collection_id,
bitwarden_identity_key=bitwarden_identity_key,
bitwarden_identity_fields=bitwarden_identity_fields,
key=key,
description=description,
)
session.add(bitwarden_sensitive_information_parameter)
await session.commit()
await session.refresh(bitwarden_sensitive_information_parameter)
return convert_to_bitwarden_sensitive_information_parameter(bitwarden_sensitive_information_parameter)
async def create_output_parameter(
self,
workflow_id: str,

View File

@@ -40,6 +40,7 @@ WORKFLOW_PARAMETER_PREFIX = "wp"
AWS_SECRET_PARAMETER_PREFIX = "asp"
OUTPUT_PARAMETER_PREFIX = "op"
BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi"
TASK_GENERATION_PREFIX = "tg"
@@ -78,6 +79,11 @@ def generate_bitwarden_login_credential_parameter_id() -> str:
return f"{BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX}_{int_id}"
def generate_bitwarden_sensitive_information_parameter_id() -> str:
int_id = generate_id()
return f"{BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX}_{int_id}"
def generate_organization_auth_token_id() -> str:
int_id = generate_id()
return f"{ORGANIZATION_AUTH_TOKEN_PREFIX}_{int_id}"

View File

@@ -22,6 +22,7 @@ from skyvern.forge.sdk.db.id import (
generate_artifact_id,
generate_aws_secret_parameter_id,
generate_bitwarden_login_credential_parameter_id,
generate_bitwarden_sensitive_information_parameter_id,
generate_org_id,
generate_organization_auth_token_id,
generate_output_parameter_id,
@@ -293,6 +294,35 @@ class BitwardenLoginCredentialParameterModel(Base):
deleted_at = Column(DateTime, nullable=True)
class BitwardenSensitiveInformationParameterModel(Base):
__tablename__ = "bitwarden_sensitive_information_parameters"
bitwarden_sensitive_information_parameter_id = Column(
String,
primary_key=True,
index=True,
default=generate_bitwarden_sensitive_information_parameter_id,
)
workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
key = Column(String, nullable=False)
description = Column(String, nullable=True)
bitwarden_client_id_aws_secret_key = Column(String, nullable=False)
bitwarden_client_secret_aws_secret_key = Column(String, nullable=False)
bitwarden_master_password_aws_secret_key = Column(String, nullable=False)
bitwarden_collection_id = Column(String, nullable=False)
bitwarden_identity_key = Column(String, nullable=False)
# This is a list of fields to extract from the Bitwarden Identity.
bitwarden_identity_fields = Column(JSON, nullable=False)
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"

View File

@@ -10,6 +10,7 @@ from skyvern.forge.sdk.db.models import (
ArtifactModel,
AWSSecretParameterModel,
BitwardenLoginCredentialParameterModel,
BitwardenSensitiveInformationParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
@@ -26,6 +27,7 @@ from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
BitwardenLoginCredentialParameter,
BitwardenSensitiveInformationParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
@@ -263,6 +265,33 @@ def convert_to_bitwarden_login_credential_parameter(
)
def convert_to_bitwarden_sensitive_information_parameter(
bitwarden_sensitive_information_parameter_model: BitwardenSensitiveInformationParameterModel,
debug_enabled: bool = False,
) -> BitwardenSensitiveInformationParameter:
if debug_enabled:
LOG.debug(
"Converting BitwardenSensitiveInformationParameterModel to BitwardenSensitiveInformationParameter",
bitwarden_sensitive_information_parameter_id=bitwarden_sensitive_information_parameter_model.bitwarden_sensitive_information_parameter_id,
)
return BitwardenSensitiveInformationParameter(
bitwarden_sensitive_information_parameter_id=bitwarden_sensitive_information_parameter_model.bitwarden_sensitive_information_parameter_id,
workflow_id=bitwarden_sensitive_information_parameter_model.workflow_id,
key=bitwarden_sensitive_information_parameter_model.key,
description=bitwarden_sensitive_information_parameter_model.description,
bitwarden_client_id_aws_secret_key=bitwarden_sensitive_information_parameter_model.bitwarden_client_id_aws_secret_key,
bitwarden_client_secret_aws_secret_key=bitwarden_sensitive_information_parameter_model.bitwarden_client_secret_aws_secret_key,
bitwarden_master_password_aws_secret_key=bitwarden_sensitive_information_parameter_model.bitwarden_master_password_aws_secret_key,
bitwarden_collection_id=bitwarden_sensitive_information_parameter_model.bitwarden_collection_id,
bitwarden_identity_key=bitwarden_sensitive_information_parameter_model.bitwarden_identity_key,
bitwarden_identity_fields=bitwarden_sensitive_information_parameter_model.bitwarden_identity_fields,
created_at=bitwarden_sensitive_information_parameter_model.created_at,
modified_at=bitwarden_sensitive_information_parameter_model.modified_at,
deleted_at=bitwarden_sensitive_information_parameter_model.deleted_at,
)
def convert_to_output_parameter(
output_parameter_model: OutputParameterModel, debug_enabled: bool = False
) -> OutputParameter: