Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/rag_endpoints/upload_security.py: 44%

1"""Security controls for vector-store file uploads. 

2 

3Content is classified by inspecting its actual bytes (magic signatures and a 

4strict UTF-8 decode), never by trusting the client-supplied filename or 

5content-type. Uploads are restricted to an allowlist of non-executable formats, 

6capped in size, screened for archives, and passed through a dependency-injected 

7malware scanner before they are accepted. Accepted uploads are given a 

8server-generated filename so the client-controlled name never reaches storage. 

9""" 

10 

11from __future__ import annotations 

12 

13import uuid 

14from collections.abc import Mapping 

15from dataclasses import dataclass 

16from enum import Enum 

17from types import MappingProxyType 

18from typing import Final, Protocol, TypeAlias, runtime_checkable 

19 

20from typing_extensions import assert_never 

21 

22MAX_UPLOAD_SIZE_BYTES: Final = 512 * 1024 * 1024 

23 

24EICAR_TEST_SIGNATURE: Final = b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" 

25 

26_ARCHIVE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = ( 

27 b"PK\x03\x04", 

28 b"PK\x05\x06", 

29 b"PK\x07\x08", 

30 b"\x1f\x8b", 

31 b"\xfd7zXZ\x00", 

32 b"7z\xbc\xaf\x27\x1c", 

33 b"Rar!\x1a\x07\x00", 

34 b"Rar!\x1a\x07\x01\x00", 

35 b"\x04\x22\x4d\x18", 

36 b"\x28\xb5\x2f\xfd", 

37) 

38 

39_ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"BZh",) 

40 

41_EXECUTABLE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = ( 

42 b"\x7fELF", 

43 b"\xca\xfe\xba\xbe", 

44 b"\xfe\xed\xfa\xce", 

45 b"\xfe\xed\xfa\xcf", 

46 b"\xce\xfa\xed\xfe", 

47 b"\xcf\xfa\xed\xfe", 

48 b"\x00asm", 

49) 

50 

51_EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"MZ", b"dex\n") 

52 

53_TAR_USTAR_MAGIC: Final = b"ustar" 

54_TAR_USTAR_OFFSET: Final = 257 

55 

56_UTF8_BOM: Final = b"\xef\xbb\xbf" 

57 

58 

59class DetectedFormat(str, Enum): 

60 PDF = "pdf" 

61 TEXT = "text" 

62 

63 

64class DisallowedKind(str, Enum): 

65 ARCHIVE = "archive" 

66 EXECUTABLE = "executable" 

67 UNKNOWN_BINARY = "unknown_binary" 

68 

69 

70class RejectionReason(str, Enum): 

71 EMPTY_FILE = "empty_file" 

72 FILE_TOO_LARGE = "file_too_large" 

73 ARCHIVE_NOT_ALLOWED = "archive_not_allowed" 

74 EXECUTABLE_NOT_ALLOWED = "executable_not_allowed" 

75 UNSUPPORTED_FORMAT = "unsupported_format" 

76 MALWARE_DETECTED = "malware_detected" 

77 MALWARE_SCAN_ERROR = "malware_scan_error" 

78 

79 

80class ScanVerdict(str, Enum): 

81 CLEAN = "clean" 

82 INFECTED = "infected" 

83 ERROR = "error" 

84 

85 

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

87class ScanResult: 

88 verdict: ScanVerdict 

89 signature: str | None = None 

90 

91 

92@runtime_checkable 

93class MalwareScanner(Protocol): 

94 def scan(self, content: bytes) -> ScanResult: ... 94 ↛ exitline 94 didn't return from function 'scan' because

95 

96 

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

98class EicarTestMalwareScanner: 

99 """Placeholder scanner that only flags the EICAR anti-malware test file. 

100 

101 It exists to prove the scan hook is wired end to end and to satisfy the 

102 EICAR retest; it provides no real protection. Inject a scanner backed by a 

103 real engine through the ``scanner`` parameter of :func:`validate_upload` to 

104 screen production uploads. 

105 """ 

106 

107 def scan(self, content: bytes) -> ScanResult: 

108 if EICAR_TEST_SIGNATURE in content: 

109 return ScanResult(verdict=ScanVerdict.INFECTED, signature="EICAR-STANDARD-ANTIVIRUS-TEST-FILE") 

110 return ScanResult(verdict=ScanVerdict.CLEAN) 

111 

112 

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

114class AllowedContent: 

115 format: DetectedFormat 

116 

117 

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

119class DisallowedContent: 

120 kind: DisallowedKind 

121 

122 

123ContentInspection: TypeAlias = AllowedContent | DisallowedContent 

124 

125 

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

127class SecuredUpload: 

128 safe_filename: str 

129 content_type: str 

130 detected_format: DetectedFormat 

131 size_bytes: int 

132 

133 

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

135class RejectedUpload: 

136 reason: RejectionReason 

137 message: str 

138 

139 

140UploadValidation: TypeAlias = SecuredUpload | RejectedUpload 

141 

142_SAFE_EXTENSION: Final[Mapping[DetectedFormat, str]] = MappingProxyType( 

143 { 

144 DetectedFormat.PDF: "pdf", 

145 DetectedFormat.TEXT: "txt", 

146 } 

147) 

148 

149_SAFE_CONTENT_TYPE: Final[Mapping[DetectedFormat, str]] = MappingProxyType( 

150 { 

151 DetectedFormat.PDF: "application/pdf", 

152 DetectedFormat.TEXT: "text/plain", 

153 } 

154) 

155 

156 

157def _starts_with_any(content: bytes, prefixes: tuple[bytes, ...]) -> bool: 

158 return any(content.startswith(prefix) for prefix in prefixes) 

159 

160 

161def _is_archive(content: bytes) -> bool: 

162 if _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES): 

163 return True 

164 tar_magic_end: Final = _TAR_USTAR_OFFSET + len(_TAR_USTAR_MAGIC) 

165 if len(content) >= tar_magic_end and content[_TAR_USTAR_OFFSET:tar_magic_end] == _TAR_USTAR_MAGIC: 

166 return True 

167 return _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content) 

168 

169 

170def _is_utf8_text(content: bytes) -> bool: 

171 if b"\x00" in content: 

172 return False 

173 try: 

174 content.decode("utf-8") 

175 except UnicodeDecodeError: 

176 return False 

177 return True 

178 

179 

180def _is_executable_binary(content: bytes) -> bool: 

181 if _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES): 

182 return True 

183 return _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content) 

184 

185 

186def _looks_like_shebang(content: bytes) -> bool: 

187 body: Final = content.removeprefix(_UTF8_BOM).lstrip() 

188 return body.startswith(b"#!") 

189 

190 

191def inspect_content(content: bytes) -> ContentInspection: 

192 if _looks_like_shebang(content): 

193 return DisallowedContent(DisallowedKind.EXECUTABLE) 

194 if content.startswith(b"%PDF-"): 

195 return AllowedContent(DetectedFormat.PDF) 

196 if _is_archive(content): 

197 return DisallowedContent(DisallowedKind.ARCHIVE) 

198 if _is_executable_binary(content): 

199 return DisallowedContent(DisallowedKind.EXECUTABLE) 

200 if _is_utf8_text(content): 

201 return AllowedContent(DetectedFormat.TEXT) 

202 return DisallowedContent(DisallowedKind.UNKNOWN_BINARY) 

203 

204 

205def generate_safe_filename(detected_format: DetectedFormat) -> str: 

206 return f"{uuid.uuid4().hex}.{_SAFE_EXTENSION[detected_format]}" 

207 

208 

209def _reject_disallowed(kind: DisallowedKind) -> RejectedUpload: 

210 match kind: 

211 case DisallowedKind.ARCHIVE: 

212 return RejectedUpload( 

213 RejectionReason.ARCHIVE_NOT_ALLOWED, 

214 "Archive uploads are not allowed.", 

215 ) 

216 case DisallowedKind.EXECUTABLE: 

217 return RejectedUpload( 

218 RejectionReason.EXECUTABLE_NOT_ALLOWED, 

219 "Executable uploads are not allowed.", 

220 ) 

221 case DisallowedKind.UNKNOWN_BINARY: 

222 return RejectedUpload( 

223 RejectionReason.UNSUPPORTED_FORMAT, 

224 "Only PDF and UTF-8 text documents are accepted.", 

225 ) 

226 assert_never(kind) 

227 

228 

229def _scan_rejection(content: bytes, scanner: MalwareScanner) -> RejectedUpload | None: 

230 result: Final = scanner.scan(content) 

231 match result.verdict: 

232 case ScanVerdict.CLEAN: 

233 return None 

234 case ScanVerdict.INFECTED: 

235 return RejectedUpload( 

236 RejectionReason.MALWARE_DETECTED, 

237 f"Uploaded file was flagged by malware scanning ({result.signature or 'unknown signature'}).", 

238 ) 

239 case ScanVerdict.ERROR: 

240 return RejectedUpload( 

241 RejectionReason.MALWARE_SCAN_ERROR, 

242 "Malware scanning could not complete; upload rejected.", 

243 ) 

244 assert_never(result.verdict) 

245 

246 

247def validate_upload( 

248 *, 

249 content: bytes, 

250 scanner: MalwareScanner, 

251 max_size_bytes: int = MAX_UPLOAD_SIZE_BYTES, 

252) -> UploadValidation: 

253 size: Final = len(content) 

254 if size == 0: 

255 return RejectedUpload(RejectionReason.EMPTY_FILE, "Uploaded file is empty.") 

256 if size > max_size_bytes: 

257 return RejectedUpload( 

258 RejectionReason.FILE_TOO_LARGE, 

259 f"Uploaded file is {size} bytes, exceeding the {max_size_bytes}-byte limit.", 

260 ) 

261 

262 inspection: Final = inspect_content(content) 

263 if isinstance(inspection, DisallowedContent): 

264 return _reject_disallowed(inspection.kind) 

265 

266 scan_rejection: Final = _scan_rejection(content, scanner) 

267 if scan_rejection is not None: 

268 return scan_rejection 

269 

270 return SecuredUpload( 

271 safe_filename=generate_safe_filename(inspection.format), 

272 content_type=_SAFE_CONTENT_TYPE[inspection.format], 

273 detected_format=inspection.format, 

274 size_bytes=size, 

275 ) 

276 

277 

278def _sanitize_header_filename(filename: str) -> str: 

279 stripped: Final = "".join(char for char in filename if char not in '"\\\r\n').strip() 

280 return stripped or "download" 

281 

282 

283def safe_download_headers(filename: str) -> Mapping[str, str]: 

284 return MappingProxyType( 

285 { 

286 "Content-Disposition": f'attachment; filename="{_sanitize_header_filename(filename)}"', 

287 "X-Content-Type-Options": "nosniff", 

288 } 

289 )