refac

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 {}