from datetime import timedelta from pathlib import Path from typing import Optional from flask import abort, jsonify, send_from_directory from db import get_db IMAGE_MIME_PREFIX = "image/" AUDIO_MIME_PREFIX = "audio/" VIDEO_MIME_PREFIX = "video/" IMAGE_EXTENSIONS = {"png", "jpg", "jpeg", "gif", "webp"} AUDIO_EXTENSIONS = {"mp3", "wav", "ogg", "m4a", "aac", "flac"} VIDEO_EXTENSIONS = {"mp4", "webm", "mov", "m4v", "ogv"} def attachment_kind_from_name_and_mime( original_name: Optional[str], mime_type: Optional[str], ) -> str: mime_type = str(mime_type or "").lower() name = str(original_name or "") ext = name.rsplit(".", 1)[-1].lower() if "." in name else "" if mime_type.startswith(IMAGE_MIME_PREFIX) or ext in IMAGE_EXTENSIONS: return "image" if mime_type.startswith(AUDIO_MIME_PREFIX) or ext in AUDIO_EXTENSIONS: return "audio" if mime_type.startswith(VIDEO_MIME_PREFIX) or ext in VIDEO_EXTENSIONS: return "video" return "file" def attachment_payload_from_row(row) -> Optional[dict]: if not row or not row["id"]: return None kind = attachment_kind_from_name_and_mime( row["original_name"], row["mime_type"], ) return { "id": row["id"], "original_name": row["original_name"], "mime_type": row["mime_type"], "size_bytes": row["size_bytes"], "kind": kind, "url": f"/api/attachments/{row['id']}/download", } def attachment_payload_legacy_from_message_select_row(row) -> Optional[dict]: if not row["attachment_id"]: return None kind = attachment_kind_from_name_and_mime( row["attachment_original_name"], row["attachment_mime_type"], ) return { "id": row["attachment_id"], "original_name": row["attachment_original_name"], "mime_type": row["attachment_mime_type"], "size_bytes": row["attachment_size_bytes"], "kind": kind, "url": f"/api/attachments/{row['attachment_id']}/download", } def attachment_list_for_message(message_id: int) -> list[dict]: with get_db() as db: rows = db.execute( """ SELECT a.id, a.original_name, a.mime_type, a.size_bytes FROM message_attachments AS ma JOIN attachments AS a ON a.id = ma.attachment_id WHERE ma.message_id = ? ORDER BY ma.position ASC, a.id ASC """, (message_id,), ).fetchall() return [attachment_payload_from_row(row) for row in rows] def cleanup_orphaned_attachments( upload_dir: Path, utc_now_func, orphan_ttl_hours: int, app_logger, ) -> None: cutoff = (utc_now_func() - timedelta(hours=orphan_ttl_hours)).isoformat() with get_db() as db: orphaned = db.execute( """ SELECT a.id, a.stored_name FROM attachments AS a WHERE a.created_at < ? AND NOT EXISTS ( SELECT 1 FROM messages AS m WHERE m.attachment_id = a.id ) AND NOT EXISTS ( SELECT 1 FROM message_attachments AS ma WHERE ma.attachment_id = a.id ) """, (cutoff,), ).fetchall() for attachment in orphaned: db.execute( """ DELETE FROM attachments WHERE id = ? """, (attachment["id"],), ) db.commit() for attachment in orphaned: try: (upload_dir / attachment["stored_name"]).unlink(missing_ok=True) except OSError: app_logger.warning( "Could not remove orphaned attachment file: %s", attachment["stored_name"], ) def maybe_cleanup_orphaned_attachments( state: dict, cleanup_interval_seconds: int, cleanup_func, app_logger, ) -> None: import time now_value = time.time() last = state.get("last_attachment_cleanup", 0.0) if now_value - last < cleanup_interval_seconds: return state["last_attachment_cleanup"] = now_value try: cleanup_func() except Exception: app_logger.exception("Could not clean orphaned attachments.") def resolve_owned_attachment_ids( attachment_ids: object, user_id: int, ) -> tuple[Optional[list[int]], Optional[tuple]]: if attachment_ids is None: return [], None if not isinstance(attachment_ids, list): return None, (jsonify({"error": "attachment_ids must be an array."}), 400) parsed = [] seen = set() for item in attachment_ids[:10]: try: attachment_id = int(item) except (TypeError, ValueError): return None, (jsonify({"error": "Invalid attachment."}), 400) if attachment_id in seen: continue seen.add(attachment_id) with get_db() as db: row = db.execute( """ SELECT id, uploader_user_id FROM attachments WHERE id = ? """, (attachment_id,), ).fetchone() if not row: return None, (jsonify({"error": "Attachment not found."}), 404) if row["uploader_user_id"] != user_id: return None, ( jsonify({"error": "You can only attach files you uploaded."}), 403, ) parsed.append(attachment_id) return parsed, None def resolve_owned_attachment( attachment_id: object, user_id: int, ) -> tuple[Optional[int], Optional[tuple]]: if attachment_id is None: return None, None try: parsed_id = int(attachment_id) except (TypeError, ValueError): return None, (jsonify({"error": "Invalid attachment."}), 400) with get_db() as db: attachment = db.execute( """ SELECT id, uploader_user_id FROM attachments WHERE id = ? """, (parsed_id,), ).fetchone() if not attachment: return None, (jsonify({"error": "Attachment not found."}), 404) if attachment["uploader_user_id"] != user_id: return None, ( jsonify({"error": "You can only attach files you uploaded."}), 403, ) return parsed_id, None def can_download_attachment( attachment_id: int, user: Optional[dict], can_access_room, can_access_channel, get_channel_by_id, ) -> bool: if not user: return False with get_db() as db: room_rows = db.execute( """ SELECT DISTINCT r.id, r.slug, r.name, r.description, r.visibility, r.created_by_user_id, r.is_active FROM message_attachments AS ma JOIN messages AS m ON m.id = ma.message_id JOIN rooms AS r ON r.id = m.room_id WHERE ma.attachment_id = ? AND r.is_active = 1 """, (attachment_id,), ).fetchall() channel_rows = db.execute( """ SELECT DISTINCT m.channel_id FROM message_attachments AS ma JOIN messages AS m ON m.id = ma.message_id WHERE ma.attachment_id = ? AND m.channel_id IS NOT NULL """, (attachment_id,), ).fetchall() dm_rows = db.execute( """ SELECT DISTINCT dc.id FROM message_attachments AS ma JOIN messages AS m ON m.id = ma.message_id JOIN direct_conversations AS dc ON dc.id = m.direct_conversation_id JOIN direct_conversation_members AS dcm ON dcm.direct_conversation_id = dc.id WHERE ma.attachment_id = ? AND dcm.user_id = ? """, (attachment_id, user["id"]), ).fetchall() if dm_rows: return True if any(can_access_room(room, user) for room in room_rows): return True for item in channel_rows: channel = get_channel_by_id(item["channel_id"]) if channel and can_access_channel(channel, user): return True return False def send_attachment_file( upload_dir: Path, attachment_id: int, user: Optional[dict], can_access_room, can_access_channel, get_channel_by_id, ): with get_db() as db: attachment = db.execute( """ SELECT id, stored_name, original_name, mime_type FROM attachments WHERE id = ? """, (attachment_id,), ).fetchone() if not attachment: abort(404) if not can_download_attachment( attachment_id, user, can_access_room, can_access_channel, get_channel_by_id, ): abort(403) return send_from_directory( upload_dir, attachment["stored_name"], as_attachment=True, download_name=attachment["original_name"], mimetype=attachment["mime_type"], )