glo/attachments.py

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"],
    )