diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index 20ea5d26af..9fdbc08c30 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -683,6 +683,8 @@ WEBUI_SECRET_KEY = os.getenv( os.getenv('WEBUI_JWT_SECRET_KEY', ''), ) +ENABLE_VALVE_ENCRYPTION = os.getenv('ENABLE_VALVE_ENCRYPTION', 'False').lower() == 'true' + WEBUI_SESSION_COOKIE_SAME_SITE = os.getenv('WEBUI_SESSION_COOKIE_SAME_SITE', 'lax') WEBUI_SESSION_COOKIE_SECURE = os.getenv('WEBUI_SESSION_COOKIE_SECURE', 'false').lower() == 'true' WEBUI_AUTH_COOKIE_SAME_SITE = os.getenv('WEBUI_AUTH_COOKIE_SAME_SITE', WEBUI_SESSION_COOKIE_SAME_SITE) diff --git a/backend/open_webui/models/functions.py b/backend/open_webui/models/functions.py index c419dd3f93..a68b689e54 100644 --- a/backend/open_webui/models/functions.py +++ b/backend/open_webui/models/functions.py @@ -7,7 +7,8 @@ import time # local imports from open_webui.internal.db import Base, JSONField, get_async_db_context -from open_webui.models.users import UserModel, UserResponse, Users +from open_webui.models.users import UserResponse, Users +from open_webui.utils.valves import decrypt_valves, encrypt_valves from pydantic import BaseModel, ConfigDict from sqlalchemy import BigInteger, Boolean, Column, Index, String, Text, delete, select, update from sqlalchemy.ext.asyncio import AsyncSession @@ -143,7 +144,8 @@ class FunctionsTable: functions: list[FunctionWithValvesModel], db: AsyncSession | None = None, ) -> list[FunctionWithValvesModel]: - # Synchronize functions for a user by updating existing ones, inserting new ones, and removing those that are no longer present. + # Synchronize functions by updating existing ones, inserting new ones, + # and removing those that are no longer present. try: async with get_async_db_context(db) as db: # Get existing functions @@ -156,24 +158,19 @@ class FunctionsTable: # Update or insert functions for func in functions: + func_data = func.model_dump() + func_data['valves'] = encrypt_valves(func_data['valves']) if func_data.get('valves') else None + func_data['user_id'] = user_id + func_data['updated_at'] = int(time.time()) + if func.id in existing_ids: await db.execute( update(Function) .filter_by(id=func.id) - .values( - **func.model_dump(), - user_id=user_id, - updated_at=int(time.time()), - ) + .values(**func_data) ) else: - new_func = Function( - **{ - **func.model_dump(), - 'user_id': user_id, - 'updated_at': int(time.time()), - } - ) + new_func = Function(**func_data) db.add(new_func) # Remove functions that are no longer present @@ -227,7 +224,15 @@ class FunctionsTable: functions = result.scalars().all() if include_valves: - return [FunctionWithValvesModel.model_validate(function) for function in functions] + return [ + FunctionWithValvesModel.model_validate( + { + **FunctionModel.model_validate(function).model_dump(), + 'valves': decrypt_valves(function.valves), + } + ) + for function in functions + ] else: return [FunctionModel.model_validate(function) for function in functions] @@ -283,7 +288,7 @@ class FunctionsTable: async with get_async_db_context(db) as db: try: function = await db.get(Function, id) - return function.valves if function.valves else {} + return decrypt_valves(function.valves if function else None) except Exception as e: log.exception(f'Error getting function valves by id {id}: {e}') return None @@ -300,7 +305,7 @@ class FunctionsTable: async with get_async_db_context(db) as db: result = await db.execute(select(Function.id, Function.valves).filter(Function.id.in_(ids))) functions = result.all() - return {f.id: (f.valves if f.valves else {}) for f in functions} + return {f.id: decrypt_valves(f.valves) for f in functions} except Exception as e: log.exception(f'Error batch-fetching function valves: {e}') return {} @@ -311,7 +316,7 @@ class FunctionsTable: async with get_async_db_context(db) as db: try: function = await db.get(Function, id) - function.valves = valves + function.valves = encrypt_valves(valves) function.updated_at = int(time.time()) await db.commit() await db.refresh(function) @@ -355,8 +360,8 @@ class FunctionsTable: if 'valves' not in user_settings['functions']: user_settings['functions']['valves'] = {} - return user_settings['functions']['valves'].get(id, {}) - except Exception as e: + return decrypt_valves(user_settings['functions']['valves'].get(id)) + except Exception: log.exception(f'Error getting user values by id {id} and user id {user_id}') return None @@ -373,12 +378,12 @@ class FunctionsTable: if 'valves' not in user_settings['functions']: user_settings['functions']['valves'] = {} - user_settings['functions']['valves'][id] = valves + user_settings['functions']['valves'][id] = encrypt_valves(valves) # Update the user settings in the database await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db) - return user_settings['functions']['valves'][id] + return valves except Exception as e: log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}') return None diff --git a/backend/open_webui/models/tools.py b/backend/open_webui/models/tools.py index 4362f81216..d6309fdf2a 100644 --- a/backend/open_webui/models/tools.py +++ b/backend/open_webui/models/tools.py @@ -10,6 +10,7 @@ from open_webui.internal.db import Base, JSONField, get_async_db_context from open_webui.models.access_grants import AccessGrantModel, AccessGrants from open_webui.models.groups import Groups from open_webui.models.users import UserResponse, Users +from open_webui.utils.valves import decrypt_valves, encrypt_valves from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import BigInteger, Column, String, Text, delete, select, update from sqlalchemy.ext.asyncio import AsyncSession @@ -232,8 +233,8 @@ class ToolsTable: try: async with get_async_db_context(db) as db: tool = await db.get(Tool, id) - return tool.valves if tool.valves else {} - except Exception as e: + return decrypt_valves(tool.valves if tool else None) + except Exception: log.exception(f'Error getting tool valves by id {id}') return None @@ -242,7 +243,11 @@ class ToolsTable: ) -> ToolValves | None: try: async with get_async_db_context(db) as db: - await db.execute(update(Tool).filter_by(id=id).values(valves=valves, updated_at=int(time.time()))) + await db.execute( + update(Tool) + .filter_by(id=id) + .values(valves=encrypt_valves(valves), updated_at=int(time.time())) + ) await db.commit() return await self.get_tool_by_id(id, db=db) except Exception: @@ -261,7 +266,7 @@ class ToolsTable: if 'valves' not in user_settings['tools']: user_settings['tools']['valves'] = {} - return user_settings['tools']['valves'].get(id, {}) + return decrypt_valves(user_settings['tools']['valves'].get(id)) except Exception as e: log.exception(f'Error getting user values by id {id} and user_id {user_id}: {e}') return None @@ -279,12 +284,12 @@ class ToolsTable: if 'valves' not in user_settings['tools']: user_settings['tools']['valves'] = {} - user_settings['tools']['valves'][id] = valves + user_settings['tools']['valves'][id] = encrypt_valves(valves) # Update the user settings in the database await Users.update_user_by_id(user_id, {'settings': user_settings}, db=db) - return user_settings['tools']['valves'][id] + return valves except Exception as e: log.exception(f'Error updating user valves by id {id} and user_id {user_id}: {e}') return None diff --git a/backend/open_webui/utils/valves.py b/backend/open_webui/utils/valves.py new file mode 100644 index 0000000000..c6e2212059 --- /dev/null +++ b/backend/open_webui/utils/valves.py @@ -0,0 +1,41 @@ +import base64 +import hashlib +import json +import logging +from functools import lru_cache + +from cryptography.fernet import Fernet, InvalidToken +from open_webui.env import ENABLE_VALVE_ENCRYPTION, WEBUI_SECRET_KEY + +log = logging.getLogger(__name__) + + +@lru_cache(maxsize=1) +def _fernet() -> Fernet: + key = WEBUI_SECRET_KEY.encode() + if len(WEBUI_SECRET_KEY) != 44: + key = base64.urlsafe_b64encode(hashlib.sha256(key).digest()) + return Fernet(key) + + +def encrypt_valves(valves: dict) -> dict | str: + if not ENABLE_VALVE_ENCRYPTION: + return valves + return _fernet().encrypt(json.dumps(valves).encode()).decode() + + +def decrypt_valves(valves) -> dict: + if not valves: + return {} + if isinstance(valves, dict): + return valves + if not isinstance(valves, str): + return {} + + try: + decrypted = json.loads(_fernet().decrypt(valves.encode()).decode()) + except (InvalidToken, json.JSONDecodeError) as e: + log.warning('Failed to decrypt valves: %s', type(e).__name__) + return {} + + return decrypted if isinstance(decrypted, dict) else {}