Underground-Digital/Workflow-Engine

1import json2 3from core.helper import encrypter4from extensions.ext_database import db5from models.source import DataSourceApiKeyAuthBinding6from services.auth.api_key_auth_factory import ApiKeyAuthFactory7 8 9class ApiKeyAuthService:10    @staticmethod11    def get_provider_auth_list(tenant_id: str) -> list:12        data_source_api_key_bindings = (13            db.session.query(DataSourceApiKeyAuthBinding)14            .filter(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.disabled.is_(False))15            .all()16        )17        return data_source_api_key_bindings18 19    @staticmethod20    def create_provider_auth(tenant_id: str, args: dict):21        auth_result = ApiKeyAuthFactory(args["provider"], args["credentials"]).validate_credentials()22        if auth_result:23            # Encrypt the api key24            api_key = encrypter.encrypt_token(tenant_id, args["credentials"]["config"]["api_key"])25            args["credentials"]["config"]["api_key"] = api_key26 27            data_source_api_key_binding = DataSourceApiKeyAuthBinding()28            data_source_api_key_binding.tenant_id = tenant_id29            data_source_api_key_binding.category = args["category"]30            data_source_api_key_binding.provider = args["provider"]31            data_source_api_key_binding.credentials = json.dumps(args["credentials"], ensure_ascii=False)32            db.session.add(data_source_api_key_binding)33            db.session.commit()34 35    @staticmethod36    def get_auth_credentials(tenant_id: str, category: str, provider: str):37        data_source_api_key_bindings = (38            db.session.query(DataSourceApiKeyAuthBinding)39            .filter(40                DataSourceApiKeyAuthBinding.tenant_id == tenant_id,41                DataSourceApiKeyAuthBinding.category == category,42                DataSourceApiKeyAuthBinding.provider == provider,43                DataSourceApiKeyAuthBinding.disabled.is_(False),44            )45            .first()46        )47        if not data_source_api_key_bindings:48            return None49        credentials = json.loads(data_source_api_key_bindings.credentials)50        return credentials51 52    @staticmethod53    def delete_provider_auth(tenant_id: str, binding_id: str):54        data_source_api_key_binding = (55            db.session.query(DataSourceApiKeyAuthBinding)56            .filter(DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.id == binding_id)57            .first()58        )59        if data_source_api_key_binding:60            db.session.delete(data_source_api_key_binding)61            db.session.commit()62 63    @classmethod64    def validate_api_key_auth_args(cls, args):65        if "category" not in args or not args["category"]:66            raise ValueError("category is required")67        if "provider" not in args or not args["provider"]:68            raise ValueError("provider is required")69        if "credentials" not in args or not args["credentials"]:70            raise ValueError("credentials is required")71        if not isinstance(args["credentials"], dict):72            raise ValueError("credentials must be a dictionary")73        if "auth_type" not in args["credentials"] or not args["credentials"]["auth_type"]:74            raise ValueError("auth_type is required")75