store totp_identifier to credentials with fallback for login runs (#4154)

This commit is contained in:
Marc Kelechava
2025-12-01 16:19:37 -08:00
committed by GitHub
parent acce1c869d
commit 7100b7e004
6 changed files with 49 additions and 16 deletions

View File

@@ -173,6 +173,7 @@ class WorkflowRunContext:
self._aws_client = aws_client
self.organization_id: str | None = None
self.include_secrets_in_templates: bool = False
self.credential_totp_identifiers: dict[str, str] = {}
def get_parameter(self, key: str) -> Parameter:
return self.parameters[key]
@@ -295,6 +296,10 @@ class WorkflowRunContext:
credential_item = await credential_service.get_credential_item(db_credential)
credential = credential_item.credential
credential_totp_identifier = getattr(credential, "totp_identifier", None)
if credential_totp_identifier:
self.credential_totp_identifiers[parameter.key] = credential_totp_identifier
self.parameters[parameter.key] = parameter
self.values[parameter.key] = {
"context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)",
@@ -319,6 +324,9 @@ class WorkflowRunContext:
self.secrets[totp_secret_value] = parse_totp_secret(credential.totp)
self.values[parameter.key]["totp"] = totp_secret_id
def get_credential_totp_identifier(self, parameter_key: str) -> str | None:
return self.credential_totp_identifiers.get(parameter_key)
async def register_secret_workflow_parameter_value(
self,
parameter: WorkflowParameter,

View File

@@ -588,19 +588,22 @@ class BaseTaskBlock(Block):
)
self.url = task_url_parameter_value
if (
self.totp_identifier
and workflow_run_context.has_parameter(self.totp_identifier)
and workflow_run_context.has_value(self.totp_identifier)
):
totp_identifier_parameter_value = workflow_run_context.get_value(self.totp_identifier)
if totp_identifier_parameter_value:
LOG.info(
"TOTP identifier is parameterized, using parameter value",
totp_identifier_parameter_value=totp_identifier_parameter_value,
totp_identifier_parameter_key=self.totp_identifier,
)
self.totp_identifier = totp_identifier_parameter_value
if self.totp_identifier:
if workflow_run_context.has_parameter(self.totp_identifier) and workflow_run_context.has_value(
self.totp_identifier
):
totp_identifier_parameter_value = workflow_run_context.get_value(self.totp_identifier)
if totp_identifier_parameter_value:
self.totp_identifier = totp_identifier_parameter_value
else:
for parameter in self.get_all_parameters(workflow_run_id):
parameter_key = getattr(parameter, "key", None)
if not parameter_key:
continue
credential_totp_identifier = workflow_run_context.get_credential_totp_identifier(parameter_key)
if credential_totp_identifier:
self.totp_identifier = credential_totp_identifier
break
if self.download_suffix and workflow_run_context.has_parameter(self.download_suffix):
download_suffix_parameter_value = workflow_run_context.get_value(self.download_suffix)