Implement BitwardenLoginCredentialParameter (#151)

This commit is contained in:
Kerem Yilmaz 2024-04-03 16:01:03 -07:00 committed by GitHub
parent 999eda9b5d
commit 1d1e29b813
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 392 additions and 8 deletions

View file

@ -0,0 +1,67 @@
"""Create bitwarden credential parameter table
Revision ID: 4630ab8c198e
Revises: ffe2f57bd288
Create Date: 2024-04-03 22:57:03.231654+00:00
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "4630ab8c198e"
down_revision: Union[str, None] = "ffe2f57bd288"
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(
"bitwarden_login_credential_parameters",
sa.Column("bitwarden_login_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("bitwarden_client_id_aws_secret_key", sa.String(), nullable=False),
sa.Column("bitwarden_client_secret_aws_secret_key", sa.String(), nullable=False),
sa.Column("bitwarden_master_password_aws_secret_key", sa.String(), nullable=False),
sa.Column("url_parameter_key", sa.String(), nullable=False),
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.ForeignKeyConstraint(
["workflow_id"],
["workflows.workflow_id"],
),
sa.PrimaryKeyConstraint("bitwarden_login_credential_parameter_id"),
)
op.create_index(
op.f("ix_bitwarden_login_credential_parameters_bitwarden_login_credential_parameter_id"),
"bitwarden_login_credential_parameters",
["bitwarden_login_credential_parameter_id"],
unique=False,
)
op.create_index(
op.f("ix_bitwarden_login_credential_parameters_workflow_id"),
"bitwarden_login_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_bitwarden_login_credential_parameters_workflow_id"), table_name="bitwarden_login_credential_parameters"
)
op.drop_index(
op.f("ix_bitwarden_login_credential_parameters_bitwarden_login_credential_parameter_id"),
table_name="bitwarden_login_credential_parameters",
)
op.drop_table("bitwarden_login_credential_parameters")
# ### end Alembic commands ###

View file

@ -185,3 +185,28 @@ class DownloadFileMaxSizeExceeded(SkyvernException):
def __init__(self, max_size: int) -> None:
self.max_size = max_size
super().__init__(f"Download file size exceeded the maximum allowed size of {max_size} MB.")
class BitwardenBaseError(SkyvernException):
def __init__(self, message: str) -> None:
super().__init__(f"Bitwarden error: {message}")
class BitwardenLoginError(BitwardenBaseError):
def __init__(self, message: str) -> None:
super().__init__(f"Error logging in to Bitwarden: {message}")
class BitwardenUnlockError(BitwardenBaseError):
def __init__(self, message: str) -> None:
super().__init__(f"Error unlocking Bitwarden: {message}")
class BitwardenListItemsError(BitwardenBaseError):
def __init__(self, message: str) -> None:
super().__init__(f"Error listing items in Bitwarden: {message}")
class BitwardenLogoutError(BitwardenBaseError):
def __init__(self, message: str) -> None:
super().__init__(f"Error logging out of Bitwarden: {message}")

View file

@ -13,6 +13,7 @@ from skyvern.forge.sdk.db.exceptions import NotFoundError
from skyvern.forge.sdk.db.models import (
ArtifactModel,
AWSSecretParameterModel,
BitwardenLoginCredentialParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
@ -28,6 +29,7 @@ from skyvern.forge.sdk.db.utils import (
_custom_json_serializer,
convert_to_artifact,
convert_to_aws_secret_parameter,
convert_to_bitwarden_login_credential_parameter,
convert_to_organization,
convert_to_organization_auth_token,
convert_to_output_parameter,
@ -43,6 +45,7 @@ from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step,
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
BitwardenLoginCredentialParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
@ -844,6 +847,31 @@ class AgentDB:
await session.refresh(aws_secret_parameter)
return convert_to_aws_secret_parameter(aws_secret_parameter)
async def create_bitwarden_login_credential_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,
url_parameter_key: str,
key: str,
description: str | None = None,
) -> BitwardenLoginCredentialParameter:
async with self.Session() as session:
bitwarden_login_credential_parameter = BitwardenLoginCredentialParameterModel(
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,
url_parameter_key=url_parameter_key,
key=key,
description=description,
)
session.add(bitwarden_login_credential_parameter)
await session.commit()
await session.refresh(bitwarden_login_credential_parameter)
return convert_to_bitwarden_login_credential_parameter(bitwarden_login_credential_parameter)
async def create_output_parameter(
self,
workflow_id: str,

View file

@ -38,6 +38,7 @@ WORKFLOW_RUN_PREFIX = "wr"
WORKFLOW_PARAMETER_PREFIX = "wp"
AWS_SECRET_PARAMETER_PREFIX = "asp"
OUTPUT_PARAMETER_PREFIX = "op"
BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
def generate_workflow_id() -> str:
@ -65,6 +66,11 @@ def generate_output_parameter_id() -> str:
return f"{OUTPUT_PARAMETER_PREFIX}_{int_id}"
def generate_bitwarden_login_credential_parameter_id() -> str:
int_id = generate_id()
return f"{BITWARDEN_LOGIN_CREDENTIAL_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

@ -8,6 +8,7 @@ from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
from skyvern.forge.sdk.db.id import (
generate_artifact_id,
generate_aws_secret_parameter_id,
generate_bitwarden_login_credential_parameter_id,
generate_org_id,
generate_organization_auth_token_id,
generate_output_parameter_id,
@ -177,6 +178,24 @@ class AWSSecretParameterModel(Base):
deleted_at = Column(DateTime, nullable=True)
class BitwardenLoginCredentialParameterModel(Base):
__tablename__ = "bitwarden_login_credential_parameters"
bitwarden_login_credential_parameter_id = Column(
String, primary_key=True, index=True, default=generate_bitwarden_login_credential_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)
url_parameter_key = Column(String, 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

@ -9,6 +9,7 @@ from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
from skyvern.forge.sdk.db.models import (
ArtifactModel,
AWSSecretParameterModel,
BitwardenLoginCredentialParameterModel,
OrganizationAuthTokenModel,
OrganizationModel,
OutputParameterModel,
@ -24,6 +25,7 @@ from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step,
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import (
AWSSecretParameter,
BitwardenLoginCredentialParameter,
OutputParameter,
WorkflowParameter,
WorkflowParameterType,
@ -211,6 +213,30 @@ def convert_to_aws_secret_parameter(
)
def convert_to_bitwarden_login_credential_parameter(
bitwarden_login_credential_parameter_model: BitwardenLoginCredentialParameterModel, debug_enabled: bool = False
) -> BitwardenLoginCredentialParameter:
if debug_enabled:
LOG.debug(
"Converting BitwardenLoginCredentialParameterModel to BitwardenLoginCredentialParameter",
bitwarden_login_credential_parameter_id=bitwarden_login_credential_parameter_model.bitwarden_login_credential_parameter_id,
)
return BitwardenLoginCredentialParameter(
bitwarden_login_credential_parameter_id=bitwarden_login_credential_parameter_model.bitwarden_login_credential_parameter_id,
workflow_id=bitwarden_login_credential_parameter_model.workflow_id,
key=bitwarden_login_credential_parameter_model.key,
description=bitwarden_login_credential_parameter_model.description,
bitwarden_client_id_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_client_id_aws_secret_key,
bitwarden_client_secret_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_client_secret_aws_secret_key,
bitwarden_master_password_aws_secret_key=bitwarden_login_credential_parameter_model.bitwarden_master_password_aws_secret_key,
url_parameter_key=bitwarden_login_credential_parameter_model.url_parameter_key,
created_at=bitwarden_login_credential_parameter_model.created_at,
modified_at=bitwarden_login_credential_parameter_model.modified_at,
deleted_at=bitwarden_login_credential_parameter_model.deleted_at,
)
def convert_to_output_parameter(
output_parameter_model: OutputParameterModel, debug_enabled: bool = False
) -> OutputParameter:

View file

@ -0,0 +1,97 @@
import json
import os
import subprocess
import structlog
from skyvern.exceptions import BitwardenListItemsError, BitwardenLoginError, BitwardenLogoutError, BitwardenUnlockError
LOG = structlog.get_logger()
class BitwardenService:
@staticmethod
def run_command(command: list[str], additional_env: dict[str, str] | None = None) -> subprocess.CompletedProcess:
"""
Run a CLI command with the specified additional environment variables and return the result.
"""
env = os.environ.copy() # Copy the current environment
# Make sure node isn't returning warnings. Warnings are sent through stderr and we raise exceptions on stderr.
env["NODE_NO_WARNINGS"] = "1"
if additional_env:
env.update(additional_env) # Update with any additional environment variables
return subprocess.run(command, capture_output=True, text=True, env=env)
@staticmethod
def get_secret_value_from_url(
client_id: str,
client_secret: str,
master_password: str,
url: str,
) -> dict[str, str]:
"""
Get the secret value from the Bitwarden CLI.
"""
# Step 1: Set up environment variables and log in
env = {"BW_CLIENTID": client_id, "BW_CLIENTSECRET": client_secret, "BW_PASSWORD": master_password}
login_command = ["bw", "login", "--apikey"]
login_result = BitwardenService.run_command(login_command, env)
# Print both stdout and stderr for debugging
if login_result.stderr:
raise BitwardenLoginError(login_result.stderr)
# Step 2: Unlock the vault
unlock_command = ["bw", "unlock", "--passwordenv", "BW_PASSWORD"]
unlock_result = BitwardenService.run_command(unlock_command, env)
if unlock_result.stderr:
raise BitwardenUnlockError(unlock_result.stderr)
# Extract session key
try:
session_key = unlock_result.stdout.split('"')[1]
except IndexError:
raise BitwardenUnlockError("Unable to extract session key.")
if not session_key:
raise BitwardenUnlockError("Session key is empty.")
# Step 3: Retrieve the items
list_command = ["bw", "list", "items", "--url", url, "--session", session_key]
items_result = BitwardenService.run_command(list_command)
if items_result.stderr:
raise BitwardenListItemsError(items_result.stderr)
# Parse the items and extract credentials
try:
items = json.loads(items_result.stdout)
except json.JSONDecodeError:
raise BitwardenListItemsError("Failed to parse items JSON. Output: " + items_result.stdout)
if not items:
raise BitwardenListItemsError("No items found in Bitwarden.")
credentials = [
{"username": item["login"]["username"], "password": item["login"]["password"]}
for item in items
if "login" in item
]
# Step 4: Log out
BitwardenService.logout()
# Todo: Handle multiple credentials, for now just return the last one
return credentials[-1] if credentials else {}
@staticmethod
def logout() -> None:
"""
Log out of the Bitwarden CLI.
"""
logout_command = ["bw", "logout"]
logout_result = BitwardenService.run_command(logout_command)
if logout_result.stderr:
raise BitwardenLogoutError(logout_result.stderr)

View file

@ -3,8 +3,9 @@ from typing import TYPE_CHECKING, Any
import structlog
from skyvern.exceptions import WorkflowRunContextNotInitialized
from skyvern.exceptions import BitwardenBaseError, WorkflowRunContextNotInitialized
from skyvern.forge.sdk.api.aws import AsyncAWSClient
from skyvern.forge.sdk.services.bitwarden import BitwardenService
from skyvern.forge.sdk.workflow.exceptions import OutputParameterKeyCollisionError
from skyvern.forge.sdk.workflow.models.parameter import (
PARAMETER_TYPE,
@ -113,6 +114,46 @@ class WorkflowRunContext:
random_secret_id = self.generate_random_secret_id()
self.secrets[random_secret_id] = secret_value
self.values[parameter.key] = random_secret_id
elif parameter.parameter_type == ParameterType.BITWARDEN_LOGIN_CREDENTIAL:
try:
# Get the Bitwarden login credentials from AWS secrets
client_id = await aws_client.get_secret(parameter.bitwarden_client_id_aws_secret_key)
client_secret = await aws_client.get_secret(parameter.bitwarden_client_secret_aws_secret_key)
master_password = await aws_client.get_secret(parameter.bitwarden_master_password_aws_secret_key)
except Exception as e:
LOG.error(f"Failed to get Bitwarden login credentials from AWS secrets. Error: {e}")
raise e
if self.has_parameter(parameter.url_parameter_key) and self.has_value(parameter.url_parameter_key):
url = self.values[parameter.url_parameter_key]
else:
LOG.error(f"URL parameter {parameter.url_parameter_key} not found or has no value")
raise ValueError(f"URL parameter for Bitwarden login credentials not found or has no value")
try:
secret_credentials = BitwardenService.get_secret_value_from_url(
client_id,
client_secret,
master_password,
url,
)
if secret_credentials:
random_secret_id = self.generate_random_secret_id()
# username secret
username_secret_id = f"{random_secret_id}_username"
self.secrets[username_secret_id] = secret_credentials["username"]
# password secret
password_secret_id = f"{random_secret_id}_password"
self.secrets[password_secret_id] = secret_credentials["password"]
self.values[parameter.key] = {
"username": username_secret_id,
"password": password_secret_id,
}
except BitwardenBaseError as e:
BitwardenService.logout()
LOG.error(f"Failed to get secret from Bitwarden. Error: {e}")
raise e
elif parameter.parameter_type == ParameterType.CONTEXT:
# ContextParameter values will be set within the blocks
return
@ -133,6 +174,9 @@ class WorkflowRunContext:
aws_client: AsyncAWSClient,
parameters: list[PARAMETER_TYPE],
) -> None:
# BitwardenLoginCredentialParameter should be processed last since it requires the URL parameter to be set
parameters.sort(key=lambda x: x.parameter_type != ParameterType.BITWARDEN_LOGIN_CREDENTIAL)
for parameter in parameters:
if parameter.key in self.parameters:
LOG.debug(f"Parameter {parameter.key} already registered, skipping")

View file

@ -275,7 +275,7 @@ class ForLoopBlock(Block):
return context_parameters
def get_loop_over_parameter_values(self, workflow_run_context: WorkflowRunContext) -> list[Any]:
if isinstance(self.loop_over, WorkflowParameter):
if isinstance(self.loop_over, WorkflowParameter) or isinstance(self.loop_over, OutputParameter):
parameter_value = workflow_run_context.get_value(self.loop_over.key)
if isinstance(parameter_value, list):
return parameter_value
@ -284,7 +284,6 @@ class ForLoopBlock(Block):
return [parameter_value]
else:
# TODO (kerem): Implement this for context parameters
# TODO (kerem): Implement this for output parameters
raise NotImplementedError
async def execute(self, workflow_run_id: str, **kwargs: dict) -> OutputParameter | None:

View file

@ -11,6 +11,7 @@ class ParameterType(StrEnum):
WORKFLOW = "workflow"
CONTEXT = "context"
AWS_SECRET = "aws_secret"
BITWARDEN_LOGIN_CREDENTIAL = "bitwarden_login_credential"
OUTPUT = "output"
@ -37,6 +38,23 @@ class AWSSecretParameter(Parameter):
deleted_at: datetime | None = None
class BitwardenLoginCredentialParameter(Parameter):
parameter_type: Literal[ParameterType.BITWARDEN_LOGIN_CREDENTIAL] = ParameterType.BITWARDEN_LOGIN_CREDENTIAL
# parameter fields
bitwarden_login_credential_parameter_id: str
workflow_id: str
# bitwarden cli required fields
bitwarden_client_id_aws_secret_key: str
bitwarden_client_secret_aws_secret_key: str
bitwarden_master_password_aws_secret_key: str
# url to request the login credentials from bitwarden
url_parameter_key: str
created_at: datetime
modified_at: datetime
deleted_at: datetime | None = None
class WorkflowParameterType(StrEnum):
STRING = "string"
INTEGER = "integer"
@ -92,5 +110,7 @@ class OutputParameter(Parameter):
deleted_at: datetime | None = None
ParameterSubclasses = Union[WorkflowParameter, ContextParameter, AWSSecretParameter, OutputParameter]
ParameterSubclasses = Union[
WorkflowParameter, ContextParameter, AWSSecretParameter, BitwardenLoginCredentialParameter, OutputParameter
]
PARAMETER_TYPE = Annotated[ParameterSubclasses, Field(discriminator="parameter_type")]

View file

@ -22,6 +22,21 @@ class AWSSecretParameterYAML(ParameterYAML):
aws_key: str
class BitwardenLoginCredentialParameterYAML(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"
# This pattern already works in block.py but since the ParameterType is not defined in this file, mypy is not able
# to infer the type of the parameter_type attribute.
parameter_type: Literal[ParameterType.BITWARDEN_LOGIN_CREDENTIAL] = ParameterType.BITWARDEN_LOGIN_CREDENTIAL # type: ignore
# bitwarden cli required fields
bitwarden_client_id_aws_secret_key: str
bitwarden_client_secret_aws_secret_key: str
bitwarden_master_password_aws_secret_key: str
# parameter key for the url to request the login credentials from bitwarden
url_parameter_key: str
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"
@ -80,7 +95,7 @@ class ForLoopBlockYAML(BlockYAML):
block_type: Literal[BlockType.FOR_LOOP] = BlockType.FOR_LOOP # type: ignore
loop_over_parameter_key: str
loop_block: BlockYAML
loop_block: "BLOCK_YAML_SUBCLASSES"
class CodeBlockYAML(BlockYAML):
@ -135,7 +150,13 @@ class SendEmailBlockYAML(BlockYAML):
file_attachments: list[str] | None = None
PARAMETER_YAML_SUBCLASSES = AWSSecretParameterYAML | WorkflowParameterYAML | ContextParameterYAML | OutputParameterYAML
PARAMETER_YAML_SUBCLASSES = (
AWSSecretParameterYAML
| BitwardenLoginCredentialParameterYAML
| WorkflowParameterYAML
| ContextParameterYAML
| OutputParameterYAML
)
PARAMETER_YAML_TYPES = Annotated[PARAMETER_YAML_SUBCLASSES, Field(discriminator="parameter_type")]
BLOCK_YAML_SUBCLASSES = (

View file

@ -351,6 +351,26 @@ class WorkflowService:
workflow_id=workflow_id, aws_key=aws_key, key=key, description=description
)
async def create_bitwarden_login_credential_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,
url_parameter_key: str,
key: str,
description: str | None = None,
) -> Parameter:
return await app.DATABASE.create_bitwarden_login_credential_parameter(
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,
url_parameter_key=url_parameter_key,
key=key,
description=description,
)
async def create_output_parameter(
self, workflow_id: str, key: str, description: str | None = None
) -> OutputParameter:
@ -643,6 +663,16 @@ class WorkflowService:
key=parameter.key,
description=parameter.description,
)
elif parameter.parameter_type == ParameterType.BITWARDEN_LOGIN_CREDENTIAL:
parameters[parameter.key] = await self.create_bitwarden_login_credential_parameter(
workflow_id=workflow.workflow_id,
bitwarden_client_id_aws_secret_key=parameter.bitwarden_client_id_aws_secret_key,
bitwarden_client_secret_aws_secret_key=parameter.bitwarden_client_secret_aws_secret_key,
bitwarden_master_password_aws_secret_key=parameter.bitwarden_master_password_aws_secret_key,
url_parameter_key=parameter.url_parameter_key,
key=parameter.key,
description=parameter.description,
)
elif parameter.parameter_type == ParameterType.WORKFLOW:
parameters[parameter.key] = await self.create_workflow_parameter(
workflow_id=workflow.workflow_id,
@ -708,10 +738,12 @@ class WorkflowService:
max_retries=block_yaml.max_retries,
)
elif block_yaml.block_type == BlockType.FOR_LOOP:
loop_block = await WorkflowService.block_yaml_to_block(block_yaml.loop_block, parameters)
loop_over_parameter = parameters[block_yaml.loop_over_parameter_key]
return ForLoopBlock(
label=block_yaml.label,
loop_over_parameter_key=parameters[block_yaml.loop_over_parameter_key],
loop_block=WorkflowService.block_yaml_to_block(block_yaml.loop_block, parameters),
loop_over=loop_over_parameter,
loop_block=loop_block,
output_parameter=output_parameter,
)
elif block_yaml.block_type == BlockType.CODE: