Coverage for open_webui/utils/valves.py: 29%

1import base64 

2import hashlib 

3import logging 

4from functools import lru_cache 

5 

6from cryptography.fernet import Fernet, InvalidToken 

7from open_webui.env import ENABLE_VALVE_ENCRYPTION, WEBUI_SECRET_KEY 

8from open_webui.utils.json_codec import JSONCodec 

9 

10log = logging.getLogger(__name__) 

11 

12 

13@lru_cache(maxsize=1) 

14def _fernet() -> Fernet: 

15 key = WEBUI_SECRET_KEY.encode() 

16 if len(WEBUI_SECRET_KEY) != 44: 

17 key = base64.urlsafe_b64encode(hashlib.sha256(key).digest()) 

18 return Fernet(key) 

19 

20 

21def encrypt_valves(valves: dict) -> dict | str: 

22 if not ENABLE_VALVE_ENCRYPTION: 

23 return valves 

24 return _fernet().encrypt(JSONCodec.dumps(valves).encode()).decode() 

25 

26 

27def decrypt_valves(valves) -> dict: 

28 if not valves: 

29 return {} 

30 if isinstance(valves, dict): 

31 return valves 

32 if not isinstance(valves, str): 

33 return {} 

34 

35 try: 

36 decrypted = JSONCodec.loads(_fernet().decrypt(valves.encode()).decode()) 

37 except (InvalidToken, JSONCodec.JSONDecodeError) as e: 

38 log.warning('Failed to decrypt valves: %s', type(e).__name__) 

39 return {} 

40 

41 return decrypted if isinstance(decrypted, dict) else {}