Clean up Azure migration (#3895)
Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
parent
b1fe444984
commit
d80d49131c
@@ -21,6 +21,7 @@ from skyvern.forge.sdk.routes.code_samples import (
|
||||
from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
|
||||
from skyvern.forge.sdk.schemas.credentials import (
|
||||
CreateCredentialRequest,
|
||||
Credential,
|
||||
CredentialResponse,
|
||||
CredentialType,
|
||||
CredentialVaultType,
|
||||
@@ -226,7 +227,7 @@ async def create_credential(
|
||||
),
|
||||
current_org: Organization = Depends(org_auth_service.get_current_org),
|
||||
) -> CredentialResponse:
|
||||
credential_service = await _get_credential_vault_service(current_org.organization_id)
|
||||
credential_service = await _get_credential_vault_service()
|
||||
|
||||
credential = await credential_service.create_credential(organization_id=current_org.organization_id, data=data)
|
||||
|
||||
@@ -348,9 +349,13 @@ async def get_credential(
|
||||
),
|
||||
current_org: Organization = Depends(org_auth_service.get_current_org),
|
||||
) -> CredentialResponse:
|
||||
credential_service = await _get_credential_vault_service(current_org.organization_id)
|
||||
credential = await app.DATABASE.get_credential(
|
||||
credential_id=credential_id, organization_id=current_org.organization_id
|
||||
)
|
||||
if not credential:
|
||||
raise HTTPException(status_code=404, detail="Credential not found")
|
||||
|
||||
return await credential_service.get_credential(current_org.organization_id, credential_id)
|
||||
return _convert_to_response(credential)
|
||||
|
||||
|
||||
@legacy_base_router.get("/credentials")
|
||||
@@ -395,9 +400,8 @@ async def get_credentials(
|
||||
openapi_extra={"x-fern-sdk-parameter-name": "page_size"},
|
||||
),
|
||||
) -> list[CredentialResponse]:
|
||||
credential_service = await _get_credential_vault_service(current_org.organization_id)
|
||||
|
||||
return await credential_service.get_credentials(current_org.organization_id, page, page_size)
|
||||
credentials = await app.DATABASE.get_credentials(current_org.organization_id, page=page, page_size=page_size)
|
||||
return [_convert_to_response(credential) for credential in credentials]
|
||||
|
||||
|
||||
@base_router.get(
|
||||
@@ -606,10 +610,8 @@ async def update_azure_client_secret_credential(
|
||||
)
|
||||
|
||||
|
||||
async def _get_credential_vault_service(organization_id: str) -> CredentialVaultService:
|
||||
org_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
|
||||
|
||||
if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN or org_collection:
|
||||
async def _get_credential_vault_service() -> CredentialVaultService:
|
||||
if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN:
|
||||
return app.BITWARDEN_CREDENTIAL_VAULT_SERVICE
|
||||
elif settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.AZURE_VAULT:
|
||||
if not app.AZURE_CREDENTIAL_VAULT_SERVICE:
|
||||
@@ -617,3 +619,30 @@ async def _get_credential_vault_service(organization_id: str) -> CredentialVault
|
||||
return app.AZURE_CREDENTIAL_VAULT_SERVICE
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Credential storage not supported")
|
||||
|
||||
|
||||
def _convert_to_response(credential: Credential) -> CredentialResponse:
|
||||
if credential.credential_type == CredentialType.PASSWORD:
|
||||
credential_response = PasswordCredentialResponse(
|
||||
username=credential.username or credential.credential_id,
|
||||
totp_type=credential.totp_type,
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential.credential_type,
|
||||
name=credential.name,
|
||||
)
|
||||
elif credential.credential_type == CredentialType.CREDIT_CARD:
|
||||
credential_response = CreditCardCredentialResponse(
|
||||
last_four=credential.card_last4 or "****",
|
||||
brand=credential.card_brand or "Card Brand",
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential.credential_type,
|
||||
name=credential.name,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Credential type not supported")
|
||||
|
||||
@@ -3,7 +3,6 @@ from typing import Annotated, Literal, Union
|
||||
|
||||
import structlog
|
||||
from azure.identity.aio import ClientSecretCredential
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, Field, TypeAdapter
|
||||
|
||||
from skyvern.forge import app
|
||||
@@ -12,13 +11,10 @@ from skyvern.forge.sdk.schemas.credentials import (
|
||||
CreateCredentialRequest,
|
||||
Credential,
|
||||
CredentialItem,
|
||||
CredentialResponse,
|
||||
CredentialType,
|
||||
CredentialVaultType,
|
||||
CreditCardCredential,
|
||||
CreditCardCredentialResponse,
|
||||
PasswordCredential,
|
||||
PasswordCredentialResponse,
|
||||
)
|
||||
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
|
||||
|
||||
@@ -107,17 +103,6 @@ class AzureCredentialVaultService(CredentialVaultService):
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
|
||||
credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id)
|
||||
if not credential:
|
||||
raise HTTPException(status_code=404, detail="Credential not found")
|
||||
|
||||
return _convert_to_response(credential)
|
||||
|
||||
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
|
||||
credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size)
|
||||
return [_convert_to_response(credential) for credential in credentials]
|
||||
|
||||
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
|
||||
secret_json_str = await self._client.get_secret(secret_name=db_credential.item_id, vault_name=self._vault_name)
|
||||
if secret_json_str is None:
|
||||
@@ -186,30 +171,3 @@ class AzureCredentialVaultService(CredentialVaultService):
|
||||
secret_name=secret_name,
|
||||
secret_value=secret_value,
|
||||
)
|
||||
|
||||
|
||||
def _convert_to_response(credential: Credential) -> CredentialResponse:
|
||||
if credential.credential_type == CredentialType.PASSWORD:
|
||||
credential_response = PasswordCredentialResponse(
|
||||
username=credential.username or credential.credential_id,
|
||||
totp_type=credential.totp_type,
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential.credential_type,
|
||||
name=credential.name,
|
||||
)
|
||||
elif credential.credential_type == CredentialType.CREDIT_CARD:
|
||||
credential_response = CreditCardCredentialResponse(
|
||||
last_four=credential.card_last4 or "****",
|
||||
brand=credential.card_brand or "Card Brand",
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential.credential_type,
|
||||
name=credential.name,
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Credential type not supported")
|
||||
|
||||
@@ -6,11 +6,7 @@ from skyvern.forge.sdk.schemas.credentials import (
|
||||
CreateCredentialRequest,
|
||||
Credential,
|
||||
CredentialItem,
|
||||
CredentialResponse,
|
||||
CredentialType,
|
||||
CredentialVaultType,
|
||||
CreditCardCredentialResponse,
|
||||
PasswordCredentialResponse,
|
||||
)
|
||||
from skyvern.forge.sdk.services.bitwarden import BitwardenService
|
||||
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
|
||||
@@ -63,88 +59,5 @@ class BitwardenCredentialVaultService(CredentialVaultService):
|
||||
await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id)
|
||||
await BitwardenService.delete_credential_item(credential.item_id)
|
||||
|
||||
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
|
||||
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
|
||||
if not organization_bitwarden_collection:
|
||||
raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
|
||||
|
||||
credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id)
|
||||
if not credential:
|
||||
raise HTTPException(status_code=404, detail="Credential not found")
|
||||
|
||||
credential_item = await BitwardenService.get_credential_item(credential.item_id)
|
||||
if not credential_item:
|
||||
raise HTTPException(status_code=404, detail="Credential not found")
|
||||
|
||||
if credential_item.credential_type == CredentialType.PASSWORD:
|
||||
credential_response = PasswordCredentialResponse(
|
||||
username=credential_item.credential.username,
|
||||
totp_type=credential.totp_type,
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential_item.credential_type,
|
||||
name=credential_item.name,
|
||||
)
|
||||
if credential_item.credential_type == CredentialType.CREDIT_CARD:
|
||||
credential_response = CreditCardCredentialResponse(
|
||||
last_four=credential_item.credential.card_number[-4:],
|
||||
brand=credential_item.credential.card_brand,
|
||||
)
|
||||
return CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=credential_item.credential_type,
|
||||
name=credential_item.name,
|
||||
)
|
||||
raise HTTPException(status_code=400, detail="Invalid credential type")
|
||||
|
||||
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
|
||||
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
|
||||
if not organization_bitwarden_collection:
|
||||
return []
|
||||
|
||||
credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size)
|
||||
items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id)
|
||||
|
||||
response_items = []
|
||||
for credential in credentials:
|
||||
item = next((item for item in items if item.item_id == credential.item_id), None)
|
||||
if not item:
|
||||
LOG.warning(
|
||||
"Credential item not found in vault",
|
||||
credential_id=credential.credential_id,
|
||||
item_id=credential.item_id,
|
||||
)
|
||||
continue
|
||||
if item.credential_type == CredentialType.PASSWORD:
|
||||
credential_response = PasswordCredentialResponse(
|
||||
username=item.credential.username,
|
||||
totp_type=credential.totp_type,
|
||||
)
|
||||
response_items.append(
|
||||
CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=item.credential_type,
|
||||
name=item.name,
|
||||
)
|
||||
)
|
||||
elif item.credential_type == CredentialType.CREDIT_CARD:
|
||||
credential_response = CreditCardCredentialResponse(
|
||||
last_four=item.credential.card_number[-4:],
|
||||
brand=item.credential.card_brand,
|
||||
)
|
||||
response_items.append(
|
||||
CredentialResponse(
|
||||
credential=credential_response,
|
||||
credential_id=credential.credential_id,
|
||||
credential_type=item.credential_type,
|
||||
name=item.name,
|
||||
)
|
||||
)
|
||||
return response_items
|
||||
|
||||
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
|
||||
return await BitwardenService.get_credential_item(db_credential.item_id)
|
||||
|
||||
@@ -5,7 +5,6 @@ from skyvern.forge.sdk.schemas.credentials import (
|
||||
CreateCredentialRequest,
|
||||
Credential,
|
||||
CredentialItem,
|
||||
CredentialResponse,
|
||||
CredentialType,
|
||||
CredentialVaultType,
|
||||
)
|
||||
@@ -32,14 +31,6 @@ class CredentialVaultService(ABC):
|
||||
Default implementation does nothing. Override in subclasses as needed.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
|
||||
"""Retrieve a credential with masked sensitive data."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
|
||||
"""Retrieve all credentials for an organization with pagination."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
|
||||
"""Retrieve the full credential data from the vault."""
|
||||
|
||||
Reference in New Issue
Block a user