feat(userbot,api,frontend): file links, media downloads, story hold-pause, drop scheduled dupes

This commit is contained in:
hh committed 2026-08-13 03:17:36 +02:00
1 parent 683b9a31a3
commit 3e08698b62
38 files changed
+2889 -134

No files matched your search

+50
View File
@@ -493,3 +493,53 @@ class Dialog(SQLModel, table=True):
onupdate=func.now(),
)
)
class FileShare(SQLModel, table=True):
__tablename__ = "file_shares"
id: int | None = Field(default=None, primary_key=True)
account_id: int
token: str = Field(unique=True)
kind: str
storage_key: str
file_name: str
mime: str | None = None
file_size: int | None = Field(default=None, sa_column=Column(BigInteger))
title: str | None = None
chat_id: int | None = Field(default=None, sa_column=Column(BigInteger))
message_id: int | None = Field(default=None, sa_column=Column(BigInteger))
peer_id: int | None = Field(default=None, sa_column=Column(BigInteger))
story_id: int | None = Field(default=None, sa_column=Column(BigInteger))
expires_at: datetime | None = Field(
default=None, sa_column=Column(DateTime(timezone=True))
)
max_downloads: int | None = None
download_count: int = 0
last_download_at: datetime | None = Field(
default=None, sa_column=Column(DateTime(timezone=True))
)
revoked_at: datetime | None = Field(
default=None, sa_column=Column(DateTime(timezone=True))
)
created_at: datetime = Field(
sa_column=Column(
DateTime(timezone=True), nullable=False, server_default=func.now()
)
)
class FileShareHit(SQLModel, table=True):
__tablename__ = "file_share_hits"
id: int | None = Field(default=None, sa_column=Column(BigInteger, primary_key=True))
share_id: int = Field(foreign_key="file_shares.id")
ts: datetime = Field(
sa_column=Column(
DateTime(timezone=True), nullable=False, server_default=func.now()
)
)
method: str
ip: str | None = None
user_agent: str | None = None
counted: bool = False
+134
View File
@@ -0,0 +1,134 @@
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",
}
_INLINE_MIME_PREFIXES = ("image/", "video/", "audio/", "text/")
_INLINE_MIMES = {"application/pdf", "application/json"}
_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",
)
_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 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 is_inline_mime(mime: str | None) -> bool:
if not mime:
return False
return mime in _INLINE_MIMES or mime.startswith(_INLINE_MIME_PREFIXES)
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)
+34 -8
View File
@@ -1,18 +1,41 @@
import asyncpg
from utils.files import media_file_name
from utils.read.message_view import load_raw
from utils.read.models import MediaVersionView, MediaView
_MEDIA_COLS = (
"id, account_id, chat_id, message_id, kind, storage_key, file_size, "
"mime, ttl_seconds, downloaded, extracted_text, created_at"
MEDIA_COLS = (
"m.id, m.account_id, m.chat_id, m.message_id, m.kind, m.storage_key, "
"m.file_size, m.mime, m.ttl_seconds, m.downloaded, m.extracted_text, "
"m.created_at, src.original_name"
)
ORIGINAL_NAME_JOIN = """
LEFT JOIN LATERAL (
SELECT msg.raw->m.kind->>'file_name' AS original_name
FROM messages msg
WHERE msg.account_id = m.account_id AND msg.chat_id = m.chat_id
AND msg.message_id = m.message_id
ORDER BY msg.date DESC LIMIT 1
) src ON true
"""
_VERSION_COLS = "id, kind, storage_key, file_size, mime, observed_at"
_WEB_PAGE_MEDIA_KINDS = ("photo", "video", "animation", "document", "audio")
def media_view(row: asyncpg.Record) -> MediaView:
fields = dict(row)
original = fields.pop("original_name", None)
return MediaView(
**fields,
file_name=media_file_name(
fields["kind"], fields["mime"], fields["message_id"], original
),
)
async def _web_page_media_stub(
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> MediaView | None:
@@ -47,29 +70,32 @@ async def _web_page_media_stub(
downloaded=False,
extracted_text=None,
created_at=row["date"],
file_name=media_file_name(
kind, obj.get("mime_type"), message_id, obj.get("file_name")
),
)
async def get_media(pool: asyncpg.Pool, media_id: int) -> MediaView | None:
row = await pool.fetchrow(
f"SELECT {_MEDIA_COLS} FROM media WHERE id = $1", # noqa: S608
f"SELECT {MEDIA_COLS} FROM media m {ORIGINAL_NAME_JOIN} WHERE m.id = $1", # noqa: S608
media_id,
)
return MediaView(**dict(row)) if row else None
return media_view(row) if row else None
async def get_message_media(
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> MediaView | None:
row = await pool.fetchrow(
f"SELECT {_MEDIA_COLS} FROM media " # noqa: S608
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
f"SELECT {MEDIA_COLS} FROM media m {ORIGINAL_NAME_JOIN} " # noqa: S608
"WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3",
account_id,
chat_id,
message_id,
)
if row is not None:
return MediaView(**dict(row))
return media_view(row)
return await _web_page_media_stub(pool, account_id, chat_id, message_id)
+4 -1
View File
@@ -5,6 +5,7 @@ from typing import Any
import asyncpg
from pydantic import ValidationError
from utils.files import media_file_name
from utils.read.models import (
ContactView,
EntityView,
@@ -344,6 +345,7 @@ def media_ref_from(
obj = obj if isinstance(obj, dict) else {}
width = obj.get("width") or obj.get("length")
height = obj.get("height") or obj.get("length")
mime = (media_row["mime"] if media_row else None) or obj.get("mime_type")
return MediaRef(
message_id=message_id,
id=media_row["id"] if media_row else None,
@@ -352,11 +354,12 @@ def media_ref_from(
width=width,
height=height,
duration=obj.get("duration"),
mime=(media_row["mime"] if media_row else None) or obj.get("mime_type"),
mime=mime,
file_size=(media_row["file_size"] if media_row else None)
or obj.get("file_size"),
ttl_seconds=media_row["ttl_seconds"] if media_row else None,
extracted_text=media_row["extracted_text"] if media_row else None,
file_name=media_file_name(kind, mime, message_id, obj.get("file_name")),
)
+31
View File
@@ -92,6 +92,7 @@ class MediaRef(BaseModel):
file_size: int | None = None
ttl_seconds: int | None = None
extracted_text: str | None = None
file_name: str | None = None
class ReactionCount(BaseModel):
@@ -222,6 +223,7 @@ class MediaView(BaseModel):
downloaded: bool
extracted_text: str | None
created_at: datetime
file_name: str | None = None
class MediaVersionView(BaseModel):
@@ -391,3 +393,32 @@ class AlertView(BaseModel):
payload: dict[str, Any]
seen: bool
created_at: datetime
class FileShareView(BaseModel):
id: int
account_id: int
token: str
kind: str
file_name: str
mime: str | None
file_size: int | None
title: str | None
chat_id: int | None
message_id: int | None
peer_id: int | None
story_id: int | None
expires_at: datetime | None
max_downloads: int | None
download_count: int
last_download_at: datetime | None
revoked_at: datetime | None
created_at: datetime
class FileShareHitView(BaseModel):
ts: datetime
method: str
ip: str | None
user_agent: str | None
counted: bool
+5 -9
View File
@@ -3,28 +3,24 @@ from datetime import datetime, timedelta
import asyncpg
from utils.read.accounts import self_user_id
from utils.read.media import MEDIA_COLS, ORIGINAL_NAME_JOIN, media_view
from utils.read.models import ChatLinkView, DayCount, MediaView, MessageAt, Page
_MEDIA_COLS = (
"id, account_id, chat_id, message_id, kind, storage_key, file_size, "
"mime, ttl_seconds, downloaded, extracted_text, created_at"
)
async def chat_media(
pool: asyncpg.Pool, account_id: int, chat_id: int, kinds: list[str], page: Page
) -> list[MediaView]:
rows = await pool.fetch(
f"SELECT {_MEDIA_COLS} FROM media " # noqa: S608
"WHERE account_id = $1 AND chat_id = $2 AND kind = ANY($3) "
"ORDER BY message_id DESC LIMIT $4 OFFSET $5",
f"SELECT {MEDIA_COLS} FROM media m {ORIGINAL_NAME_JOIN} " # noqa: S608
"WHERE m.account_id = $1 AND m.chat_id = $2 AND m.kind = ANY($3) "
"ORDER BY m.message_id DESC LIMIT $4 OFFSET $5",
account_id,
chat_id,
kinds,
page.capped_limit,
page.offset,
)
return [MediaView(**dict(row)) for row in rows]
return [media_view(row) for row in rows]
async def chat_links(
+204
View File
@@ -0,0 +1,204 @@
import secrets
from datetime import datetime
import asyncpg
from utils.read.models import FileShareHitView, FileShareView, Page
_COLS = (
"id, account_id, token, kind, file_name, mime, file_size, title, "
"chat_id, message_id, peer_id, story_id, expires_at, max_downloads, "
"download_count, last_download_at, revoked_at, created_at"
)
_TOKEN_BYTES = 12
_SERVE = """
SELECT id, storage_key, file_name, mime, expires_at, max_downloads,
download_count, revoked_at
FROM file_shares WHERE token = $1
"""
_INSERT = (
"INSERT INTO file_shares " # noqa: S608
"(account_id, token, kind, storage_key, file_name, mime, file_size, title, "
"chat_id, message_id, peer_id, story_id, expires_at, max_downloads) "
"VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) "
f"RETURNING {_COLS}"
)
_ACTIVE = "revoked_at IS NULL AND (expires_at IS NULL OR expires_at > now())"
_LOOKUP = (
f"SELECT {_COLS} FROM file_shares " # noqa: S608
f"WHERE account_id = $1 AND storage_key = $2 AND {_ACTIVE} "
"AND (max_downloads IS NULL OR download_count < max_downloads) "
"ORDER BY created_at DESC LIMIT 1"
)
_CONSUME = """
UPDATE file_shares
SET download_count = download_count + 1, last_download_at = now()
WHERE id = $1 AND (max_downloads IS NULL OR download_count < max_downloads)
RETURNING download_count
"""
def new_token() -> str:
return secrets.token_urlsafe(_TOKEN_BYTES)
def _view(row: asyncpg.Record) -> FileShareView:
return FileShareView(**dict(row))
async def list_shares(
pool: asyncpg.Pool, account_id: int, page: Page, *, active_only: bool = False
) -> list[FileShareView]:
where = "account_id = $1"
if active_only:
where += f" AND {_ACTIVE}"
rows = await pool.fetch(
f"SELECT {_COLS} FROM file_shares WHERE {where} " # noqa: S608
"ORDER BY created_at DESC LIMIT $2 OFFSET $3",
account_id,
page.capped_limit,
page.offset,
)
return [_view(row) for row in rows]
async def get_share(pool: asyncpg.Pool, share_id: int) -> FileShareView | None:
row = await pool.fetchrow(
f"SELECT {_COLS} FROM file_shares WHERE id = $1", # noqa: S608
share_id,
)
return _view(row) if row else None
async def find_active_share(
pool: asyncpg.Pool, account_id: int, storage_key: str
) -> FileShareView | None:
row = await pool.fetchrow(_LOOKUP, account_id, storage_key)
return _view(row) if row else None
async def create_share( # noqa: PLR0913
pool: asyncpg.Pool,
account_id: int,
kind: str,
storage_key: str,
file_name: str,
*,
mime: str | None = None,
file_size: int | None = None,
title: str | None = None,
chat_id: int | None = None,
message_id: int | None = None,
peer_id: int | None = None,
story_id: int | None = None,
expires_at: datetime | None = None,
max_downloads: int | None = None,
) -> FileShareView:
row = await pool.fetchrow(
_INSERT,
account_id,
new_token(),
kind,
storage_key,
file_name,
mime,
file_size,
title,
chat_id,
message_id,
peer_id,
story_id,
expires_at,
max_downloads,
)
return _view(row)
async def update_share(
pool: asyncpg.Pool,
share_id: int,
*,
expires_at: datetime | None,
max_downloads: int | None,
) -> FileShareView | None:
row = await pool.fetchrow(
"UPDATE file_shares SET expires_at = $2, max_downloads = $3 " # noqa: S608
f"WHERE id = $1 RETURNING {_COLS}",
share_id,
expires_at,
max_downloads,
)
return _view(row) if row else None
async def revoke_share(pool: asyncpg.Pool, share_id: int) -> FileShareView | None:
row = await pool.fetchrow(
"UPDATE file_shares SET revoked_at = now() " # noqa: S608
f"WHERE id = $1 AND revoked_at IS NULL RETURNING {_COLS}",
share_id,
)
if row is not None:
return _view(row)
return await get_share(pool, share_id)
async def delete_share(pool: asyncpg.Pool, share_id: int) -> bool:
result = await pool.execute("DELETE FROM file_shares WHERE id = $1", share_id)
return result.endswith("1")
async def rotate_token(pool: asyncpg.Pool, share_id: int) -> FileShareView | None:
row = await pool.fetchrow(
"UPDATE file_shares SET token = $2, revoked_at = NULL " # noqa: S608
f"WHERE id = $1 RETURNING {_COLS}",
share_id,
new_token(),
)
return _view(row) if row else None
async def share_by_token(pool: asyncpg.Pool, token: str) -> asyncpg.Record | None:
return await pool.fetchrow(_SERVE, token)
async def consume_download(pool: asyncpg.Pool, share_id: int) -> bool:
return await pool.fetchval(_CONSUME, share_id) is not None
async def record_hit( # noqa: PLR0913
pool: asyncpg.Pool,
share_id: int,
method: str,
ip: str | None,
user_agent: str | None,
*,
counted: bool,
) -> None:
await pool.execute(
"INSERT INTO file_share_hits (share_id, method, ip, user_agent, counted) "
"VALUES ($1, $2, $3, $4, $5)",
share_id,
method,
ip,
user_agent,
counted,
)
async def list_hits(
pool: asyncpg.Pool, share_id: int, page: Page
) -> list[FileShareHitView]:
rows = await pool.fetch(
"SELECT ts, method, ip, user_agent, counted FROM file_share_hits "
"WHERE share_id = $1 ORDER BY ts DESC LIMIT $2 OFFSET $3",
share_id,
page.capped_limit,
page.offset,
)
return [FileShareHitView(**dict(row)) for row in rows]