beavergram/backend/src/utils/files.py

import re
from datetime import UTC, datetime, timedelta
from urllib.parse import quote

_MIME_EXTENSIONS = {
    "application/pdf": ".pdf",
    "application/x-tgsticker": ".tgs",
    "application/zip": ".zip",
    "audio/mpeg": ".mp3",
    "audio/ogg": ".ogg",
    "image/gif": ".gif",
    "image/jpeg": ".jpg",
    "image/png": ".png",
    "image/webp": ".webp",
    "text/plain": ".txt",
    "video/mp4": ".mp4",
    "video/quicktime": ".mov",
    "video/webm": ".webm",
}

_KIND_EXTENSIONS = {
    "animation": ".mp4",
    "audio": ".mp3",
    "gif": ".mp4",
    "photo": ".jpg",
    "sticker": ".webp",
    "video": ".mp4",
    "video_note": ".mp4",
    "voice": ".ogg",
}

_EXTENSION_MIMES = {
    ".flac": "audio/flac",
    ".gif": "image/gif",
    ".heic": "image/heic",
    ".jpeg": "image/jpeg",
    ".jpg": "image/jpeg",
    ".m4a": "audio/mp4",
    ".mov": "video/quicktime",
    ".mp3": "audio/mpeg",
    ".mp4": "video/mp4",
    ".ogg": "audio/ogg",
    ".pdf": "application/pdf",
    ".png": "image/png",
    ".svg": "image/svg+xml",
    ".txt": "text/plain",
    ".wav": "audio/wav",
    ".webm": "video/webm",
    ".webp": "image/webp",
}

_KIND_MIMES = {
    "animation": "video/mp4",
    "gif": "video/mp4",
    "photo": "image/jpeg",
    "video": "video/mp4",
    "video_note": "video/mp4",
    "voice": "audio/ogg",
}

_GENERIC_MIMES = {
    "application/octet-stream",
    "application/binary",
    "binary/octet-stream",
    "",
}

_INLINE_MIME_PREFIXES = ("image/", "video/", "audio/", "text/")
_INLINE_MIMES = {"application/pdf", "application/json"}
_ACTIVE_MIMES = {
    "text/html",
    "application/xhtml+xml",
    "image/svg+xml",
    "text/xml",
    "application/xml",
}

_PREVIEW_AGENTS = (
    "telegrambot",
    "twitterbot",
    "facebookexternalhit",
    "whatsapp",
    "discordbot",
    "slackbot",
    "skypeuripreview",
    "vkshare",
    "redditbot",
    "linkedinbot",
    "embedly",
    "quora link preview",
    "googlebot",
    "bingbot",
    "applebot",
    "yandexbot",
    "duckduckbot",
    "petalbot",
    "ahrefsbot",
    "semrushbot",
    "headlesschrome",
    "python-requests",
)

_TRANSLIT = {
    "а": "a",
    "б": "b",
    "в": "v",
    "г": "g",
    "д": "d",
    "е": "e",
    "ё": "e",
    "ж": "zh",
    "з": "z",
    "и": "i",
    "й": "y",
    "к": "k",
    "л": "l",
    "м": "m",
    "н": "n",
    "о": "o",
    "п": "p",
    "р": "r",
    "с": "s",
    "т": "t",
    "у": "u",
    "ф": "f",
    "х": "h",
    "ц": "c",
    "ч": "ch",
    "ш": "sh",
    "щ": "sch",
    "ъ": "",
    "ы": "y",
    "ь": "",
    "э": "e",
    "ю": "yu",
    "я": "ya",
}

_URL_UNSAFE = re.compile(r"[^A-Za-z0-9._-]+")
_URL_REPEATS = re.compile(r"_{2,}")

_UNSAFE_CHARS = re.compile(r'[\\/:*?"<>|\x00-\x1f]+')
_SPACES = re.compile(r"\s+")
_NAME_LIMIT = 120


def extension_for(kind: str | None, mime: str | None) -> str:
    if mime and mime in _MIME_EXTENSIONS:
        return _MIME_EXTENSIONS[mime]
    if kind and kind in _KIND_EXTENSIONS:
        return _KIND_EXTENSIONS[kind]
    if mime and "/" in mime:
        tail = mime.rsplit("/", 1)[1].split(";")[0].strip()
        if tail.isalnum():
            return f".{tail}"
    return ".bin"


def sanitize_name(name: str) -> str:
    cleaned = _SPACES.sub(" ", _UNSAFE_CHARS.sub("_", name)).strip(" .")
    if len(cleaned) > _NAME_LIMIT:
        head, dot, tail = cleaned.rpartition(".")
        cleaned = (
            f"{head[: _NAME_LIMIT - len(tail) - 1]}{dot}{tail}"
            if dot and len(tail) < 12  # noqa: PLR2004
            else cleaned[:_NAME_LIMIT]
        )
    return cleaned or "file"


def url_slug(file_name: str) -> str:
    lowered = "".join(
        _TRANSLIT.get(ch, _TRANSLIT.get(ch.lower(), ch)) for ch in file_name
    )
    slug = _URL_REPEATS.sub("_", _URL_UNSAFE.sub("_", lowered)).strip("_.")
    if not slug:
        return "file"
    return slug[:_NAME_LIMIT]


def media_file_name(
    kind: str | None, mime: str | None, message_id: int, original: str | None = None
) -> str:
    if original:
        name = sanitize_name(original)
        return name if "." in name else f"{name}{extension_for(kind, mime)}"
    return f"{kind or 'media'}_{message_id}{extension_for(kind, mime)}"


def story_file_name(peer_id: int, story_id: int, media_kind: str | None) -> str:
    return f"story_{peer_id}_{story_id}{extension_for(media_kind or 'photo', None)}"


def resolve_mime(
    kind: str | None, mime: str | None, file_name: str | None = None
) -> str:
    if mime and mime.lower() not in _GENERIC_MIMES:
        return mime
    if file_name and "." in file_name:
        by_extension = _EXTENSION_MIMES.get(f".{file_name.rsplit('.', 1)[1].lower()}")
        if by_extension:
            return by_extension
    return _KIND_MIMES.get(kind or "", mime or "application/octet-stream")


def is_inline_mime(mime: str | None) -> bool:
    if not mime:
        return False
    return mime in _INLINE_MIMES or mime.startswith(_INLINE_MIME_PREFIXES)


def is_active_mime(mime: str | None) -> bool:
    return bool(mime) and mime.split(";", 1)[0].strip().lower() in _ACTIVE_MIMES


def content_disposition(file_name: str, *, attachment: bool) -> str:
    kind = "attachment" if attachment else "inline"
    ascii_name = file_name.encode("ascii", "replace").decode("ascii").replace('"', "_")
    return f"{kind}; filename=\"{ascii_name}\"; filename*=UTF-8''{quote(file_name)}"


def is_preview_agent(user_agent: str | None) -> bool:
    if not user_agent:
        return False
    lowered = user_agent.lower()
    return any(marker in lowered for marker in _PREVIEW_AGENTS)


def counts_as_download(
    method: str, user_agent: str | None, range_header: str | None
) -> bool:
    if method.upper() != "GET" or is_preview_agent(user_agent):
        return False
    if range_header is None:
        return True
    return range_header.replace(" ", "").startswith("bytes=0-")


def expiry_from_seconds(seconds: int | None) -> datetime | None:
    if seconds is None or seconds <= 0:
        return None
    return datetime.now(UTC) + timedelta(seconds=seconds)