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

1# What is this? 

2## If litellm license in env, checks if it's valid 

3import base64 

4import json 

5import os 

6from datetime import datetime 

7from typing import TYPE_CHECKING, Final 

8 

9import httpx 

10 

11from litellm._logging import verbose_proxy_logger 

12from litellm.constants import NON_LLM_CONNECTION_TIMEOUT 

13from litellm.llms.custom_httpx.http_handler import HTTPHandler 

14 

15if TYPE_CHECKING: 15 ↛ 16line 15 didn't jump to line 16 because the condition on line 15 was never true

16 from litellm.proxy._types import EnterpriseLicenseData 

17 

18 

19AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" 

20LICENSE_ALL_FEATURES: Final = "*" 

21AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." 

22 

23 

24class LicenseCheck: 

25 """ 

26 - Check if license in env 

27 - Returns if license is valid 

28 """ 

29 

30 base_url = "https://license.litellm.ai" 

31 

32 def __init__(self) -> None: 

33 self.license_str = os.getenv("LITELLM_LICENSE", None) 

34 verbose_proxy_logger.debug("License Str value - %s", self.license_str) 

35 self.http_handler = HTTPHandler(timeout=NON_LLM_CONNECTION_TIMEOUT) 

36 self._premium_check_logged = False 

37 self.public_key = None 

38 self.read_public_key() 

39 self.airgapped_license_data: EnterpriseLicenseData | None = None 

40 

41 def read_public_key(self): 

42 try: 

43 from cryptography.hazmat.primitives import serialization 

44 

45 # current dir 

46 current_dir: Final = os.path.dirname(os.path.realpath(__file__)) 

47 

48 # check if public_key.pem exists 

49 _path_to_public_key: Final = os.path.join(current_dir, "public_key.pem") 

50 if os.path.exists(_path_to_public_key): 50 ↛ 54line 50 didn't jump to line 54 because the condition on line 50 was always true

51 with open(_path_to_public_key, "rb") as key_file: 

52 self.public_key = serialization.load_pem_public_key(key_file.read()) 

53 else: 

54 self.public_key = None 

55 except Exception as e: 

56 verbose_proxy_logger.error("Error reading public key: %s", e) 

57 

58 def _verify(self, license_str: str) -> bool: 

59 verbose_proxy_logger.debug( 

60 "litellm.proxy.auth.litellm_license.py::_verify - Checking license against %s/verify_license - %s", 

61 self.base_url, 

62 license_str, 

63 ) 

64 url: Final = f"{self.base_url}/verify_license/{license_str}" 

65 

66 response: httpx.Response | None = None 

67 try: # don't impact user, if call fails 

68 num_retries: Final = 3 

69 for i in range(num_retries): 

70 try: 

71 response = self.http_handler.get(url=url) 

72 if response is None: 

73 raise Exception("No response from license server") 

74 response.raise_for_status() 

75 except httpx.HTTPStatusError: 

76 if i == num_retries - 1: 

77 raise 

78 

79 if response is None: 

80 raise Exception("No response from license server") 

81 

82 response_json: Final = response.json() 

83 

84 premium: Final = response_json["verify"] 

85 

86 assert isinstance(premium, bool) 

87 

88 verbose_proxy_logger.debug( 

89 "litellm.proxy.auth.litellm_license.py::_verify - License=%s is premium=%s", license_str, premium 

90 ) 

91 return premium 

92 except Exception as e: 

93 verbose_proxy_logger.exception( 

94 "litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License=%s via api. - %s", 

95 license_str, 

96 e, 

97 ) 

98 return False 

99 

100 def is_premium(self) -> bool: 

101 """ 

102 1. verify_license_without_api_request: checks if license was generate using private / public key pair 

103 2. _verify: checks if license is valid calling litellm API. This is the old way we were generating/validating license 

104 """ 

105 try: 

106 if not self._premium_check_logged: 

107 verbose_proxy_logger.debug( 

108 "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - LiteLLM License=%s", 

109 self.license_str, 

110 ) 

111 

112 if self.license_str is None: 112 ↛ 115line 112 didn't jump to line 115 because the condition on line 112 was always true

113 self.license_str = os.getenv("LITELLM_LICENSE", None) 

114 

115 if not self._premium_check_logged: 

116 verbose_proxy_logger.debug( 

117 "litellm.proxy.auth.litellm_license.py::is_premium() - Updated 'self.license_str' - %s", 

118 self.license_str, 

119 ) 

120 self._premium_check_logged = True 

121 

122 if self.license_str is None: 122 ↛ 124line 122 didn't jump to line 124 because the condition on line 122 was always true

123 return False 

124 elif ( 

125 self.verify_license_without_api_request(public_key=self.public_key, license_key=self.license_str) 

126 is True 

127 ) or self._verify(license_str=self.license_str) is True: 

128 return True 

129 return False 

130 except Exception: 

131 return False 

132 

133 def is_over_limit(self, total_users: int) -> bool: 

134 """ 

135 Check if the license is over the limit 

136 """ 

137 if self.airgapped_license_data is None: 137 ↛ 139line 137 didn't jump to line 139 because the condition on line 137 was always true

138 return False 

139 if "max_users" not in self.airgapped_license_data or not isinstance( 

140 self.airgapped_license_data["max_users"], int 

141 ): 

142 return False 

143 return total_users > self.airgapped_license_data["max_users"] 

144 

145 def is_team_count_over_limit(self, team_count: int) -> bool: 

146 """ 

147 Check if the license is over the limit 

148 """ 

149 if self.airgapped_license_data is None: 149 ↛ 152line 149 didn't jump to line 152 because the condition on line 149 was always true

150 return False 

151 

152 _max_teams_in_license: Final[int | None] = self.airgapped_license_data.get("max_teams") 

153 if "max_teams" not in self.airgapped_license_data or not isinstance(_max_teams_in_license, int): 

154 return False 

155 return team_count > _max_teams_in_license 

156 

157 def grants_feature(self, feature: str) -> bool: 

158 if self.airgapped_license_data is None: 158 ↛ 160line 158 didn't jump to line 160 because the condition on line 158 was always true

159 return False 

160 allowed_features: Final = self.airgapped_license_data.get("allowed_features") 

161 granted: Final = allowed_features if isinstance(allowed_features, list) else (allowed_features,) 

162 return feature in granted or LICENSE_ALL_FEATURES in granted 

163 

164 def auto_router_capability_limit(self) -> int | None: 

165 """ 

166 How many auto-routers may claim each gated classifier or customization capability: 

167 unlimited (None) only when the signed license lists the auto_router feature or the 

168 "*" wildcard that grants every feature, otherwise one per capability. A license verified 

169 through the API carries no feature list, so it does not lift the limit either. 

170 """ 

171 if self.grants_feature(AUTO_ROUTER_LICENSE_FEATURE): 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true

172 return None 

173 return 1 

174 

175 def verify_license_without_api_request(self, public_key, license_key): 

176 try: 

177 from cryptography.hazmat.primitives import hashes 

178 from cryptography.hazmat.primitives.asymmetric import padding 

179 

180 from litellm.proxy._types import EnterpriseLicenseData 

181 

182 # Decode the license key - add padding if needed for base64 

183 # Base64 strings need to be a multiple of 4 characters 

184 padding_needed: Final = len(license_key) % 4 

185 if padding_needed: 

186 license_key += "=" * (4 - padding_needed) 

187 

188 decoded: Final = base64.b64decode(license_key) 

189 message, signature = decoded.split(b".", 1) 

190 

191 # Verify the signature 

192 public_key.verify( 

193 signature, 

194 message, 

195 padding.PSS( 

196 mgf=padding.MGF1(hashes.SHA256()), 

197 salt_length=padding.PSS.MAX_LENGTH, 

198 ), 

199 hashes.SHA256(), 

200 ) 

201 

202 # Decode and parse the data 

203 license_data: Final = json.loads(message.decode()) 

204 

205 # debug information provided in license data 

206 verbose_proxy_logger.debug("License data: %s", license_data) 

207 

208 # Check expiration date 

209 expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") 

210 if expiration_date < datetime.now(): 

211 self.airgapped_license_data = None 

212 return False, "License has expired" 

213 

214 self.airgapped_license_data = EnterpriseLicenseData(**license_data) 

215 

216 return True 

217 

218 except Exception as e: 

219 self.airgapped_license_data = None 

220 verbose_proxy_logger.debug( 

221 "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s", 

222 e, 

223 ) 

224 return False