Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/encrypt_decrypt_utils.py: 74%

1import base64 

2import os 

3from collections.abc import Mapping 

4from typing import Final, Literal, cast 

5 

6from pydantic import TypeAdapter, ValidationError 

7 

8from litellm._logging import verbose_proxy_logger 

9 

10# Versioned ciphertext marker for AES-256-GCM values. 

11# Format: "v2:gcm:" + base64url(nonce(12) || ciphertext || tag(16)). 

12# Legacy XSalsa20-Poly1305 (nacl) values carry no marker; the colon in the 

13# prefix can never appear in base64url(nacl output), so the prefix check is an 

14# unambiguous discriminator between the two formats on read. 

15_V2_GCM_PREFIX: Final = "v2:gcm:" 

16 

17# general_settings key selecting the at-rest encryption algorithm for new writes. 

18# Default preserves the legacy algorithm so existing deployments are byte-for-byte 

19# unchanged until they explicitly opt in. Decrypt is always format-detecting, so 

20# flipping this flag forward (or back) never strands previously-written data. 

21_ENCRYPTION_ALGORITHM_SETTING: Final = "encryption_algorithm" 

22_ALGO_AES_GCM: Final = "aes-256-gcm" 

23_ALGO_XSALSA20: Final = "xsalsa20-poly1305" 

24 

25 

26def _get_salt_key(): 

27 from litellm.proxy.proxy_server import master_key 

28 

29 salt_key = os.getenv("LITELLM_SALT_KEY", None) 

30 

31 if salt_key is None: 31 ↛ 32line 31 didn't jump to line 32 because the condition on line 31 was never true

32 salt_key = master_key 

33 

34 return salt_key 

35 

36 

37def _get_encryption_algorithm() -> str: 

38 """ 

39 Resolve the configured at-rest encryption algorithm for *new writes*. 

40 

41 Read from ``general_settings.encryption_algorithm`` at write time. Defaults to 

42 the legacy XSalsa20-Poly1305 algorithm so deployments that have not opted in 

43 keep producing byte-for-byte identical ciphertext. 

44 """ 

45 try: 

46 from litellm.proxy.proxy_server import general_settings 

47 

48 algo: Final = general_settings.get(_ENCRYPTION_ALGORITHM_SETTING, _ALGO_XSALSA20) 

49 except Exception: 

50 # general_settings may not be importable in some contexts (e.g. SDK-only 

51 # use of these helpers). Fall back to the legacy algorithm. 

52 return _ALGO_XSALSA20 

53 

54 if isinstance(algo, str) and algo.lower() == _ALGO_AES_GCM: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true

55 return _ALGO_AES_GCM 

56 return _ALGO_XSALSA20 

57 

58 

59def _derive_key(signing_key: str) -> bytes: 

60 """Derive a 32-byte key from the salt/master key (shared by both algorithms). 

61 

62 Known limitation: this is a single-pass, unsalted ``SHA-256`` of the key, not 

63 a dedicated KDF (HKDF/PBKDF2). It is the *same* derivation the legacy nacl 

64 path already uses, so the AES path introduces no new weakness and stays 

65 interoperable with existing key sourcing; AES-256-GCM's per-value 12-byte 

66 random nonce gives the unique (key, nonce) pairs GCM requires. Moving both 

67 algorithms to HKDF-SHA256 would be more defensible in an audit but is a 

68 separate, coordinated change (it must re-derive or re-encrypt existing data). 

69 """ 

70 import hashlib 

71 

72 return hashlib.sha256(signing_key.encode()).digest() 

73 

74 

75def _seal_aes_gcm(value: str, signing_key: str, aad: bytes | None) -> bytes: 

76 from cryptography.hazmat.primitives.ciphers.aead import AESGCM 

77 

78 nonce: Final = os.urandom(12) 

79 # AESGCM.encrypt returns ciphertext || tag(16); wire format is nonce || that. 

80 return nonce + AESGCM(_derive_key(signing_key)).encrypt(nonce, value.encode("utf-8"), aad) 

81 

82 

83def _open_aes_gcm(sealed: bytes, signing_key: str, aad: bytes | None) -> str: 

84 from cryptography.hazmat.primitives.ciphers.aead import AESGCM 

85 

86 # An empty plaintext still serializes to nonce(12) || tag(16) = 28 bytes, so a 

87 # short/empty buffer here is a corrupt value: let AESGCM.decrypt raise and be 

88 # swallowed by the caller (returns None/original), same as legacy. 

89 return AESGCM(_derive_key(signing_key)).decrypt(sealed[:12], sealed[12:], aad).decode("utf-8") 

90 

91 

92def _encrypt_aes_gcm(value: str, signing_key: str) -> str: 

93 """Encrypt under AES-256-GCM and return the versioned ``v2:gcm:`` string.""" 

94 sealed: Final = _seal_aes_gcm(value=value, signing_key=signing_key, aad=None) 

95 return _V2_GCM_PREFIX + base64.urlsafe_b64encode(sealed).decode("utf-8") 

96 

97 

98def _decrypt_aes_gcm(value: str, signing_key: str) -> str: 

99 """Decrypt a versioned ``v2:gcm:`` string produced by :func:`_encrypt_aes_gcm`.""" 

100 sealed: Final = base64.urlsafe_b64decode(value[len(_V2_GCM_PREFIX) :]) 

101 return _open_aes_gcm(sealed=sealed, signing_key=signing_key, aad=None) 

102 

103 

104def encrypt_bearer_token(value: str, prefix: str) -> str: 

105 """AES-256-GCM as unpadded base64url behind ``prefix``, which is also the AAD so a token can't change kind.""" 

106 salt_key: Final = _get_salt_key() 

107 if not isinstance(salt_key, str): 

108 raise ValueError("Set LITELLM_SALT_KEY or a master key to mint bearer tokens") # noqa: TRY004 # missing config, not a bad argument type 

109 sealed: Final = _seal_aes_gcm(value=value, signing_key=salt_key, aad=prefix.encode("utf-8")) 

110 return prefix + base64.urlsafe_b64encode(sealed).decode("ascii").rstrip("=") 

111 

112 

113def decrypt_bearer_token(token: str, prefix: str) -> str | None: 

114 """None unless ``token`` came from :func:`encrypt_bearer_token` with the same ``prefix``.""" 

115 salt_key: Final = _get_salt_key() 

116 if not isinstance(salt_key, str) or not token.startswith(prefix): 116 ↛ 118line 116 didn't jump to line 118 because the condition on line 116 was always true

117 return None 

118 encoded: Final = token.removeprefix(prefix) 

119 try: 

120 sealed: Final = base64.b64decode(encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True) 

121 return _open_aes_gcm(sealed=sealed, signing_key=salt_key, aad=prefix.encode("utf-8")) 

122 except Exception: # noqa: BLE001 # base64 and AES-GCM each raise their own "not a token" type 

123 return None 

124 

125 

126def encrypt_value_helper(value: str, new_encryption_key: str | None = None): 

127 signing_key: Final = new_encryption_key or _get_salt_key() 

128 

129 try: 

130 if isinstance(value, str): 

131 if _get_encryption_algorithm() == _ALGO_AES_GCM: 131 ↛ 134line 131 didn't jump to line 134 because the condition on line 131 was never true

132 # AES path: the v2:gcm: output is already a base64url string, so it 

133 # is returned directly with no extra base64 wrapper. 

134 return _encrypt_aes_gcm(value=value, signing_key=cast(str, signing_key)) 

135 

136 encrypted_value = encrypt_value(value=value, signing_key=signing_key) 

137 # Use urlsafe_b64encode for URL-safe base64 encoding (replaces + with - and / with _) 

138 encrypted_value = base64.urlsafe_b64encode(encrypted_value).decode("utf-8") 

139 

140 return encrypted_value 

141 

142 verbose_proxy_logger.debug( 

143 "Invalid value type passed to encrypt_value: %s for Value: %s\n Value must be a string", type(value), value 

144 ) 

145 # if it's not a string - do not encrypt it and return the value 

146 return value 

147 except Exception as e: 

148 raise e 

149 

150 

151def _legacy_ciphertext_bytes(value: str) -> bytes: 

152 # Try URL-safe base64 decoding first (new format) 

153 # Fall back to standard base64 decoding for backwards compatibility (old format) 

154 try: 

155 return base64.urlsafe_b64decode(value) 

156 except Exception: 

157 return base64.b64decode(value) 

158 

159 

160def _decrypt_with_signing_key(value: str, signing_key: str) -> str: 

161 # Versioned AES-256-GCM values are detected before any base64 decode. 

162 # The prefix is the algorithm tag the legacy nacl format never carried. 

163 if value.startswith(_V2_GCM_PREFIX): 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true

164 return _decrypt_aes_gcm(value=value, signing_key=signing_key) 

165 

166 return decrypt_value(value=_legacy_ciphertext_bytes(value), signing_key=signing_key) 

167 

168 

169def decrypt_if_encrypted_with(value: str, signing_key: str) -> str | None: 

170 """None unless value is a ciphertext under signing_key.""" 

171 try: 

172 # base64 decoding skips characters outside its alphabet, so "" and "*" decode to no bytes, 

173 # which decrypt_value reads as an empty plaintext under any key. 

174 decodes_to_nothing: Final = not value.startswith(_V2_GCM_PREFIX) and not _legacy_ciphertext_bytes(value) 

175 return None if decodes_to_nothing else _decrypt_with_signing_key(value=value, signing_key=signing_key) 

176 except Exception: # noqa: BLE001 # base64, nacl and AES-GCM each raise their own "not a ciphertext" type 

177 return None 

178 

179 

180def decrypt_value_helper( 

181 value: str, 

182 key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key. 

183 exception_type: Literal["debug", "error"] = "error", 

184 return_original_value: bool = False, 

185) -> str | None: 

186 signing_key: Final = _get_salt_key() 

187 

188 try: 

189 if isinstance(value, str): 

190 return _decrypt_with_signing_key(value=value, signing_key=cast(str, signing_key)) 

191 

192 # if it's not str - do not decrypt it, return the value 

193 return value 

194 except Exception as e: 

195 error_message = f"Error decrypting value for key: {key}, Did your master_key/salt key change recently? \nError: {e}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" 

196 if exception_type == "debug": 

197 verbose_proxy_logger.debug(error_message) 

198 return value if return_original_value else None 

199 

200 verbose_proxy_logger.debug("Unable to decrypt value for key: %s, returning None", key) 

201 if return_original_value: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true

202 return value 

203 else: 

204 verbose_proxy_logger.exception(error_message) 

205 # [Non-Blocking Exception. - this should not block decrypting other values] 

206 return None 

207 

208 

209def encrypt_value(value: str, signing_key: str): 

210 import hashlib 

211 

212 import nacl.secret 

213 import nacl.utils 

214 

215 # get 32 byte master key # 

216 hash_object: Final = hashlib.sha256(signing_key.encode()) 

217 hash_bytes: Final = hash_object.digest() 

218 

219 # initialize secret box # 

220 box: Final = nacl.secret.SecretBox(hash_bytes) 

221 

222 # encode message # 

223 value_bytes: Final = value.encode("utf-8") 

224 

225 encrypted: Final = box.encrypt(value_bytes) 

226 

227 return encrypted 

228 

229 

230def decrypt_value(value: bytes, signing_key: str) -> str: 

231 import hashlib 

232 

233 import nacl.secret 

234 import nacl.utils 

235 

236 # get 32 byte master key # 

237 hash_object: Final = hashlib.sha256(signing_key.encode()) 

238 hash_bytes: Final = hash_object.digest() 

239 

240 # initialize secret box # 

241 box: Final = nacl.secret.SecretBox(hash_bytes) 

242 

243 # Convert the bytes object to a string 

244 try: 

245 if len(value) == 0: 

246 return "" 

247 

248 plaintext = box.decrypt(value) 

249 plaintext = plaintext.decode("utf-8") 

250 return plaintext 

251 except Exception as e: 

252 raise e 

253 

254 

255class SecretMapDecodeError(RuntimeError): 

256 pass 

257 

258 

259_SECRET_MAP: Final = TypeAdapter(Mapping[str, str]) 

260_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str) 

261_SECRET_STRING: Final = TypeAdapter(str) 

262 

263 

264def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str: 

265 if not value: 

266 return "{}" 

267 ciphertext: Final = _SECRET_STRING.validate_python( 

268 encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True 

269 ) 

270 return _SECRET_STRING.dump_json(ciphertext).decode() 

271 

272 

273def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None: 

274 if value is None: 

275 return None 

276 try: 

277 stored: Final = ( 

278 _STORED_SECRET_MAP.validate_json(value, strict=True) 

279 if isinstance(value, str) and value.lstrip().startswith(("{", '"')) 

280 else _STORED_SECRET_MAP.validate_python(value, strict=True) 

281 ) 

282 if not isinstance(stored, str): 

283 return stored 

284 decrypted: Final = decrypt_value_helper( 

285 value=stored, key=key, exception_type="debug", return_original_value=False 

286 ) 

287 return _SECRET_MAP.validate_json(decrypted, strict=True) 

288 except ValidationError: 

289 raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None