Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/auth/master_key_boot_check.py: 53%

1import atexit 

2import sys 

3from collections.abc import Awaitable, Callable, Mapping 

4from dataclasses import dataclass, replace 

5from enum import Enum 

6from types import MappingProxyType 

7from typing import Final 

8 

9from typing_extensions import assert_never 

10 

11from litellm._logging import verbose_proxy_logger 

12 

13WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING: Final = "dangerously_permit_weak_or_unset_master_key" 

14WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR: Final = "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY" 

15MASTER_KEY_SETTING: Final = "master_key" 

16MASTER_KEY_ENV_VAR: Final = "LITELLM_MASTER_KEY" 

17SALT_KEY_ENV_VAR: Final = "LITELLM_SALT_KEY" 

18MIGRATE_FROM_MASTER_KEY_ENV_VAR: Final = "LITELLM_MIGRATE_FROM_MASTER_KEY" 

19PUBLICLY_KNOWN_MASTER_KEYS: Final = frozenset({"sk-1234"}) 

20ROTATION_DOCS_URL: Final = "https://docs.litellm.ai/docs/proxy/master_key_rotations#proxy-refuses-to-start" 

21_NEW_MASTER_KEY: Final = "sk-$(openssl rand -hex 32)" 

22GENERATE_MASTER_KEY_COMMAND: Final = f'echo "{MASTER_KEY_ENV_VAR}={_NEW_MASTER_KEY}" | tee -a .env' 

23PRINT_NEW_MASTER_KEY_COMMAND: Final = f'echo "{_NEW_MASTER_KEY}"' 

24 

25 

26class UnsafeMasterKeyReason(Enum): 

27 NOT_SET = "not_set" 

28 EMPTY = "empty" 

29 PUBLICLY_KNOWN = "publicly_known" 

30 

31 

32@dataclass(frozen=True, slots=True) 

33class ConfigFileSource: 

34 config_file_path: str | None 

35 

36 

37@dataclass(frozen=True, slots=True) 

38class EnvironmentSource: 

39 pass 

40 

41 

42MasterKeySource = ConfigFileSource | EnvironmentSource 

43 

44 

45@dataclass(frozen=True, slots=True) 

46class SafeMasterKey: 

47 pass 

48 

49 

50@dataclass(frozen=True, slots=True) 

51class UnsafeMasterKeyAllowed: 

52 reason: UnsafeMasterKeyReason 

53 

54 

55@dataclass(frozen=True, slots=True) 

56class StoredSecretsMigration: 

57 from_master_key: str 

58 encrypted_value_count: int | None 

59 

60 

61@dataclass(frozen=True, slots=True) 

62class UnsafeMasterKeyRefused: 

63 reason: UnsafeMasterKeyReason 

64 source: MasterKeySource 

65 environment_variable_is_set: bool 

66 migration: StoredSecretsMigration | None 

67 

68 

69MasterKeyBootVerdict = SafeMasterKey | UnsafeMasterKeyAllowed | UnsafeMasterKeyRefused 

70 

71 

72class UnsafeMasterKeyError(Exception): 

73 pass 

74 

75 

76def master_key_boot_verdict( 

77 *, 

78 master_key: str | None, 

79 environment_master_key: str | None, 

80 general_settings: Mapping[str, object], 

81 config_file_path: str | None, 

82 override_env_is_on: bool, 

83 salt_key_is_set: bool, 

84 database_is_configured: bool, 

85) -> MasterKeyBootVerdict: 

86 reason: Final = _unsafe_reason(master_key) 

87 if reason is None: 87 ↛ 89line 87 didn't jump to line 89 because the condition on line 87 was always true

88 return SafeMasterKey() 

89 if override_env_is_on or general_settings.get(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING) is True: 

90 return UnsafeMasterKeyAllowed(reason=reason) 

91 config_file_only_relays_the_environment: Final = master_key is not None and master_key == environment_master_key 

92 return UnsafeMasterKeyRefused( 

93 reason=reason, 

94 source=( 

95 ConfigFileSource(config_file_path=config_file_path) 

96 if MASTER_KEY_SETTING in general_settings and not config_file_only_relays_the_environment 

97 else EnvironmentSource() 

98 ), 

99 environment_variable_is_set=environment_master_key is not None, 

100 migration=( 

101 StoredSecretsMigration(from_master_key=master_key, encrypted_value_count=None) 

102 if master_key is not None and not salt_key_is_set and database_is_configured 

103 else None 

104 ), 

105 ) 

106 

107 

108async def with_stored_secrets_counted( 

109 verdict: MasterKeyBootVerdict, count_values_encrypted_with: Callable[[str], Awaitable[int | None]] 

110) -> MasterKeyBootVerdict: 

111 if not isinstance(verdict, UnsafeMasterKeyRefused) or verdict.migration is None: 111 ↛ 113line 111 didn't jump to line 113 because the condition on line 111 was always true

112 return verdict 

113 count: Final = await count_values_encrypted_with(verdict.migration.from_master_key) 

114 return replace(verdict, migration=None if count == 0 else replace(verdict.migration, encrypted_value_count=count)) 

115 

116 

117def enforce_master_key_boot_verdict(verdict: MasterKeyBootVerdict, announce: Callable[[str], object]) -> None: 

118 match verdict: 

119 case SafeMasterKey(): 119 ↛ 121line 119 didn't jump to line 121 because the pattern on line 119 always matched

120 return 

121 case UnsafeMasterKeyAllowed(reason=reason): 

122 verbose_proxy_logger.warning( 

123 "%s is on, so the proxy is starting with %s. Never run this outside local development.", 

124 WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING, 

125 _UNSAFE_STATE[reason], 

126 ) 

127 case UnsafeMasterKeyRefused(reason=reason): 

128 announce(f"\n{render_refusal(verdict)}\n\n") 

129 raise UnsafeMasterKeyError( 

130 f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[reason]} The fix is printed once the server exits." 

131 ) 

132 case _: 

133 assert_never(verdict) 

134 

135 

136def announce_on_stderr_at_exit(message: str) -> None: 

137 """A logger would redact the key-shaped command and the lifespan traceback would bury it, so print at exit.""" 

138 atexit.register(_flush_stdout_then_write_stderr, message) 

139 

140 

141def _flush_stdout_then_write_stderr(message: str) -> None: 

142 sys.stdout.flush() 

143 sys.stderr.write(message) 

144 

145 

146def render_refusal(refusal: UnsafeMasterKeyRefused) -> str: 

147 return "\n\n".join( 

148 ( 

149 f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[refusal.reason]}\n{_source_line(refusal)}", 

150 _fix_steps(refusal), 

151 _OVERRIDE_HINT, 

152 ) 

153 ) 

154 

155 

156_UNSAFE_STATE: Final = MappingProxyType( 

157 { 

158 UnsafeMasterKeyReason.NOT_SET: "no master key, which accepts every request without authentication", 

159 UnsafeMasterKeyReason.EMPTY: "an empty master key", 

160 UnsafeMasterKeyReason.PUBLICLY_KNOWN: "a publicly known master key", 

161 } 

162) 

163 

164_REFUSAL_HEADLINE: Final = MappingProxyType( 

165 { 

166 UnsafeMasterKeyReason.NOT_SET: ( 

167 "no master key is set, so every request would be accepted without authentication." 

168 ), 

169 UnsafeMasterKeyReason.EMPTY: "the master key is empty.", 

170 UnsafeMasterKeyReason.PUBLICLY_KNOWN: "the master key is a publicly known default.", 

171 } 

172) 

173 

174_SAVE_KEY_STEP: Final = ( 

175 "Generate a key and save it to .env:\n" 

176 f" {GENERATE_MASTER_KEY_COMMAND}\n" 

177 " Not using a .env file (docker run, Kubernetes, pip install)? Pass the same value as the\n" 

178 f" {MASTER_KEY_ENV_VAR} environment variable instead." 

179) 

180 

181_REPLACE_EXPORTED_KEY_STEP: Final = ( 

182 "Generate a key:\n" 

183 f" {PRINT_NEW_MASTER_KEY_COMMAND}\n" 

184 f" Put it in place of the current {MASTER_KEY_ENV_VAR} value wherever that is set: a shell export, your\n" 

185 " container or deployment environment, or its line in .env. Do not just add it to .env, because a value\n" 

186 " already exported in the environment wins over .env." 

187) 

188 

189_RESTART_TO_MIGRATE_STEP: Final = ( 

190 "Start the proxy again. It re-encrypts the stored values with the new key, then logs that\n" 

191 f" {MIGRATE_FROM_MASTER_KEY_ENV_VAR} can be removed. Details: {ROTATION_DOCS_URL}" 

192) 

193 

194_OVERRIDE_HINT: Final = ( 

195 f"Local development only: set {WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR}=true, or\n" 

196 f"general_settings.{WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING}: true, to start anyway." 

197) 

198 

199 

200def _unsafe_reason(master_key: str | None) -> UnsafeMasterKeyReason | None: 

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

202 return UnsafeMasterKeyReason.NOT_SET 

203 stripped: Final = master_key.strip() 

204 if not stripped: 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true

205 return UnsafeMasterKeyReason.EMPTY 

206 if stripped in PUBLICLY_KNOWN_MASTER_KEYS: 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true

207 return UnsafeMasterKeyReason.PUBLICLY_KNOWN 

208 return None 

209 

210 

211def _config_label(source: ConfigFileSource) -> str: 

212 return source.config_file_path or "your config" 

213 

214 

215def _source_line(refusal: UnsafeMasterKeyRefused) -> str: 

216 match refusal.source: 

217 case ConfigFileSource() as source: 

218 if refusal.reason is UnsafeMasterKeyReason.NOT_SET: 

219 return ( 

220 f"general_settings.{MASTER_KEY_SETTING} in {_config_label(source)} is blank, or points at an " 

221 "environment variable that is not set." 

222 ) 

223 return f"It comes from general_settings.{MASTER_KEY_SETTING} in {_config_label(source)}." 

224 case EnvironmentSource(): 

225 if refusal.reason is UnsafeMasterKeyReason.NOT_SET: 

226 return ( 

227 f"Neither general_settings.{MASTER_KEY_SETTING} nor the {MASTER_KEY_ENV_VAR} " 

228 "environment variable is set." 

229 ) 

230 return f"It comes from the {MASTER_KEY_ENV_VAR} environment variable." 

231 case _: 

232 assert_never(refusal.source) 

233 

234 

235def _fix_steps(refusal: UnsafeMasterKeyRefused) -> str: 

236 steps: Final = (*_config_steps(refusal.source), *_key_steps(refusal)) 

237 numbered: Final = "\n".join(f"{number}. {step}" for number, step in enumerate(steps, start=1)) 

238 return numbered if refusal.migration is None else f"{_migration_lead(refusal.migration)}\n{numbered}" 

239 

240 

241def _config_steps(source: MasterKeySource) -> tuple[str, ...]: 

242 match source: 

243 case ConfigFileSource(): 

244 return ( 

245 f"Make sure {_config_label(source)} reads the key from the environment:\n" 

246 f" general_settings:\n {MASTER_KEY_SETTING}: os.environ/{MASTER_KEY_ENV_VAR}", 

247 ) 

248 case EnvironmentSource(): 

249 return () 

250 case _: 

251 assert_never(source) 

252 

253 

254def _key_steps(refusal: UnsafeMasterKeyRefused) -> tuple[str, ...]: 

255 if refusal.migration is None: 

256 return (_REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP,) 

257 if refusal.environment_variable_is_set: 

258 return ( 

259 f"Set the key to migrate from next to {MASTER_KEY_ENV_VAR}, wherever that is set (a shell export, your\n" 

260 " container or deployment environment, or .env):\n" 

261 f" {_migrate_from_assignment(refusal.migration)}", 

262 _REPLACE_EXPORTED_KEY_STEP, 

263 _RESTART_TO_MIGRATE_STEP, 

264 ) 

265 return ( 

266 "Save the key to migrate from and a newly generated key to .env:\n" 

267 f" echo '{_migrate_from_assignment(refusal.migration)}' | tee -a .env\n" 

268 f" {GENERATE_MASTER_KEY_COMMAND}\n" 

269 " Not using a .env file (docker run, Kubernetes, pip install)? Pass the same two values as\n" 

270 " environment variables instead.", 

271 _RESTART_TO_MIGRATE_STEP, 

272 ) 

273 

274 

275def _migrate_from_assignment(migration: StoredSecretsMigration) -> str: 

276 key: Final = migration.from_master_key 

277 value: Final = key if key == key.strip() else f'"{key}"' 

278 return f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}={value}" 

279 

280 

281def _migration_lead(migration: StoredSecretsMigration) -> str: 

282 found: Final = ( 

283 "could not be checked for values" 

284 if migration.encrypted_value_count is None 

285 else f"holds {migration.encrypted_value_count} value(s)" 

286 ) 

287 return ( 

288 f"Your database {found} encrypted with this master key,\n" 

289 f"which encrypts stored credentials while {SALT_KEY_ENV_VAR} is not set. Replacing the key alone makes them\n" 

290 "unreadable, so also tell the proxy which key to migrate from:" 

291 )