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

feat: encrypt org auth tokens with AES (#3104)

Browse Source
This commit is contained in:
LawyZheng
2025-08-05 12:36:24 +08:00
committed by GitHub
parent 977c9d4f13
commit 02576e5be3
8 changed files with 192 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
skyvern/forge/sdk/db/client.py
View File

@@ -66,6 +66,8 @@ from skyvern.forge.sdk.db.utils import (
convert_to_workflow_run_parameter,
hydrate_action,
)
from skyvern.forge.sdk.encrypt import encryptor
from skyvern.forge.sdk.encrypt.base import EncryptMethod
from skyvern.forge.sdk.log_artifacts import save_workflow_run_logs
from skyvern.forge.sdk.models import Step, StepStatus
from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
@@ -867,7 +869,7 @@ class AgentDB:
.order_by(OrganizationAuthTokenModel.created_at.desc())
)
).first():
return convert_to_organization_auth_token(token)
return await convert_to_organization_auth_token(token)
else:
return None
except SQLAlchemyError:
@@ -893,7 +895,7 @@ class AgentDB:
.order_by(OrganizationAuthTokenModel.created_at.desc())
)
).all()
return [convert_to_organization_auth_token(token) for token in tokens]
return [await convert_to_organization_auth_token(token) for token in tokens]
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
@@ -907,19 +909,27 @@ class AgentDB:
token_type: OrganizationAuthTokenType,
token: str,
valid: bool | None = True,
encrypted_method: EncryptMethod | None = None,
) -> OrganizationAuthToken | None:
try:
encrypted_token = ""
if encrypted_method is not None:
encrypted_token = await encryptor.encrypt(token, encrypted_method)
async with self.Session() as session:
query = (
select(OrganizationAuthTokenModel)
.filter_by(organization_id=organization_id)
.filter_by(token_type=token_type)
.filter_by(token=token)
)
if encrypted_token:
query = query.filter_by(encrypted_token=encrypted_token)
else:
query = query.filter_by(token=token)
if valid is not None:
query = query.filter_by(valid=valid)
if token_obj := (await session.scalars(query)).first():
return convert_to_organization_auth_token(token_obj)
return await convert_to_organization_auth_token(token_obj)
else:
return None
except SQLAlchemyError:
@@ -934,18 +944,28 @@ class AgentDB:
organization_id: str,
token_type: OrganizationAuthTokenType,
token: str,
encrypted_method: EncryptMethod | None = None,
) -> OrganizationAuthToken:
plaintext_token = token
encrypted_token = ""
if encrypted_method is not None:
encrypted_token = await encryptor.encrypt(token, encrypted_method)
plaintext_token = ""
async with self.Session() as session:
auth_token = OrganizationAuthTokenModel(
organization_id=organization_id,
token_type=token_type,
token=token,
token=plaintext_token,
encrypted_token=encrypted_token,
encrypted_method=encrypted_method.value if encrypted_method is not None else "",
)
session.add(auth_token)
await session.commit()
await session.refresh(auth_token)
return convert_to_organization_auth_token(auth_token)
return await convert_to_organization_auth_token(auth_token)
async def get_artifacts_for_task_v2(
self,