| 1 | import os |
| 2 | import io |
| 3 | import base64 |
| 4 | from PIL import Image |
| 5 | from typing import Dict, List, Optional, Tuple |
| 6 | from helpers.security import safe_filename |
| 7 | from werkzeug.datastructures import FileStorage |
| 8 | |
| 9 | from helpers.print_style import PrintStyle |
| 10 | |
| 11 | class AttachmentManager: |
| 12 | ALLOWED_EXTENSIONS = { |
| 13 | 'image': {'jpg', 'jpeg', 'png', 'bmp'}, |
| 14 | 'code': {'py', 'js', 'sh', 'html', 'css'}, |
| 15 | 'document': {'md', 'pdf', 'txt', 'csv', 'json'} |
| 16 | } |
| 17 | |
| 18 | def __init__(self, work_dir: str): |
| 19 | self.work_dir = work_dir |
| 20 | os.makedirs(work_dir, exist_ok=True) |
| 21 | |
| 22 | def is_allowed_file(self, filename: str) -> bool: |
| 23 | ext = self.get_file_extension(filename) |
| 24 | all_allowed = set().union(*self.ALLOWED_EXTENSIONS.values()) |
| 25 | return ext in all_allowed |
| 26 | |
| 27 | def get_file_type(self, filename: str) -> str: |
| 28 | ext = self.get_file_extension(filename) |
| 29 | for file_type, extensions in self.ALLOWED_EXTENSIONS.items(): |
| 30 | if ext in extensions: |
| 31 | return file_type |
| 32 | return 'unknown' |
| 33 | |
| 34 | @staticmethod |
| 35 | def get_file_extension(filename: str) -> str: |
| 36 | return filename.rsplit('.', 1)[1].lower() if '.' in filename else '' |
| 37 | |
| 38 | def validate_mime_type(self, file) -> bool: |
| 39 | try: |
| 40 | mime_type = file.content_type |
| 41 | return mime_type.split('/')[0] in ['image', 'text', 'application'] |
| 42 | except AttributeError: |
| 43 | return False |
| 44 | |
| 45 | def save_file(self, file: FileStorage, name: str) -> Tuple[str, Dict]: |
| 46 | """Save file and return path and metadata""" |
| 47 | try: |
| 48 | filename = safe_filename(name) |
| 49 | if not filename: |
| 50 | raise ValueError("Invalid filename") |
| 51 | |
| 52 | file_path = os.path.join(self.work_dir, filename) |
| 53 | |
| 54 | file_type = self.get_file_type(filename) |
| 55 | metadata = { |
| 56 | 'filename': filename, |
| 57 | 'type': file_type, |
| 58 | 'extension': self.get_file_extension(filename), |
| 59 | 'preview': None |
| 60 | } |
| 61 | |
| 62 | # Save file |
| 63 | file.save(file_path) |
| 64 | |
| 65 | # Generate preview for images |
| 66 | if file_type == 'image': |
| 67 | metadata['preview'] = self.generate_image_preview(file_path) |
| 68 | |
| 69 | return file_path, metadata |
| 70 | |
| 71 | except Exception as e: |
| 72 | PrintStyle.error(f"Error saving file {name}: {e}") |
| 73 | return None, {} # type: ignore |
| 74 | |
| 75 | def generate_image_preview(self, image_path: str, max_size: int = 800) -> Optional[str]: |
| 76 | try: |
| 77 | with Image.open(image_path) as img: |
| 78 | # Convert image if needed |
| 79 | if img.mode in ('RGBA', 'P'): |
| 80 | img = img.convert('RGB') |
| 81 | |
| 82 | # Resize for preview |
| 83 | img.thumbnail((max_size, max_size)) |
| 84 | |
| 85 | # Save to buffer |
| 86 | buffer = io.BytesIO() |
| 87 | img.save(buffer, format="JPEG", quality=70, optimize=True) |
| 88 | |
| 89 | # Convert to base64 |
| 90 | return base64.b64encode(buffer.getvalue()).decode('utf-8') |
| 91 | except Exception as e: |
| 92 | PrintStyle.error(f"Error generating preview for {image_path}: {e}") |
| 93 | return None |
| 94 |