Coverage for src/local_deep_research/security/filename_sanitizer.py: 99%

1"""Filename sanitization for file uploads. 

2 

3Wraps werkzeug's secure_filename with additional safety checks. 

4All file upload endpoints should use sanitize_filename() instead of 

5importing secure_filename directly. 

6""" 

7 

8from __future__ import annotations 

9 

10import hashlib 

11from typing import Optional 

12 

13from werkzeug.utils import secure_filename 

14 

15# Maximum filename length (including extension) 

16MAX_FILENAME_LENGTH = 255 

17 

18 

19class UnsafeFilenameError(ValueError): 

20 """Raised when a filename cannot be sanitized to a safe value.""" 

21 

22 

23def sanitize_filename( 

24 filename: Optional[str], 

25 *, 

26 allowed_extensions: Optional[set[str]] = None, 

27 max_length: int = MAX_FILENAME_LENGTH, 

28) -> str: 

29 """Sanitize an uploaded filename for safe filesystem storage. 

30 

31 Args: 

32 filename: Raw filename from the upload. 

33 allowed_extensions: Optional set of allowed extensions 

34 (lowercase, with dot, e.g. {".pdf", ".txt"}). 

35 If None, all extensions are allowed. 

36 max_length: Maximum allowed filename length. 

37 

38 Returns: 

39 Sanitized filename safe for filesystem use. If the stem (the 

40 part before the last dot) sanitizes to nothing but still 

41 contains letters or digits — as happens for non-Latin scripts 

42 such as CJK or Cyrillic, which werkzeug's secure_filename() 

43 drops entirely — the stem is replaced with a deterministic 

44 ``upload-<12 hex chars>`` name derived from a hash of the 

45 original filename, and the sanitized extension is preserved. 

46 

47 Raises: 

48 UnsafeFilenameError: If the filename is empty, becomes empty 

49 after sanitization, has a disallowed extension, 

50 ``max_length`` is not positive, or the extension alone 

51 exceeds ``max_length`` (leaving no room for a stem). 

52 """ 

53 if not filename: 

54 raise UnsafeFilenameError("No filename provided") 

55 if max_length < 1: 

56 raise UnsafeFilenameError("Maximum filename length must be positive") 

57 

58 # Strip null bytes before passing to secure_filename 

59 cleaned = filename.replace("\x00", "") 

60 

61 # Apply werkzeug's path traversal protection 

62 safe_name = secure_filename(cleaned) 

63 

64 raw_stem, separator, raw_extension = cleaned.rpartition(".") 

65 if not separator: 

66 raw_stem = cleaned 

67 

68 if not secure_filename(raw_stem) and any( 

69 character.isalnum() for character in raw_stem 

70 ): 

71 extension = "" 

72 if separator: 72 ↛ 80line 72 didn't jump to line 80 because the condition on line 72 was always true

73 extension_source = secure_filename(f"x.{raw_extension}") 

74 extension_index = extension_source.rfind(".") 

75 if extension_index > 0: 

76 extension = extension_source[extension_index:] 

77 

78 # A surrogate can occur in a filesystem-derived or decoded name. 

79 # It only contributes to the digest; the generated name stays ASCII. 

80 digest = hashlib.sha256( 

81 cleaned.encode("utf-8", errors="surrogatepass") 

82 ).hexdigest()[:12] 

83 safe_name = f"upload-{digest}{extension}" 

84 

85 if not safe_name: 

86 raise UnsafeFilenameError( 

87 "Filename contains no safe characters after sanitization" 

88 ) 

89 

90 # Enforce length limit 

91 if len(safe_name) > max_length: 

92 # Preserve extension when truncating 

93 dot_idx = safe_name.rfind(".") 

94 if dot_idx > 0: 

95 ext = safe_name[dot_idx:] 

96 stem_length = max_length - len(ext) 

97 if stem_length < 1: 

98 raise UnsafeFilenameError( 

99 "Filename extension exceeds maximum length" 

100 ) 

101 safe_name = safe_name[:stem_length] + ext 

102 else: 

103 safe_name = safe_name[:max_length] 

104 

105 # Validate extension if allowlist provided 

106 if allowed_extensions is not None: 

107 dot_idx = safe_name.rfind(".") 

108 ext = safe_name[dot_idx:].lower() if dot_idx > 0 else "" 

109 # Normalize allowlist for case-insensitive comparison 

110 normalized = {e.lower() for e in allowed_extensions} 

111 if ext not in normalized: 

112 raise UnsafeFilenameError("File type not allowed") 

113 

114 return safe_name