Coverage for integrations / channels / media / limits.py: 43.6%
140 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-12 04:49 +0000
1"""
2Media Limiter for size/type limits.
4Provides configuration and checking of media limits.
5"""
7from dataclasses import dataclass, field
8from enum import Enum
9from typing import Optional, Dict, Any, List, Set
10import logging
12logger = logging.getLogger(__name__)
15class MediaType(Enum):
16 """Types of media."""
17 IMAGE = "image"
18 VIDEO = "video"
19 AUDIO = "audio"
20 DOCUMENT = "document"
21 ARCHIVE = "archive"
22 OTHER = "other"
25@dataclass
26class MediaLimits:
27 """
28 Media limits configuration.
30 Defines maximum sizes and allowed types for different media.
31 """
32 # Maximum file sizes in bytes
33 max_image_size: int = 10 * 1024 * 1024 # 10MB
34 max_video_size: int = 50 * 1024 * 1024 # 50MB
35 max_audio_size: int = 25 * 1024 * 1024 # 25MB
36 max_document_size: int = 25 * 1024 * 1024 # 25MB
37 max_archive_size: int = 50 * 1024 * 1024 # 50MB
38 max_other_size: int = 10 * 1024 * 1024 # 10MB
40 # Maximum dimensions for images/video
41 max_image_width: int = 4096
42 max_image_height: int = 4096
43 max_video_width: int = 1920
44 max_video_height: int = 1080
46 # Duration limits in seconds
47 max_video_duration: int = 600 # 10 minutes
48 max_audio_duration: int = 3600 # 1 hour
50 # Allowed file extensions (empty = all allowed)
51 allowed_image_extensions: Set[str] = field(default_factory=lambda: {
52 'jpg', 'jpeg', 'png', 'gif', 'webp', 'bmp', 'svg'
53 })
54 allowed_video_extensions: Set[str] = field(default_factory=lambda: {
55 'mp4', 'webm', 'avi', 'mov', 'mkv', 'm4v'
56 })
57 allowed_audio_extensions: Set[str] = field(default_factory=lambda: {
58 'mp3', 'wav', 'ogg', 'flac', 'm4a', 'aac', 'wma'
59 })
60 allowed_document_extensions: Set[str] = field(default_factory=lambda: {
61 'pdf', 'doc', 'docx', 'xls', 'xlsx', 'ppt', 'pptx',
62 'txt', 'csv', 'json', 'xml', 'md', 'rtf'
63 })
64 allowed_archive_extensions: Set[str] = field(default_factory=lambda: {
65 'zip', 'tar', 'gz', 'rar', '7z', 'bz2'
66 })
68 # Blocked extensions (always blocked regardless of allowed)
69 blocked_extensions: Set[str] = field(default_factory=lambda: {
70 'exe', 'dll', 'bat', 'cmd', 'sh', 'ps1', 'vbs',
71 'scr', 'msi', 'jar', 'app', 'dmg'
72 })
74 # MIME type restrictions
75 allowed_mime_types: Set[str] = field(default_factory=set)
76 blocked_mime_types: Set[str] = field(default_factory=lambda: {
77 'application/x-executable',
78 'application/x-msdownload',
79 'application/x-sh'
80 })
82 # Additional constraints
83 max_files_per_message: int = 10
84 max_total_size: int = 100 * 1024 * 1024 # 100MB total
86 def to_dict(self) -> Dict[str, Any]:
87 """Convert limits to dictionary."""
88 return {
89 "max_image_size": self.max_image_size,
90 "max_video_size": self.max_video_size,
91 "max_audio_size": self.max_audio_size,
92 "max_document_size": self.max_document_size,
93 "max_archive_size": self.max_archive_size,
94 "max_other_size": self.max_other_size,
95 "max_image_width": self.max_image_width,
96 "max_image_height": self.max_image_height,
97 "max_video_width": self.max_video_width,
98 "max_video_height": self.max_video_height,
99 "max_video_duration": self.max_video_duration,
100 "max_audio_duration": self.max_audio_duration,
101 "allowed_image_extensions": list(self.allowed_image_extensions),
102 "allowed_video_extensions": list(self.allowed_video_extensions),
103 "allowed_audio_extensions": list(self.allowed_audio_extensions),
104 "allowed_document_extensions": list(self.allowed_document_extensions),
105 "allowed_archive_extensions": list(self.allowed_archive_extensions),
106 "blocked_extensions": list(self.blocked_extensions),
107 "max_files_per_message": self.max_files_per_message,
108 "max_total_size": self.max_total_size
109 }
111 @classmethod
112 def from_dict(cls, data: Dict[str, Any]) -> 'MediaLimits':
113 """Create MediaLimits from dictionary."""
114 # Convert lists back to sets for extensions
115 for key in ['allowed_image_extensions', 'allowed_video_extensions',
116 'allowed_audio_extensions', 'allowed_document_extensions',
117 'allowed_archive_extensions', 'blocked_extensions',
118 'allowed_mime_types', 'blocked_mime_types']:
119 if key in data and isinstance(data[key], list):
120 data[key] = set(data[key])
121 return cls(**data)
124@dataclass
125class LimitCheckResult:
126 """Result of a limit check."""
127 allowed: bool
128 reason: Optional[str] = None
129 media_type: Optional[MediaType] = None
130 details: Dict[str, Any] = field(default_factory=dict)
132 def to_dict(self) -> Dict[str, Any]:
133 return {
134 "allowed": self.allowed,
135 "reason": self.reason,
136 "media_type": self.media_type.value if self.media_type else None,
137 "details": self.details
138 }
141class MediaLimiter:
142 """
143 Media limiter for checking size/type limits.
145 Validates media against configured limits.
146 """
148 # Extension to media type mapping
149 EXTENSION_TYPE_MAP = {
150 # Images
151 'jpg': MediaType.IMAGE, 'jpeg': MediaType.IMAGE, 'png': MediaType.IMAGE,
152 'gif': MediaType.IMAGE, 'webp': MediaType.IMAGE, 'bmp': MediaType.IMAGE,
153 'svg': MediaType.IMAGE, 'ico': MediaType.IMAGE, 'tiff': MediaType.IMAGE,
154 # Video
155 'mp4': MediaType.VIDEO, 'webm': MediaType.VIDEO, 'avi': MediaType.VIDEO,
156 'mov': MediaType.VIDEO, 'mkv': MediaType.VIDEO, 'm4v': MediaType.VIDEO,
157 'wmv': MediaType.VIDEO, 'flv': MediaType.VIDEO,
158 # Audio
159 'mp3': MediaType.AUDIO, 'wav': MediaType.AUDIO, 'ogg': MediaType.AUDIO,
160 'flac': MediaType.AUDIO, 'm4a': MediaType.AUDIO, 'aac': MediaType.AUDIO,
161 'wma': MediaType.AUDIO, 'opus': MediaType.AUDIO,
162 # Documents
163 'pdf': MediaType.DOCUMENT, 'doc': MediaType.DOCUMENT, 'docx': MediaType.DOCUMENT,
164 'xls': MediaType.DOCUMENT, 'xlsx': MediaType.DOCUMENT, 'ppt': MediaType.DOCUMENT,
165 'pptx': MediaType.DOCUMENT, 'txt': MediaType.DOCUMENT, 'csv': MediaType.DOCUMENT,
166 'json': MediaType.DOCUMENT, 'xml': MediaType.DOCUMENT, 'md': MediaType.DOCUMENT,
167 'rtf': MediaType.DOCUMENT, 'odt': MediaType.DOCUMENT,
168 # Archives
169 'zip': MediaType.ARCHIVE, 'tar': MediaType.ARCHIVE, 'gz': MediaType.ARCHIVE,
170 'rar': MediaType.ARCHIVE, '7z': MediaType.ARCHIVE, 'bz2': MediaType.ARCHIVE
171 }
173 # MIME type to media type mapping
174 MIME_TYPE_MAP = {
175 'image/': MediaType.IMAGE,
176 'video/': MediaType.VIDEO,
177 'audio/': MediaType.AUDIO,
178 'application/pdf': MediaType.DOCUMENT,
179 'application/msword': MediaType.DOCUMENT,
180 'application/vnd.': MediaType.DOCUMENT,
181 'text/': MediaType.DOCUMENT,
182 'application/zip': MediaType.ARCHIVE,
183 'application/x-rar': MediaType.ARCHIVE,
184 'application/x-7z': MediaType.ARCHIVE,
185 'application/gzip': MediaType.ARCHIVE
186 }
188 def __init__(self, limits: Optional[MediaLimits] = None):
189 """
190 Initialize media limiter.
192 Args:
193 limits: Media limits configuration (uses defaults if not provided)
194 """
195 self._limits = limits or MediaLimits()
197 def check(
198 self,
199 filename: Optional[str] = None,
200 size: Optional[int] = None,
201 mime_type: Optional[str] = None,
202 width: Optional[int] = None,
203 height: Optional[int] = None,
204 duration: Optional[float] = None
205 ) -> LimitCheckResult:
206 """
207 Check if media meets limits.
209 Args:
210 filename: File name (for extension check)
211 size: File size in bytes
212 mime_type: MIME type of the file
213 width: Width (for images/video)
214 height: Height (for images/video)
215 duration: Duration in seconds (for audio/video)
217 Returns:
218 LimitCheckResult indicating if media is allowed
219 """
220 # Determine media type
221 media_type = self._get_media_type(filename, mime_type)
223 # Check extension
224 if filename:
225 ext = self._get_extension(filename)
227 # Check if blocked
228 if ext in self._limits.blocked_extensions:
229 return LimitCheckResult(
230 allowed=False,
231 reason=f"File extension '{ext}' is blocked",
232 media_type=media_type,
233 details={"extension": ext}
234 )
236 # Check if allowed for media type
237 allowed_exts = self._get_allowed_extensions(media_type)
238 if allowed_exts and ext not in allowed_exts:
239 return LimitCheckResult(
240 allowed=False,
241 reason=f"File extension '{ext}' not allowed for {media_type.value}",
242 media_type=media_type,
243 details={"extension": ext, "allowed": list(allowed_exts)}
244 )
246 # Check MIME type
247 if mime_type:
248 if mime_type in self._limits.blocked_mime_types:
249 return LimitCheckResult(
250 allowed=False,
251 reason=f"MIME type '{mime_type}' is blocked",
252 media_type=media_type,
253 details={"mime_type": mime_type}
254 )
256 if self._limits.allowed_mime_types:
257 if mime_type not in self._limits.allowed_mime_types:
258 return LimitCheckResult(
259 allowed=False,
260 reason=f"MIME type '{mime_type}' not in allowed list",
261 media_type=media_type,
262 details={"mime_type": mime_type}
263 )
265 # Check size
266 if size is not None:
267 max_size = self._get_max_size(media_type)
268 if size > max_size:
269 return LimitCheckResult(
270 allowed=False,
271 reason=f"File size ({size} bytes) exceeds limit ({max_size} bytes)",
272 media_type=media_type,
273 details={"size": size, "max_size": max_size}
274 )
276 # Check dimensions
277 if media_type == MediaType.IMAGE:
278 if width and width > self._limits.max_image_width:
279 return LimitCheckResult(
280 allowed=False,
281 reason=f"Image width ({width}px) exceeds limit ({self._limits.max_image_width}px)",
282 media_type=media_type,
283 details={"width": width, "max_width": self._limits.max_image_width}
284 )
285 if height and height > self._limits.max_image_height:
286 return LimitCheckResult(
287 allowed=False,
288 reason=f"Image height ({height}px) exceeds limit ({self._limits.max_image_height}px)",
289 media_type=media_type,
290 details={"height": height, "max_height": self._limits.max_image_height}
291 )
293 if media_type == MediaType.VIDEO:
294 if width and width > self._limits.max_video_width:
295 return LimitCheckResult(
296 allowed=False,
297 reason=f"Video width ({width}px) exceeds limit ({self._limits.max_video_width}px)",
298 media_type=media_type,
299 details={"width": width, "max_width": self._limits.max_video_width}
300 )
301 if height and height > self._limits.max_video_height:
302 return LimitCheckResult(
303 allowed=False,
304 reason=f"Video height ({height}px) exceeds limit ({self._limits.max_video_height}px)",
305 media_type=media_type,
306 details={"height": height, "max_height": self._limits.max_video_height}
307 )
309 # Check duration
310 if duration is not None:
311 if media_type == MediaType.VIDEO and duration > self._limits.max_video_duration:
312 return LimitCheckResult(
313 allowed=False,
314 reason=f"Video duration ({duration}s) exceeds limit ({self._limits.max_video_duration}s)",
315 media_type=media_type,
316 details={"duration": duration, "max_duration": self._limits.max_video_duration}
317 )
318 if media_type == MediaType.AUDIO and duration > self._limits.max_audio_duration:
319 return LimitCheckResult(
320 allowed=False,
321 reason=f"Audio duration ({duration}s) exceeds limit ({self._limits.max_audio_duration}s)",
322 media_type=media_type,
323 details={"duration": duration, "max_duration": self._limits.max_audio_duration}
324 )
326 return LimitCheckResult(
327 allowed=True,
328 media_type=media_type
329 )
331 def check_batch(
332 self,
333 files: List[Dict[str, Any]]
334 ) -> LimitCheckResult:
335 """
336 Check a batch of files against limits.
338 Args:
339 files: List of file info dicts with keys: filename, size, mime_type, etc.
341 Returns:
342 LimitCheckResult for the batch
343 """
344 # Check file count
345 if len(files) > self._limits.max_files_per_message:
346 return LimitCheckResult(
347 allowed=False,
348 reason=f"Too many files ({len(files)}), max is {self._limits.max_files_per_message}",
349 details={"count": len(files), "max_count": self._limits.max_files_per_message}
350 )
352 # Check total size
353 total_size = sum(f.get('size', 0) for f in files)
354 if total_size > self._limits.max_total_size:
355 return LimitCheckResult(
356 allowed=False,
357 reason=f"Total size ({total_size} bytes) exceeds limit ({self._limits.max_total_size} bytes)",
358 details={"total_size": total_size, "max_total_size": self._limits.max_total_size}
359 )
361 # Check each file
362 for i, file_info in enumerate(files):
363 result = self.check(**file_info)
364 if not result.allowed:
365 result.details["file_index"] = i
366 return result
368 return LimitCheckResult(allowed=True)
370 def get_limits(self) -> MediaLimits:
371 """Get current limits."""
372 return self._limits
374 def set_limits(self, limits: MediaLimits):
375 """Set new limits."""
376 self._limits = limits
378 def update_limits(self, **kwargs):
379 """Update specific limit values."""
380 for key, value in kwargs.items():
381 if hasattr(self._limits, key):
382 setattr(self._limits, key, value)
384 def _get_extension(self, filename: str) -> str:
385 """Extract lowercase extension from filename."""
386 if '.' in filename:
387 return filename.rsplit('.', 1)[-1].lower()
388 return ''
390 def _get_media_type(
391 self,
392 filename: Optional[str],
393 mime_type: Optional[str]
394 ) -> MediaType:
395 """Determine media type from filename or MIME type."""
396 # Try extension first
397 if filename:
398 ext = self._get_extension(filename)
399 if ext in self.EXTENSION_TYPE_MAP:
400 return self.EXTENSION_TYPE_MAP[ext]
402 # Try MIME type
403 if mime_type:
404 for prefix, media_type in self.MIME_TYPE_MAP.items():
405 if mime_type.startswith(prefix):
406 return media_type
408 return MediaType.OTHER
410 def _get_max_size(self, media_type: MediaType) -> int:
411 """Get maximum size for media type."""
412 size_map = {
413 MediaType.IMAGE: self._limits.max_image_size,
414 MediaType.VIDEO: self._limits.max_video_size,
415 MediaType.AUDIO: self._limits.max_audio_size,
416 MediaType.DOCUMENT: self._limits.max_document_size,
417 MediaType.ARCHIVE: self._limits.max_archive_size,
418 MediaType.OTHER: self._limits.max_other_size
419 }
420 return size_map.get(media_type, self._limits.max_other_size)
422 def _get_allowed_extensions(self, media_type: MediaType) -> Set[str]:
423 """Get allowed extensions for media type."""
424 ext_map = {
425 MediaType.IMAGE: self._limits.allowed_image_extensions,
426 MediaType.VIDEO: self._limits.allowed_video_extensions,
427 MediaType.AUDIO: self._limits.allowed_audio_extensions,
428 MediaType.DOCUMENT: self._limits.allowed_document_extensions,
429 MediaType.ARCHIVE: self._limits.allowed_archive_extensions
430 }
431 return ext_map.get(media_type, set())
433 def format_size(self, size: int) -> str:
434 """Format size in human-readable format."""
435 for unit in ['B', 'KB', 'MB', 'GB']:
436 if size < 1024:
437 return f"{size:.1f} {unit}"
438 size /= 1024
439 return f"{size:.1f} TB"
441 def get_limits_summary(self) -> Dict[str, str]:
442 """Get human-readable limits summary."""
443 return {
444 "image": f"Max {self.format_size(self._limits.max_image_size)}, {self._limits.max_image_width}x{self._limits.max_image_height}px",
445 "video": f"Max {self.format_size(self._limits.max_video_size)}, {self._limits.max_video_duration}s, {self._limits.max_video_width}x{self._limits.max_video_height}px",
446 "audio": f"Max {self.format_size(self._limits.max_audio_size)}, {self._limits.max_audio_duration}s",
447 "document": f"Max {self.format_size(self._limits.max_document_size)}",
448 "total": f"Max {self._limits.max_files_per_message} files, {self.format_size(self._limits.max_total_size)} total"
449 }