import secrets from datetime import datetime import asyncpg from utils.files import url_slug 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_ALPHABET = "abcdefghijkmnpqrstuvwxyz23456789" # noqa: S105 _TOKEN_LENGTH = 10 _SERVE = """ SELECT id, kind, 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 "".join(secrets.choice(_TOKEN_ALPHABET) for _ in range(_TOKEN_LENGTH)) def _view(row: asyncpg.Record) -> FileShareView: fields = dict(row) return FileShareView(**fields, url_name=url_slug(fields["file_name"])) 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]