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)