beavergram/backend/src/utils/read/shares.py

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]