feat(userbot,api,frontend): file links, media downloads, story hold-pause, drop scheduled dupes
This commit is contained in:
@@ -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]
|
||||
Reference in New Issue
Block a user