Compare commits

...
Author SHA1 Message Date
renovate Bot ec5979437d chore(frontend): update template to v0.8.0 2026-10-20 05:00:56 +00:00
renovate Bot 59f171a62d chore(backend): update template to v0.7.1 2026-10-20 05:00:50 +00:00
renovate Bot e3559da2a7 chore(deps): update code dependencies (dev) (#21)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-18 04:02:20 +00:00
hh dcd6786ca5 fix(backend): avatar job ends without a crash when telegram does not know the owner 2026-10-14 05:25:48 +02:00
hh 32d58193f3 fix(backend): raw update reaches every handler that declares it, contact list follows user updates 2026-10-14 05:25:45 +02:00
hh 59d0e2c8ad fix(backend): shared contact card keeps its phone number in the stored message 2026-10-14 05:25:42 +02:00
hh 850d39ad6f fix(frontend): cleanup block shows its failure in the error color 2026-10-14 05:17:58 +02:00
hh b8a38092eb fix(frontend): emptied chat keeps its storage head, shared-only chat explains why nothing can be bulk moved, back from a chat returns to the same place in the list, cleanup wording agrees in number 2026-10-14 05:15:27 +02:00
hh 2836dcb3a9 fix(backend): stale avatar is refetched again under kurigram 2.2.26, profile photos come as chat photos 2026-10-14 05:09:38 +02:00
hh 01cf01b629 fix(backend): rich messages of the new telegram layer read as plain text instead of no text 2026-10-14 05:06:35 +02:00
hh 7f124059be feat(frontend): storage shows the heaviest chats, a chat opens its own files with selective cleanup and a whole-chat move to trash 2026-10-14 05:02:06 +02:00
hh 5eacabd101 feat(backend): storage lists the heaviest chats, summary and file list take a chat filter 2026-10-14 05:02:03 +02:00
hh 0362777a51 fix(frontend): opened policy rule stays in view while its card grows, chevrons hold one column, switch and field sizes match the rest 2026-10-14 04:31:18 +02:00
hh dfc726976b fix(frontend): policy card keeps its wells in the light theme, folder rules show their content, opened rule scrolls into view, rule removal asks first 2026-10-14 04:27:59 +02:00
hh cb01d7dff8 feat(frontend): capture policy shows what each rule changes, rule card opens from folder and chat menus and can trash downloaded files, trash is paged 2026-10-14 04:20:38 +02:00
hh f646f29ff9 feat(backend): policy rules show what they inherit, files of a chat or folder go to trash in one batch, trash list is paged 2026-10-14 04:13:13 +02:00
hh 945adcf879 fix(backend): userbot follows kurigram 2.2.26, user and chat parsing is awaited, missing chat is handled 2026-10-14 03:55:33 +02:00
renovate Bot d213dfdf64 chore(deps): update dependency uv_build to >=0.13.0,<0.14.0 (dev) (#19)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-14 03:54:09 +02:00
renovate Bot 566ee09431 chore(deps): update code dependencies (dev) (#17)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-14 03:54:09 +02:00
hh 09c5f78061 fix(frontend,backend): all stories panel groups by peer again, missing trashed file answers 404, access log hides the token 2026-10-14 03:53:56 +02:00
hh 897421e7c4 fix(frontend): storage failures use the error colour next to the warning, shared mark on its own line in the list, trash key is visible 2026-10-14 03:28:29 +02:00
hh 3223f62dea fix(frontend): selected storage chip keeps its fill in the light theme, tile label uses the medium weight 2026-10-14 03:24:31 +02:00
hh c9c088ed73 fix(frontend): shared mark on its own line in trash rows, undo ignores the trailing click, trash segment has a key 2026-10-14 03:23:22 +02:00
hh 580927da0a fix(frontend,backend): storage categories keep one hue in both themes, trash counted in the total, undo keeps the list, larger touch targets 2026-10-14 03:18:59 +02:00
hh 06ae3e928f refactor(frontend): storage panel answers in place instead of toasts, empty groups hidden, weight bar animates by transform 2026-10-14 03:14:51 +02:00
hh 55634984fd feat(frontend,backend): storage panel with weight rows and trash, heavy media streams by url in the viewer 2026-10-14 03:09:37 +02:00
hh f2c245707a feat(backend): storage trash with seven day retention, trashed files are not fetched again automatically 2026-10-14 02:52:21 +02:00
renovate Bot c20e132872 chore(deps): pin timescale/timescaledb docker tag to b346edc (main) (#16)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-10 23:22:30 +00:00
hh 4d16105e92 feat(frontend)!: sveltekit 3 on the v0.7.0 template, typescript 6, adapter-auto dropped; backend base moves to trixie 2026-10-10 23:18:57 +00:00
hh b1ffebe858 fix(infra): timescaledb image 2.27.1 to 2.30.2, the high advisory fixed in 2.29.1 was still open 2026-10-10 20:50:06 +00:00
hh 2512a7988d chore(frontend): biome 2.5.15 stops flagging regexp exec results as always truthy, svelte files reformatted by it 2026-10-10 21:05:02 +02:00
hh 56146e4e8d chore(backend,infra): resolve template update conflicts, import check knows the watcher and the api login 2026-10-10 20:59:12 +02:00
hh 5fb8f8a534 style(backend): one-line import, ruff 0.17 keeps the pragma out of the line width 2026-10-10 20:59:12 +02:00
renovate Bot 98ab8a579b chore(infra): update template to v0.6.0, conflicts left 2026-10-10 20:59:12 +02:00
renovate Bot a5d10146fd chore(frontend): update template to v0.6.0 2026-10-10 20:59:12 +02:00
renovate Bot 5f1f18b343 chore(caddy): update template to v0.1.4 2026-10-10 20:59:12 +02:00
renovate Bot 8eb31ed4d3 chore(backend): update template to v0.6.0, conflicts left 2026-10-10 20:59:12 +02:00
renovate Bot e67dbdb325 chore(deps): update dependency ultracite to v7.12.4 (dev) (#13)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-10 20:59:12 +02:00
hh 869f8291fb fix(api): access log keeps the mcp path but not its query, the token rode in it 2026-10-09 02:44:01 +00:00
hh 283eab1eb0 feat(mcp): download any message file via share link 2026-10-08 21:22:59 +02:00
hh 7f4d500302 chore(infra): update template to v0.5.0 2026-10-08 21:22:59 +02:00
hh cf48d2db3e chore(frontend): update template to v0.4.2 2026-10-08 21:22:59 +02:00
hh 114058e6ae chore(backend): update template to v0.4.0 2026-10-08 21:22:59 +02:00
renovate Bot eb8187ff6e chore(deps): update dependency ultracite to v7.12.2 (dev) (#12)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-03 06:02:55 +02:00
renovate Bot 53dc805fc3 chore(deps): update dependency ultracite to v7.12.1 (dev) (#11)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-10-01 06:03:12 +02:00
renovate Bot ed8489616e chore(deps): update dependency @biomejs/biome to v2.5.14 (dev) (#10)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-09-20 06:03:12 +02:00
hh b60815d042 fix(watcher): pick an unfinished backlog back up without restarting the bot 2026-09-19 16:19:46 +00:00
hh 468f72ed11 fix(watcher): resume the bot backlog from a saved cursor instead of restarting it 2026-09-19 16:18:40 +00:00
hh 59d3f214e8 chore(renovate): hold typescript on 5.x for svelte-check 2026-09-19 00:33:22 +02:00
hh 34dcbaac2b build(frontend): pin typescript to 5.x and decode tgs stickers with pako toText 2026-09-19 00:21:15 +02:00
hh c8e5cff871 style(frontend): apply ultracite 7.12 formatting and lint rules 2026-09-19 00:21:07 +02:00
hh fc4758aa0b fix(watcher): getDifference needs a date, and only a completed drain counts 2026-09-18 18:42:38 +02:00
hh b0ed041c2e feat(watcher,api,frontend): watched bots become accounts and drain the update backlog 2026-09-18 18:41:20 +02:00
hh 4bd031f21e feat(watcher,api,frontend): archive who writes to watched bots over mtproto 2026-09-18 15:05:01 +02:00
renovate Bot 089a31db8d chore(deps): update dependency ultracite to v7.12.0 (dev) (#9)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-09-18 06:02:46 +02:00
renovate Bot 81ebbdc93e chore(deps): update code dependencies (dev) (#8)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-09-15 23:56:37 +02:00
hh a25f79543a fix(frontend): regenerate bun.lock after dependency bumps 2026-09-15 10:52:23 +00:00
renovate Bot c61ecfe8f2 chore(deps): pin ghcr.io/astral-sh/uv docker tag to 531f855 (#1)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-09-14 06:03:16 +02:00
hh 3d0b15bd2c Merge pull request 'chore(deps): update dependency typescript to v7' (#6) from renovate/typescript-7.x into main
Reviewed-on: #6
2026-09-10 20:00:14 +02:00
hh 0e54786672 Merge pull request 'fix(deps): update dependency pako to v3' (#7) from renovate/pako-3.x into main
Reviewed-on: #7
2026-09-10 20:00:06 +02:00
renovate Bot 5778611a38 fix(deps): update dependency pako to v3
renovate/stability-days Updates have met minimum release age requirement
renovate/artifacts Artifact file update failure
2026-09-10 17:58:54 +00:00
renovate Bot b138533d88 chore(deps): update dependency typescript to v7
renovate/stability-days Updates have met minimum release age requirement
renovate/artifacts Artifact file update failure
2026-09-10 17:58:51 +00:00
hh 0b6cbac945 chore(infra): update template to v0.4.0 2026-09-10 19:07:01 +02:00
hh ff80141785 refactor(renovate): use shared preset instead of a copy 2026-09-10 17:58:36 +02:00
renovate Bot d206f7e5ca chore(deps): pin oven/bun docker tag to d888c0a (#2)
Co-authored-by: renovate <renovate@hhhhh.dev>
2026-09-10 17:44:26 +02:00
hh 4ebe6de60b chore(renovate): add Renovate config 2026-09-10 15:03:44 +00:00
hh 02fb567c1f chore(backend): update template to v0.3.3 2026-09-09 01:17:16 +02:00
hh ccbc167005 chore(*): adopt project templates 2026-09-09 01:17:12 +02:00
hh 5ec0212af5 fix(userbot): mark deletions per row to stay under timescale decompression limit 2026-09-06 21:30:51 +02:00
hh ab59ef4525 feat(api,mcp): explain account_id mixed up with a telegram user id 2026-09-06 21:30:51 +02:00
hh 676680319b fix(userbot): encode absent pyrogram raw payloads as empty json 2026-09-06 21:26:13 +02:00
hh 6a6e63ff75 docs(readme): what it keeps, the parts, the mcp tools, how to run 2026-09-05 04:56:56 +02:00
hh 1fa4f18a9b perf(*): bound list queries, batch fetches, harden media downloads 2026-09-01 21:38:22 +02:00
hh 9c12628304 fix(userbot): drop stale media sessions and survive download timeouts 2026-09-01 21:05:48 +02:00
hh cf277c56cf fix(userbot): stop gracefully on sigterm and flush session peers periodically 2026-08-30 01:56:10 +02:00
hh 8f7d47476c fix(userbot): refetch avatar by unique_id when stored file_id is stale 2026-08-30 01:50:56 +02:00
hh c832513e5c fix(api,frontend): isolate active share files and allowlist link schemes 2026-08-29 12:52:54 +02:00
hh 0399145791 style(frontend): soften the document glyph nudge to a hairline 2026-08-13 04:13:03 +02:00
hh 0b8e90159f style(frontend): optically nudge the document glyph in file chips 2026-08-13 04:07:34 +02:00
hh 995322b17a fix(frontend): readable links and file chips on own bubbles 2026-08-13 04:06:44 +02:00
hh 6ae28cb321 fix(frontend): render document albums as file rows instead of broken images 2026-08-13 03:59:01 +02:00
hh e833de245a fix(api,frontend): render files by name without fetching, previewable share links 2026-08-13 03:49:29 +02:00
hh ef739838a5 fix(migrations): lift timescale decompression limit when dropping scheduled dupes 2026-08-13 03:30:07 +02:00
hh 3e08698b62 feat(userbot,api,frontend): file links, media downloads, story hold-pause, drop scheduled dupes 2026-08-13 03:17:36 +02:00
hh 683b9a31a3 fix(frontend): restore outside-click dismiss, collapse story groups by default 2026-08-06 15:28:10 +02:00
hh 004fd56f2c feat(frontend): open sidebar media on click, unbreak long-press menus 2026-08-06 15:17:25 +02:00
hh 12cc7d57e3 feat(userbot,api,frontend): backfill stories and unbreak all-stories list 2026-08-06 14:08:10 +02:00
hh 1e143c7573 fix(frontend): unbreak round video ring and playback in firefox 2026-08-06 13:52:27 +02:00
hh ee63f8b783 feat(api,userbot,frontend): fetch only new messages in backfill by default 2026-08-06 12:30:37 +02:00
hh 6b6edc9a0d fix(frontend): paint video poster frames, clip round videos to circle 2026-08-06 12:29:37 +02:00
hh 525ce024bc perf(api,frontend): index hot queries, paginate chats, fix realtime 2026-08-06 01:52:41 +02:00
hh 9d767d2531 chore(frontend): sync bun lockfile 2026-08-06 01:07:52 +02:00
hh 18220407af fix(make): enable db and migrate profiles in migrate target 2026-08-06 01:07:51 +02:00
293 changed files with 14602 additions and 2662 deletions

No files matched your search

+1 -1
View File
@@ -4,7 +4,7 @@
{ {
"hooks": [ "hooks": [
{ {
"command": "bun fix --skip=correctness/noUnusedImports", "command": "cd frontend && bun run fix --skip=correctness/noUnusedImports",
"type": "command" "type": "command"
} }
], ],
+17
View File
@@ -0,0 +1,17 @@
_commit: v0.7.1
_src_path: https://git.kotikot.com/templates/backend-python.git
author_email: h@kotikot.com
author_name: h
backend_services:
- api
- kurigram_userbot
database: postgres
dynamic_config: false
install_deps: false
layout: part
project_description: A read-only Telegram client that keeps everything
project_name: beavergram
project_slug: beavergram
serve_spa: true
use_llm: false
use_redis: false
+5
View File
@@ -0,0 +1,5 @@
_commit: v0.1.4
_src_path: https://git.kotikot.com/templates/caddy.git
project_slug: beavergram
routes:
- api
+9
View File
@@ -0,0 +1,9 @@
_commit: v0.8.0
_src_path: https://git.kotikot.com/templates/frontend-svelte.git
install_deps: false
layout: part
linter: biome
project_name: beavergram
project_slug: beavergram
serve: api
ui_library: none
+12
View File
@@ -0,0 +1,12 @@
_commit: v0.6.0
_src_path: https://git.kotikot.com/templates/infra.git
init_git: false
layout: parts
mcp_servers:
- context7
- svelte
parts:
- backend
- frontend
project_description: A read-only Telegram client that keeps everything
project_name: beavergram
+2 -16
View File
@@ -1,16 +1,2 @@
COMPOSE_PROFILES=db,userbot,api COMPOSE_FILE=docker-compose.yml:backend/docker-compose.override.yml:frontend/docker-compose.override.yml
RUN_ENVIRONMENT=prod COMPOSE_PROFILES=external,services
# Leave empty for *.localhost.
FRONTEND_DEV_HOST=
DB__HOST=postgres
DB__PORT=5432
DB__USER=beavergram
DB__PASSWORD=beavergram
DB__DB_NAME=beavergram
DB__MIN_POOL_SIZE=5
DB__MAX_POOL_SIZE=20
DB__SCRIPTS_CONNECTION_URL=postgresql://beavergram:beavergram@localhost:5433/beavergram
STORAGE__ROOT=./storage
+8 -3
View File
@@ -1,6 +1,11 @@
.idea
docker-compose.override.yml
.env .env
.env.*
!.env.example
docker-compose.override.yml
.DS_Store
.idea
docs docs
storage storage
backend/sessions
# Local playground
/t/
+4 -4
View File
@@ -1,12 +1,12 @@
{ {
"mcpServers": { "mcpServers": {
"svelte": {
"type": "http",
"url": "https://mcp.svelte.dev/mcp"
},
"context7": { "context7": {
"type": "http", "type": "http",
"url": "https://mcp.context7.com/mcp" "url": "https://mcp.context7.com/mcp"
},
"svelte": {
"type": "http",
"url": "https://mcp.svelte.dev/mcp"
} }
} }
} }
+18 -22
View File
@@ -1,31 +1,27 @@
repos: repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.15
hooks:
- id: ruff-check
types_or: [ python, pyi ]
args: [ --fix ]
- id: ruff-format
types_or: [ python, pyi ]
- repo: local - repo: local
hooks: hooks:
- id: ultracite - id: backend-fmt
name: ultracite name: backend fmt
entry: bash -c 'cd frontend && bun x ultracite fix' entry: make -C backend fmt
language: system
files: ^backend/
pass_filenames: false
- id: backend-check
name: backend check
entry: make -C backend check
language: system
files: ^backend/
pass_filenames: false
- id: frontend-fmt
name: frontend fmt
entry: make -C frontend fmt
language: system language: system
files: ^frontend/ files: ^frontend/
pass_filenames: false pass_filenames: false
- id: frontend-check
- id: svelte-check name: frontend check
name: svelte check entry: make -C frontend check
entry: bash -c 'cd frontend && bun check'
language: system language: system
files: ^frontend/ files: ^frontend/
pass_filenames: false pass_filenames: false
- id: ty
name: ty check
entry: uvx ty check
language: python
types_or: [ python, pyi ]
+3 -2
View File
@@ -6,10 +6,11 @@
- **Reusability**: Create reusable components and functions - **Reusability**: Create reusable components and functions
- **Less Code = Less Debt**: Minimize code footprint - **Less Code = Less Debt**: Minimize code footprint
- **NEVER write comments** - code should be self-documenting - **NEVER write comments** - code should be self-documenting
- **Docstrings ONLY if explicitly asked**
## Checking commands ## Checking commands
After writing code, always run: After writing code, always run:
```shell ```shell
bun check # svelte-check make fmt
bun fix # biome linter make check
``` ```
+27 -23
View File
@@ -1,30 +1,34 @@
.PHONY: recreate down restart rebuild deploy migrate session-create frontend PARTS := $(patsubst %/Makefile,%,$(wildcard */Makefile))
.PHONY: env recreate rebuild restart logs down fmt check build
env:
@fresh=; test -f .env || { awk 'FNR == 1 && NR > 1 { print "" } 1' .env.example $(addsuffix /.env.example,$(PARTS)) > .env; fresh=1; }; \
for part in $(PARTS); do \
override=$$part/docker-compose.override.yml; \
test -f $$override && continue; \
cp $$override.example $$override; \
for port in $$(sed -n 's/.*127\.0\.0\.1:\([0-9]*\):.*/\1/p' $$override | sort -u); do \
random=$$((20000 + $$(od -An -N2 -tu2 /dev/urandom) % 40000)); \
sed "s/127\.0\.0\.1:$$port:/127.0.0.1:$$random:/" $$override > $$override.tmp && mv $$override.tmp $$override; \
test -z "$$fresh" || { sed "s|localhost:$$port/|localhost:$$random/|" .env > .env.tmp && mv .env.tmp .env; }; \
done; \
done
recreate: recreate:
docker compose up -d --force-recreate docker compose up -d --force-recreate
rebuild: build
docker compose up -d
restart:
docker compose restart
logs:
docker compose logs -f --tail=100
down: down:
docker compose down docker compose down
restart: fmt check build:
$(MAKE) recreate @for part in $(PARTS); do $(MAKE) -C $$part $@ || exit 1; done
rebuild:
docker compose build
docker compose up -d
frontend:
docker compose run --rm --no-deps frontend-dev sh -c "bun install && bun run build"
deploy:
$(MAKE) frontend
$(MAKE) rebuild
migrate:
docker compose run --rm migrator $(filter-out $@,$(MAKECMDGOALS))
session-create:
cd backend && uv run python scripts/session/create.py
%:
@:
+45 -1
View File
@@ -1,2 +1,46 @@
# beavergram # beavergram
[![AI Slop Inside](https://sladge.net/badge.svg)](https://sladge.net)
A read-only Telegram client that keeps everything. A userbot on your own account logs every incoming message into Postgres, including the ones later deleted or edited, along with media, reactions, presence and stories. A web client reads that archive back, and an MCP server hands it to agents.
Useful on its own: a Telegram where nothing disappears. In the [beaver](https://git.kotikot.com/beaver) setup it is one of the agent's hands - the dispatcher reads chat history by the `chat_id` stored in a person's card.
## What it keeps
- **Messages** with edits as versions and deletions as flags, so a message removed in Telegram is still in the archive with its history.
- **Media** on disk, content-addressed, with the versions of a replaced photo.
- **Voice** transcribed through Telegram's own transcription (Premium), queued when the quota runs out.
- **Presence** as a time series (TimescaleDB), read receipts, stories.
- **Capture policy** from your Telegram folders with per-chat overrides: channels nothing, groups messages only, private chats everything, by default.
- **Annotations and watches**: your notes on chats and rules that fire alerts, both yours to write in the web client.
- **Watched bots**: bot tokens you add in the web client. Each one logs in over MTProto (`auth.importBotAuthorization`, the same door a bot-login client uses), backfills the history of the bot's dialogs once and then keeps the incoming stream: who writes to the bot, what, and when. A separate authorization, so it does not take `getUpdates` away from whoever runs that bot over the Bot API.
Nothing is ever deleted from the archive.
## Parts
- `backend/` - `userbot` (kurigram, the capture loop and jobs), `watcher` (kurigram again, one bot-token session per watched bot) and `api` (FastAPI: the web API, the MCP server, and the built frontend as static files).
- `frontend/` - SvelteKit, `adapter-static`; a chat client over the archive: folders, chat list, thread with media, versions and the deleted messages inline, search, stories, presence graphs.
- `docker-compose.yml` - includes `backend/` (Postgres with TimescaleDB, `userbot`, `api`) and `frontend/` (the dev server); profiles `external` and `services`.
## MCP
Read-only: `list_watched_bots`, `get_bot_stats`, `get_bot_updates`, `search_messages_tool`, `list_chats`, `get_chat_history` (with deleted messages included), `get_deleted_messages`, `get_message_versions`, `get_media`, `get_callbacks`, `presence_history`, `get_peer_history`, `get_stories`, `get_annotations`, `list_watches`, `list_alerts`. The one write is `set_watch`. There are no send, edit or delete tools and there will not be.
## Run
```sh
make env # .env and docker-compose.override.yml from the examples
make -C backend session # log the account in once; the session file stays in backend/sessions
make -C backend migrate # the schema
make recreate # build and start the services in COMPOSE_PROFILES
```
`COMPOSE_PROFILES` in `.env` selects what runs on the machine, `docker-compose.override.yml` in each part holds the published ports and the `caddy` network aliases. The web client and the MCP endpoint are served by `api`; the bearer for both is `AUTH__TOKEN`.
## Develop
```sh
make fmt
make check
make -C frontend dev
```
+7
View File
@@ -0,0 +1,7 @@
.venv
.env
__pycache__
*.pyc
.ruff_cache
.pytest_cache
sessions
+27 -9
View File
@@ -1,14 +1,32 @@
TG__SESSION_NAME=beavergram RUN_ENVIRONMENT=prod
TG__SESSIONS_DIR=sessions TZ=UTC
API__HOST=0.0.0.0
API__PORT=8080
AUTH__TOKEN=<AUTH__TOKEN>
STORAGE__SHARD_DEPTH=2
LOG__LEVEL=INFO LOG__LEVEL=INFO
LOG__LEVEL_EXTERNAL=WARNING LOG__LEVEL_EXTERNAL=WARNING
LOG__SHOW_TIME=false LOG__SHOW_TIME=false
LOG__SHOW_LOCALS=false
LOG__CONSOLE_WIDTH=150 LOG__CONSOLE_WIDTH=150
DB__HOST=postgres
DB__PORT=5432
DB__USER=beavergram
DB__PASSWORD=beavergram
DB__DB_NAME=beavergram
DB__MIN_POOL_SIZE=5
DB__MAX_POOL_SIZE=20
DB__SCRIPTS_CONNECTION_URL=postgresql://beavergram:beavergram@localhost:5432/beavergram
TG__SESSION_NAME=beavergram
TG__SESSIONS_DIR=sessions
BOTS__SESSIONS_DIR=sessions/bots
BOTS__BACKFILL_CHATS=0
BOTS__BACKFILL_PER_CHAT=0
API__HOST=0.0.0.0
API__PORT=8080
AUTH__TOKEN=
STORAGE__ROOT=./storage
STORAGE__SHARD_DEPTH=2
+8
View File
@@ -0,0 +1,8 @@
## Checking commands
After writing code, always run from `backend/`:
```shell
make fmt
make check
```
Mind startup imports: `make check` runs `scripts/check_imports.py`.
+6 -5
View File
@@ -1,4 +1,6 @@
FROM ghcr.io/astral-sh/uv:python3.13-bookworm-slim FROM ghcr.io/astral-sh/uv:python3.13-trixie-slim
RUN apt-get update && apt-get install -y build-essential && rm -rf /var/lib/apt/lists/*
WORKDIR /app WORKDIR /app
@@ -7,10 +9,6 @@ ENV PATH="/app/.venv/bin:$PATH" \
UV_COMPILE_BYTECODE=1 \ UV_COMPILE_BYTECODE=1 \
UV_LINK_MODE=copy UV_LINK_MODE=copy
RUN apt-get update \
&& apt-get install -y --no-install-recommends build-essential \
&& rm -rf /var/lib/apt/lists/*
COPY pyproject.toml uv.lock ./ COPY pyproject.toml uv.lock ./
RUN uv sync --frozen --no-install-project --no-dev RUN uv sync --frozen --no-install-project --no-dev
@@ -20,4 +18,7 @@ RUN uv sync --frozen --no-dev
COPY alembic.ini ./ COPY alembic.ini ./
COPY migrations ./migrations COPY migrations ./migrations
COPY scripts ./scripts
ENTRYPOINT ["python", "-m"] ENTRYPOINT ["python", "-m"]
CMD ["api"]
+28
View File
@@ -0,0 +1,28 @@
COMPOSE := cd .. && docker compose
.PHONY: fmt check build migrate revision session
fmt:
uv run ruff format
uv run ruff check --fix
check:
uv run ruff format --check
uv run ruff check
uv run ty check
uv run python scripts/check_imports.py
build:
$(COMPOSE) build api
migrate:
$(COMPOSE) run --rm migrator $(or $(filter-out $@,$(MAKECMDGOALS)),upgrade head)
revision:
$(COMPOSE) run --rm migrator revision --autogenerate -m "$(m)"
%:
@:
session:
uv run python scripts/session/create.py
@@ -0,0 +1,17 @@
services:
api:
networks:
default: {}
caddy:
aliases:
- beavergram-api
# ports:
# - "127.0.0.1:8080:${API__PORT}"
postgres:
# ports:
# - "127.0.0.1:5432:${DB__PORT}"
networks:
caddy:
external: true
+133
View File
@@ -0,0 +1,133 @@
services:
api:
build: backend
image: beavergram/backend
profiles: [api, services]
restart: unless-stopped
env_file:
- path: .env
required: false
- path: backend/.env
required: false
environment:
RUN_ENVIRONMENT: prod
TZ: ${TZ:-UTC}
STORAGE__ROOT: /app/storage
volumes:
- ./backend/src:/app/src
- ./backend/scripts:/app/scripts
- ./backend/sessions:/app/sessions
- ./frontend/build:/app/static:ro
- ${STORAGE__ROOT:-./storage}:/app/storage
command: [api]
depends_on:
postgres:
condition: service_healthy
required: false
networks:
default:
database:
userbot:
build: backend
image: beavergram/backend
profiles: [userbot, services]
restart: unless-stopped
env_file:
- path: .env
required: false
- path: backend/.env
required: false
environment:
RUN_ENVIRONMENT: prod
TZ: ${TZ:-UTC}
STORAGE__ROOT: /app/storage
volumes:
- ./backend/src:/app/src
- ./backend/scripts:/app/scripts
- ./backend/sessions:/app/sessions
- ${STORAGE__ROOT:-./storage}:/app/storage
command: [userbot]
depends_on:
postgres:
condition: service_healthy
required: false
networks:
database:
watcher:
build: backend
image: beavergram/backend
profiles: [watcher, services]
restart: unless-stopped
env_file:
- path: .env
required: false
- path: backend/.env
required: false
environment:
RUN_ENVIRONMENT: prod
TZ: ${TZ:-UTC}
volumes:
- ./backend/src:/app/src
- ./backend/sessions:/app/sessions
command: [watcher]
depends_on:
postgres:
condition: service_healthy
required: false
networks:
database:
postgres:
image: timescale/timescaledb:2.30.2-pg17@sha256:b346edcdb51a1fd6020e3965e0bd1c9f3406fa6d5fbce1e28f4852587ef934e2
profiles: [postgres, external]
restart: unless-stopped
environment:
POSTGRES_USER: ${DB__USER:-beavergram}
POSTGRES_PASSWORD: ${DB__PASSWORD:-beavergram}
POSTGRES_DB: ${DB__DB_NAME:-beavergram}
PGPORT: ${DB__PORT:-5432}
volumes:
- pgdata:/var/lib/postgresql/data
- ./backend/migrations/init:/docker-entrypoint-initdb.d:ro
healthcheck:
test: ["CMD-SHELL", "pg_isready -U ${DB__USER:-beavergram}"]
interval: 5s
timeout: 5s
retries: 10
networks:
database:
aliases:
- ${DB__HOST:-postgres}
migrator:
build: backend
image: beavergram/backend
profiles: [migrate]
env_file:
- path: .env
required: false
- path: backend/.env
required: false
environment:
RUN_ENVIRONMENT: prod
TZ: ${TZ:-UTC}
volumes:
- ./backend/src:/app/src
- ./backend/migrations:/app/migrations
- ./backend/alembic.ini:/app/alembic.ini
entrypoint: [alembic]
command: [upgrade, head]
depends_on:
postgres:
condition: service_healthy
required: false
networks:
database:
volumes:
pgdata:
networks:
database:
+10 -1
View File
@@ -20,10 +20,15 @@ def _async_url() -> str:
return env.db.connection_url.replace("postgresql://", "postgresql+asyncpg://", 1) return env.db.connection_url.replace("postgresql://", "postgresql+asyncpg://", 1)
def include_name(name: str | None, type_: str, _parents: object) -> bool:
return not (type_ == "table" and name is not None and name.startswith("taskiq_"))
def run_migrations_offline() -> None: def run_migrations_offline() -> None:
context.configure( context.configure(
url=_async_url(), url=_async_url(),
target_metadata=target_metadata, target_metadata=target_metadata,
include_name=include_name,
literal_binds=True, literal_binds=True,
dialect_opts={"paramstyle": "named"}, dialect_opts={"paramstyle": "named"},
) )
@@ -32,7 +37,11 @@ def run_migrations_offline() -> None:
def _do_run_migrations(connection) -> None: # noqa: ANN001 def _do_run_migrations(connection) -> None: # noqa: ANN001
context.configure(connection=connection, target_metadata=target_metadata) context.configure(
connection=connection,
target_metadata=target_metadata,
include_name=include_name,
)
with context.begin_transaction(): with context.begin_transaction():
context.run_migrations() context.run_migrations()
@@ -0,0 +1,49 @@
"""read path indexes
Revision ID: a6d2f8b41c9e
Revises: f2b8d3c9a51e
Create Date: 2026-09-01 12:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "a6d2f8b41c9e"
down_revision: str | None = "f2b8d3c9a51e"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
op.execute("CREATE INDEX ix_jobs_account_id ON jobs (account_id, id DESC)")
op.execute(
"CREATE INDEX ix_messages_deleted ON messages "
"(account_id, deleted_at DESC) WHERE deleted_at IS NOT NULL"
)
op.execute(
"CREATE INDEX ix_messages_pinned ON messages "
"(account_id, chat_id, date DESC, message_id DESC) "
"WHERE raw->>'service' LIKE '%PINNED_MESSAGE%'"
)
op.execute("CREATE INDEX ix_watches_account ON watches (account_id, id DESC)")
op.execute("CREATE INDEX ix_alerts_account_ts ON alerts (account_id, ts DESC)")
op.execute(
"CREATE INDEX ix_stories_account_date ON stories "
"(account_id, date DESC NULLS LAST)"
)
op.execute(
"CREATE INDEX ix_annotations_account_msg ON annotations "
"(account_id, chat_id, message_id)"
)
def downgrade() -> None:
op.execute("DROP INDEX IF EXISTS ix_annotations_account_msg")
op.execute("DROP INDEX IF EXISTS ix_stories_account_date")
op.execute("DROP INDEX IF EXISTS ix_alerts_account_ts")
op.execute("DROP INDEX IF EXISTS ix_watches_account")
op.execute("DROP INDEX IF EXISTS ix_messages_pinned")
op.execute("DROP INDEX IF EXISTS ix_messages_deleted")
op.execute("DROP INDEX IF EXISTS ix_jobs_account_id")
@@ -0,0 +1,25 @@
"""retry bot backlog
Revision ID: a7f3d2b81c04
Revises: e2c9a41f7b68
Create Date: 2026-09-18 17:30:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "a7f3d2b81c04"
down_revision: str | None = "e2c9a41f7b68"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
"""No drain has succeeded yet, so let every bot try the queue once more."""
op.execute("UPDATE watched_bots SET backfilled_at = NULL")
def downgrade() -> None:
pass
@@ -0,0 +1,34 @@
"""bot backlog cursor
Revision ID: b6d1f4a82e57
Revises: a7f3d2b81c04
Create Date: 2026-09-19 10:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "b6d1f4a82e57"
down_revision: str | None = "a7f3d2b81c04"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
op.execute(
"ALTER TABLE watched_bots "
"ADD COLUMN backfill_pts integer, "
"ADD COLUMN backfill_qts integer, "
"ADD COLUMN backfill_date integer"
)
def downgrade() -> None:
op.execute(
"ALTER TABLE watched_bots "
"DROP COLUMN IF EXISTS backfill_pts, "
"DROP COLUMN IF EXISTS backfill_qts, "
"DROP COLUMN IF EXISTS backfill_date"
)
@@ -0,0 +1,41 @@
"""hot path indexes
Revision ID: c1f6b3d84a92
Revises: b9e4d1a70c26
Create Date: 2026-08-06 12:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "c1f6b3d84a92"
down_revision: str | None = "b9e4d1a70c26"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
op.execute(
"CREATE INDEX ix_messages_chat_date ON messages "
"(account_id, chat_id, date DESC, message_id DESC)"
)
op.execute("CREATE INDEX ix_avatars_owner ON avatars (account_id, owner_id)")
op.execute("CREATE INDEX ix_media_message ON media (account_id, message_id)")
op.execute(
"CREATE INDEX ix_chat_history_chat_ts ON chat_history "
"(account_id, chat_id, ts DESC)"
)
op.execute(
"CREATE INDEX ix_read_receipts_chat ON read_receipts "
"(account_id, chat_id, kind, message_id DESC)"
)
def downgrade() -> None:
op.execute("DROP INDEX IF EXISTS ix_read_receipts_chat")
op.execute("DROP INDEX IF EXISTS ix_chat_history_chat_ts")
op.execute("DROP INDEX IF EXISTS ix_media_message")
op.execute("DROP INDEX IF EXISTS ix_avatars_owner")
op.execute("DROP INDEX IF EXISTS ix_messages_chat_date")
@@ -0,0 +1,85 @@
"""bot watcher
Revision ID: c8f3b6a2d571
Revises: a6d2f8b41c9e
Create Date: 2026-09-18 12:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "c8f3b6a2d571"
down_revision: str | None = "a6d2f8b41c9e"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_BOTS = """
CREATE TABLE watched_bots (
bot_id bigint PRIMARY KEY,
username text,
title text,
enabled boolean NOT NULL DEFAULT true,
status text NOT NULL DEFAULT 'pending',
status_detail text,
backfilled_at timestamptz,
checked_at timestamptz,
added_at timestamptz NOT NULL DEFAULT now(),
updated_at timestamptz NOT NULL DEFAULT now()
)
"""
_UPDATES = """
CREATE TABLE bot_updates (
bot_id bigint NOT NULL,
kind text NOT NULL,
chat_id bigint NOT NULL,
item_id bigint NOT NULL,
ts timestamptz NOT NULL,
source text NOT NULL,
outgoing boolean NOT NULL DEFAULT false,
sender_id bigint,
sender_username text,
sender_name text,
sender_is_bot boolean,
chat_type text,
chat_title text,
text text,
media_kind text,
edited_at timestamptz,
received_at timestamptz NOT NULL DEFAULT now(),
raw jsonb NOT NULL,
PRIMARY KEY (bot_id, kind, chat_id, item_id, ts)
)
"""
def upgrade() -> None:
op.execute(_BOTS)
op.execute(_UPDATES)
op.execute(
"SELECT create_hypertable('bot_updates', by_range('ts', INTERVAL '1 month'))"
)
op.execute(
"ALTER TABLE bot_updates SET ("
"timescaledb.enable_columnstore = true, "
"timescaledb.segmentby = 'bot_id', "
"timescaledb.orderby = 'ts DESC')"
)
op.execute(
"CALL add_columnstore_policy('bot_updates', after => INTERVAL '30 days')"
)
op.execute("CREATE INDEX ix_bot_updates_recent ON bot_updates (bot_id, ts DESC)")
op.execute(
"CREATE INDEX ix_bot_updates_sender ON bot_updates "
"(bot_id, sender_id, ts DESC) WHERE NOT outgoing"
)
op.execute(
"CREATE INDEX ix_bot_updates_chat ON bot_updates (bot_id, chat_id, ts DESC)"
)
def downgrade() -> None:
op.execute("DROP TABLE IF EXISTS bot_updates")
op.execute("DROP TABLE IF EXISTS watched_bots")
@@ -0,0 +1,57 @@
"""media trash
Revision ID: c9e2a7f4b1d8
Revises: b6d1f4a82e57
Create Date: 2026-10-14 03:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "c9e2a7f4b1d8"
down_revision: str | None = "b6d1f4a82e57"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_TRASH = """
CREATE TABLE media_trash (
storage_key text PRIMARY KEY,
account_id integer NOT NULL,
chat_id bigint NOT NULL,
message_id bigint NOT NULL,
kind text NOT NULL,
mime text,
file_name text NOT NULL,
file_size bigint,
trashed_at timestamptz NOT NULL DEFAULT now()
)
"""
def upgrade() -> None:
op.execute(
"ALTER TABLE media ADD COLUMN evicted_at timestamptz, "
"ADD COLUMN evicted_key text"
)
op.execute(
"CREATE INDEX ix_media_evicted_unique ON media (unique_id) "
"WHERE evicted_at IS NOT NULL"
)
op.execute(
"CREATE INDEX ix_media_evicted_key ON media (evicted_key) "
"WHERE evicted_key IS NOT NULL"
)
op.execute(_TRASH)
op.execute(
"CREATE INDEX ix_media_trash_account ON media_trash "
"(account_id, trashed_at DESC)"
)
def downgrade() -> None:
op.execute("DROP TABLE IF EXISTS media_trash")
op.execute("DROP INDEX IF EXISTS ix_media_evicted_key")
op.execute("DROP INDEX IF EXISTS ix_media_evicted_unique")
op.execute("ALTER TABLE media DROP COLUMN evicted_key, DROP COLUMN evicted_at")
@@ -0,0 +1,119 @@
"""chat stats
Revision ID: d4a7e2b91f38
Revises: c1f6b3d84a92
Create Date: 2026-08-06 12:30:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "d4a7e2b91f38"
down_revision: str | None = "c1f6b3d84a92"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_APPLY = """
CREATE FUNCTION chat_stats_apply() RETURNS trigger LANGUAGE plpgsql AS $$
BEGIN
IF TG_OP = 'INSERT' THEN
INSERT INTO chat_stats AS cs (account_id, chat_id, message_count,
last_date, last_message_id,
last_text, last_sender_id)
VALUES (NEW.account_id, NEW.chat_id, 1,
CASE WHEN NEW.date <= now() + interval '1 day'
THEN NEW.date END,
CASE WHEN NEW.date <= now() + interval '1 day'
THEN NEW.message_id END,
NEW.text, NEW.sender_id)
ON CONFLICT (account_id, chat_id) DO UPDATE SET
message_count = cs.message_count + 1,
last_date = CASE WHEN chat_stats_newer(cs, EXCLUDED)
THEN EXCLUDED.last_date ELSE cs.last_date END,
last_message_id = CASE WHEN chat_stats_newer(cs, EXCLUDED)
THEN EXCLUDED.last_message_id
ELSE cs.last_message_id END,
last_text = CASE WHEN chat_stats_newer(cs, EXCLUDED)
THEN EXCLUDED.last_text ELSE cs.last_text END,
last_sender_id = CASE WHEN chat_stats_newer(cs, EXCLUDED)
THEN EXCLUDED.last_sender_id
ELSE cs.last_sender_id END;
ELSE
UPDATE chat_stats
SET last_text = NEW.text, last_sender_id = NEW.sender_id
WHERE account_id = NEW.account_id
AND chat_id = NEW.chat_id
AND last_message_id = NEW.message_id;
END IF;
RETURN NULL;
END;
$$
"""
_NEWER = """
CREATE FUNCTION chat_stats_newer(current chat_stats, incoming chat_stats)
RETURNS boolean LANGUAGE sql IMMUTABLE AS $$
SELECT incoming.last_date IS NOT NULL
AND (current.last_date IS NULL
OR (incoming.last_date, incoming.last_message_id)
> (current.last_date, current.last_message_id))
$$
"""
_BACKFILL = """
INSERT INTO chat_stats (account_id, chat_id, message_count, last_date,
last_message_id, last_text, last_sender_id)
SELECT agg.account_id, agg.chat_id, agg.message_count,
last.date, last.message_id, last.text, last.sender_id
FROM (
SELECT account_id, chat_id, count(*) AS message_count
FROM messages GROUP BY account_id, chat_id
) agg
LEFT JOIN LATERAL (
SELECT date, message_id, text, sender_id FROM messages m
WHERE m.account_id = agg.account_id AND m.chat_id = agg.chat_id
AND m.date <= now() + interval '1 day'
ORDER BY m.date DESC, m.message_id DESC LIMIT 1
) last ON true
"""
def upgrade() -> None:
op.execute(
"CREATE TABLE chat_stats ("
"account_id integer NOT NULL, "
"chat_id bigint NOT NULL, "
"message_count bigint NOT NULL DEFAULT 0, "
"last_date timestamptz, "
"last_message_id bigint, "
"last_text text, "
"last_sender_id bigint, "
"PRIMARY KEY (account_id, chat_id))"
)
op.execute(_BACKFILL)
op.execute(
"CREATE INDEX ix_chat_stats_recent ON chat_stats "
"(account_id, last_date DESC, chat_id DESC)"
)
op.execute(_NEWER)
op.execute(_APPLY)
op.execute(
"CREATE TRIGGER messages_chat_stats_insert AFTER INSERT ON messages "
"FOR EACH ROW EXECUTE FUNCTION chat_stats_apply()"
)
op.execute(
"CREATE TRIGGER messages_chat_stats_update AFTER UPDATE ON messages "
"FOR EACH ROW WHEN (OLD.text IS DISTINCT FROM NEW.text "
"OR OLD.sender_id IS DISTINCT FROM NEW.sender_id) "
"EXECUTE FUNCTION chat_stats_apply()"
)
def downgrade() -> None:
op.execute("DROP TRIGGER IF EXISTS messages_chat_stats_update ON messages")
op.execute("DROP TRIGGER IF EXISTS messages_chat_stats_insert ON messages")
op.execute("DROP FUNCTION IF EXISTS chat_stats_apply()")
op.execute("DROP FUNCTION IF EXISTS chat_stats_newer(chat_stats, chat_stats)")
op.execute("DROP TABLE IF EXISTS chat_stats")
@@ -0,0 +1,75 @@
"""drop scheduled message duplicates
Revision ID: e1c7a4b62d90
Revises: d4a7e2b91f38
Create Date: 2026-08-13 10:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "e1c7a4b62d90"
down_revision: str | None = "d4a7e2b91f38"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_STAGE = """
CREATE TEMP TABLE scheduled_keys ON COMMIT DROP AS
SELECT DISTINCT s.account_id, s.chat_id, s.message_id
FROM messages s
WHERE s.raw->>'scheduled' = 'true'
AND NOT EXISTS (
SELECT 1 FROM messages r
WHERE r.account_id = s.account_id AND r.chat_id = s.chat_id
AND r.message_id = s.message_id
AND r.raw->>'scheduled' IS DISTINCT FROM 'true'
)
"""
_DELETE_CHILD = """
DELETE FROM {table} t
USING scheduled_keys k
WHERE t.account_id = k.account_id AND t.chat_id = k.chat_id
AND t.message_id = k.message_id
"""
_DELETE_MESSAGES = """
DELETE FROM messages m
USING scheduled_keys k
WHERE m.account_id = k.account_id AND m.chat_id = k.chat_id
AND m.message_id = k.message_id AND m.raw->>'scheduled' = 'true'
"""
_RECOUNT = """
UPDATE chat_stats cs
SET message_count = fresh.message_count
FROM (
SELECT k.account_id, k.chat_id,
(SELECT count(*) FROM messages m
WHERE m.account_id = k.account_id AND m.chat_id = k.chat_id)
AS message_count
FROM (SELECT DISTINCT account_id, chat_id FROM scheduled_keys) k
) fresh
WHERE cs.account_id = fresh.account_id AND cs.chat_id = fresh.chat_id
"""
_CHILD_TABLES = ("media", "media_versions", "message_versions", "links", "callbacks")
_ALLOW_DECOMPRESSION = (
"SET LOCAL timescaledb.max_tuples_decompressed_per_dml_transaction = 0"
)
def upgrade() -> None:
op.execute(_ALLOW_DECOMPRESSION)
op.execute(_STAGE)
for table in _CHILD_TABLES:
op.execute(_DELETE_CHILD.format(table=table))
op.execute(_DELETE_MESSAGES)
op.execute(_RECOUNT)
def downgrade() -> None:
pass
@@ -0,0 +1,61 @@
"""bot accounts
Revision ID: e2c9a41f7b68
Revises: c8f3b6a2d571
Create Date: 2026-09-18 16:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "e2c9a41f7b68"
down_revision: str | None = "c8f3b6a2d571"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_ACCOUNTS = """
INSERT INTO accounts
(tg_user_id, label, phone, session_name, is_active, is_bot, raw)
SELECT b.bot_id,
coalesce(b.title, '@' || b.username, b.bot_id::text),
NULL,
b.bot_id::text,
b.enabled,
true,
jsonb_build_object('id', b.bot_id, 'username', b.username,
'first_name', b.title)
FROM watched_bots b
ON CONFLICT (tg_user_id) DO NOTHING
"""
_LINK = """
UPDATE watched_bots b SET account_id = a.account_id
FROM accounts a WHERE a.tg_user_id = b.bot_id
"""
_MERGE_MESSAGES = """
INSERT INTO messages
(account_id, chat_id, message_id, date, sender_id, text, raw, has_media)
SELECT b.account_id, u.chat_id, u.item_id, u.ts, u.sender_id, u.text, u.raw,
u.media_kind IS NOT NULL
FROM bot_updates u
JOIN watched_bots b ON b.bot_id = u.bot_id
WHERE u.kind = 'message' AND b.account_id IS NOT NULL
ON CONFLICT (account_id, chat_id, message_id, date) DO NOTHING
"""
def upgrade() -> None:
op.execute("ALTER TABLE accounts ADD COLUMN is_bot boolean NOT NULL DEFAULT false")
op.execute("ALTER TABLE watched_bots ADD COLUMN account_id integer")
op.execute(_ACCOUNTS)
op.execute(_LINK)
op.execute(_MERGE_MESSAGES)
op.execute("DROP TABLE IF EXISTS bot_updates")
def downgrade() -> None:
op.execute("ALTER TABLE watched_bots DROP COLUMN IF EXISTS account_id")
op.execute("ALTER TABLE accounts DROP COLUMN IF EXISTS is_bot")
@@ -0,0 +1,73 @@
"""file shares
Revision ID: f2b8d3c9a51e
Revises: e1c7a4b62d90
Create Date: 2026-08-13 10:30:00.000000
"""
from collections.abc import Sequence
from alembic import op
revision: str = "f2b8d3c9a51e"
down_revision: str | None = "e1c7a4b62d90"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_SHARES = """
CREATE TABLE file_shares (
id serial PRIMARY KEY,
account_id integer NOT NULL,
token text NOT NULL UNIQUE,
kind text NOT NULL,
storage_key text NOT NULL,
file_name text NOT NULL,
mime text,
file_size bigint,
title text,
chat_id bigint,
message_id bigint,
peer_id bigint,
story_id bigint,
expires_at timestamptz,
max_downloads integer,
download_count integer NOT NULL DEFAULT 0,
last_download_at timestamptz,
revoked_at timestamptz,
created_at timestamptz NOT NULL DEFAULT now()
)
"""
_HITS = """
CREATE TABLE file_share_hits (
id bigserial PRIMARY KEY,
share_id integer NOT NULL REFERENCES file_shares (id) ON DELETE CASCADE,
ts timestamptz NOT NULL DEFAULT now(),
method text NOT NULL,
ip text,
user_agent text,
counted boolean NOT NULL DEFAULT false
)
"""
def upgrade() -> None:
op.execute(_SHARES)
op.execute(
"CREATE INDEX ix_file_shares_account ON file_shares "
"(account_id, created_at DESC)"
)
op.execute(
"CREATE INDEX ix_file_shares_subject ON file_shares "
"(account_id, storage_key) WHERE revoked_at IS NULL"
)
op.execute(_HITS)
op.execute(
"CREATE INDEX ix_file_share_hits_share ON file_share_hits (share_id, ts DESC)"
)
def downgrade() -> None:
op.execute("DROP TABLE IF EXISTS file_share_hits")
op.execute("DROP TABLE IF EXISTS file_shares")
+4 -3
View File
@@ -23,22 +23,23 @@ dependencies = [
] ]
[build-system] [build-system]
requires = ["uv_build>=0.11.14,<0.12.0"] requires = ["uv_build>=0.13.0,<0.14.0"]
build-backend = "uv_build" build-backend = "uv_build"
[tool.uv.build-backend] [tool.uv.build-backend]
module-name = ["userbot", "api", "utils", "dependencies"] module-name = ["userbot", "api", "watcher", "utils", "dependencies"]
[tool.ruff] [tool.ruff]
target-version = "py313" target-version = "py313"
[tool.ruff.lint] [tool.ruff.lint]
select = ["ALL"] select = ["ALL"]
ignore = ["CPY", "D1", "D203", "D212", "COM812", "BLE", "S101", "PT", "RUF001", "RUF002", "RUF003", "TRY301"] ignore = ["CPY", "D1", "D203", "D212", "COM812", "BLE", "S101", "PT", "PLR0917", "RUF001", "RUF002", "RUF003", "TRY301"]
unfixable = ["F401"] unfixable = ["F401"]
[tool.ruff.lint.per-file-ignores] [tool.ruff.lint.per-file-ignores]
"scripts/*" = ["INP"] "scripts/*" = ["INP"]
"scripts/check_imports.py" = ["T201"]
"t/*" = ["ALL"] "t/*" = ["ALL"]
"*.lock" = ["ALL"] "*.lock" = ["ALL"]
"migrations/*" = ["INP", "D", "E501"] "migrations/*" = ["INP", "D", "E501"]
+45
View File
@@ -0,0 +1,45 @@
import os
import subprocess
import sys
HEAVY = ["aiogram", "pyrogram", "fastapi", "google.genai"]
SERVICES = {
"api": (["api.app"], {"fastapi", "pyrogram"}),
"userbot": (["userbot.runner", "userbot.handlers"], {"pyrogram"}),
"watcher": (["watcher.runner"], {"pyrogram"}),
}
PROBE = """
import importlib, sys
for module in sys.argv[1].split(","):
importlib.import_module(module)
print(" ".join(name for name in sys.argv[2:] if name in sys.modules))
"""
ENV = {**os.environ, "AUTH__TOKEN": "check"}
def loaded(modules: list[str]) -> tuple[set[str], str]:
result = subprocess.run( # noqa: S603
[sys.executable, "-c", PROBE, ",".join(modules), *HEAVY],
capture_output=True,
text=True,
cwd="src",
env=ENV,
check=False,
)
error = (result.stderr.strip().splitlines() or ["exit code"])[-1]
return set(result.stdout.split()) & set(HEAVY), error if result.returncode else ""
def main() -> int:
failed = False
for service, (modules, allowed) in SERVICES.items():
found, error = loaded(modules)
extra = sorted(found - allowed)
status = error or ("imports " + ", ".join(extra) if extra else "")
print(f"{service:10s} {'FAIL: ' + status if status else 'ok'}")
failed |= bool(status)
return int(failed)
if __name__ == "__main__":
sys.exit(main())
+23
View File
@@ -0,0 +1,23 @@
import logging
import re
QUERYLESS_PREFIXES = ("/mcp",)
_PATH_ARG = 2
_TOKEN = re.compile(r"(?<=[?&]token=)[^&]+")
class ScrubQueryFilter(logging.Filter):
def filter(self, record: logging.LogRecord) -> bool:
args = record.args
if isinstance(args, tuple) and len(args) > _PATH_ARG:
path = str(args[_PATH_ARG])
if path.startswith(QUERYLESS_PREFIXES):
path = path.partition("?")[0]
else:
path = _TOKEN.sub("***", path)
record.args = (*args[:_PATH_ARG], path, *args[_PATH_ARG + 1 :])
return True
def scrub_access_log() -> None:
logging.getLogger("uvicorn.access").addFilter(ScrubQueryFilter())
+24 -2
View File
@@ -1,3 +1,4 @@
import asyncio
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from pathlib import Path from pathlib import Path
@@ -9,6 +10,7 @@ from fastapi.responses import FileResponse, RedirectResponse
from fastmcp.utilities.lifespan import combine_lifespans from fastmcp.utilities.lifespan import combine_lifespans
from starlette.applications import Starlette from starlette.applications import Starlette
from api.access_log import scrub_access_log
from api.auth import BearerAuthMiddleware from api.auth import BearerAuthMiddleware
from api.mcp.server import mcp from api.mcp.server import mcp
from api.realtime import hub from api.realtime import hub
@@ -18,10 +20,12 @@ from api.routers import (
annotations, annotations,
avatars, avatars,
backfill, backfill,
bots,
chats, chats,
custom_emoji, custom_emoji,
discover, discover,
events, events,
files,
folders, folders,
media, media,
peers, peers,
@@ -29,18 +33,25 @@ from api.routers import (
presence, presence,
profile, profile,
search, search,
shares,
social, social,
storage,
stories, stories,
watches, watches,
) )
from dependencies.container import container from dependencies.container import container
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS, NO_STORE_HEADERS
from utils.env import env from utils.env import env
from utils.storage import ContentAddressedStorage
from utils.trash import purge_loop
if env.auth.token is None: if env.auth.token is None:
msg = "AUTH__TOKEN is required for the api process" msg = "AUTH__TOKEN is required for the api process"
raise RuntimeError(msg) raise RuntimeError(msg)
_token = env.auth.token.get_secret_value() _token = env.auth.token.get_secret_value()
scrub_access_log()
mcp_app = mcp.http_app(path="/") mcp_app = mcp.http_app(path="/")
@@ -48,7 +59,11 @@ mcp_app = mcp.http_app(path="/")
async def lifespan(app_: Starlette) -> AsyncGenerator[None]: async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
pool = await container.get(asyncpg.Pool) pool = await container.get(asyncpg.Pool)
await hub.start(pool) await hub.start(pool)
purger = asyncio.create_task(
purge_loop(pool, await container.get(ContentAddressedStorage))
)
yield yield
purger.cancel()
await hub.stop() await hub.stop()
await app_.state.dishka_container.close() await app_.state.dishka_container.close()
@@ -70,6 +85,7 @@ async def health(pool: FromDishka[asyncpg.Pool]) -> dict[str, bool]:
app.include_router(accounts.router) app.include_router(accounts.router)
app.include_router(analytics.router) app.include_router(analytics.router)
app.include_router(bots.router)
app.include_router(policy.router) app.include_router(policy.router)
app.include_router(folders.router) app.include_router(folders.router)
app.include_router(backfill.router) app.include_router(backfill.router)
@@ -87,6 +103,9 @@ app.include_router(peers.router)
app.include_router(discover.router) app.include_router(discover.router)
app.include_router(annotations.router) app.include_router(annotations.router)
app.include_router(watches.router) app.include_router(watches.router)
app.include_router(shares.router)
app.include_router(storage.router)
app.include_router(files.router)
app.mount("/mcp", mcp_app) app.mount("/mcp", mcp_app)
@@ -105,8 +124,11 @@ if _spa_dir.is_dir():
async def serve_spa(spa_path: str) -> FileResponse: async def serve_spa(spa_path: str) -> FileResponse:
candidate = (_spa_dir / spa_path).resolve() candidate = (_spa_dir / spa_path).resolve()
if spa_path and candidate.is_relative_to(_spa_dir) and candidate.is_file(): if spa_path and candidate.is_relative_to(_spa_dir) and candidate.is_file():
return FileResponse(candidate) immutable = spa_path.startswith("_app/immutable/")
return FileResponse(_spa_index) return FileResponse(
candidate, headers=IMMUTABLE_HEADERS if immutable else DAY_HEADERS
)
return FileResponse(_spa_index, headers=NO_STORE_HEADERS)
app.add_middleware(BearerAuthMiddleware, token=_token) app.add_middleware(BearerAuthMiddleware, token=_token)
+1 -1
View File
@@ -9,7 +9,7 @@ from pyrogram.errors import AuthTokenExpired, SessionPasswordNeeded
from pyrogram.qrlogin import QRLogin from pyrogram.qrlogin import QRLogin
from pyrogram.types import User from pyrogram.types import User
from userbot import PyroClient from userbot.modules.client import PyroClient
from utils.env import env from utils.env import env
LOGIN_TTL_SECONDS = 900 LOGIN_TTL_SECONDS = 900
+113
View File
@@ -0,0 +1,113 @@
import asyncio
from pathlib import Path
from typing import Any
import asyncpg
from fastmcp.exceptions import ToolError
from fastmcp.server.dependencies import get_http_request
from fastmcp.utilities.types import Image
from mcp.types import TextContent
from utils.files import expiry_from_seconds
from utils.jobs import enqueue
from utils.read import chats, media, shares
from utils.read.models import FileShareView, MediaView
from utils.storage import ContentAddressedStorage
LINK_TTL_SECONDS = 24 * 3600
MAX_WAIT_SECONDS = 120
_POLL_SECONDS = 1.0
_INLINE_IMAGE_BYTES = 5 * 1024 * 1024
_INLINE_TEXT_BYTES = 256 * 1024
_INLINE_IMAGE_MIMES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
_TEXT_MIMES = {"application/json", "application/xml", "application/x-yaml"}
_FAILED_JOB = ("failed", "canceled")
async def _job_error(pool: asyncpg.Pool, job_id: int) -> str | None:
row = await pool.fetchrow("SELECT status, error FROM jobs WHERE id = $1", job_id)
if row is None or row["status"] in _FAILED_JOB:
return (row and row["error"]) or "download job is gone"
return None
async def fetch_media(
pool: asyncpg.Pool,
account_id: int,
chat_id: int,
message_id: int,
wait_seconds: float,
) -> MediaView | None:
item = await media.get_message_media(pool, account_id, chat_id, message_id)
if item is None:
msg = f"message {chat_id}/{message_id} has no media"
raise ToolError(msg)
if item.downloaded and item.storage_key:
return item
job_id = await enqueue(
pool, account_id, "fetch_media", {"chat_id": chat_id, "message_id": message_id}
)
loop = asyncio.get_running_loop()
deadline = loop.time() + wait_seconds
while loop.time() < deadline:
await asyncio.sleep(_POLL_SECONDS)
item = await media.get_message_media(pool, account_id, chat_id, message_id)
if item is not None and item.downloaded and item.storage_key:
return item
error = await _job_error(pool, job_id)
if error is not None:
msg = f"download failed: {error}"
raise ToolError(msg)
return None
def _origin() -> str:
request = get_http_request()
proto = request.headers.get("x-forwarded-proto", request.url.scheme)
return f"{proto}://{request.url.netloc}"
async def share_media(pool: asyncpg.Pool, item: MediaView, mime: str) -> FileShareView:
storage_key = item.storage_key or ""
share = await shares.find_active_share(pool, item.account_id, storage_key)
if share is None:
chat = await chats.get_chat(pool, item.account_id, item.chat_id)
share = await shares.create_share(
pool,
item.account_id,
"media",
storage_key,
item.file_name or f"media_{item.id}",
mime=mime,
file_size=item.file_size,
title=chat.title if chat else None,
chat_id=item.chat_id,
message_id=item.message_id,
expires_at=expiry_from_seconds(LINK_TTL_SECONDS),
)
return share
def _is_text(mime: str) -> bool:
return mime.startswith("text/") or mime in _TEXT_MIMES
def inline_content(storage: ContentAddressedStorage, key: str, mime: str) -> list[Any]:
size = Path(storage.url(key)).stat().st_size
if mime in _INLINE_IMAGE_MIMES and size <= _INLINE_IMAGE_BYTES:
return [Image(data=storage.get(key), format=mime.rsplit("/", maxsplit=1)[-1])]
if _is_text(mime) and size <= _INLINE_TEXT_BYTES:
text = storage.get(key).decode("utf-8", errors="replace")
return [TextContent(type="text", text=text)]
return []
def describe(item: MediaView, mime: str, share: FileShareView) -> str:
size = f"{item.file_size} bytes" if item.file_size else "size unknown"
expires = share.expires_at.isoformat() if share.expires_at else "never"
return (
f"{item.file_name} ({item.kind}, {mime}, {size})\n"
f"Download: {_origin()}/f/{share.token}/{share.url_name}\n"
f"Public link, expires {expires}; fetch it with curl/wget or any HTTP "
"client, no auth needed."
)
+79 -18
View File
@@ -4,13 +4,24 @@ from typing import Any
import asyncpg import asyncpg
from fastmcp import FastMCP from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp.tools.base import ToolResult
from mcp.types import CallToolRequestParams, TextContent
from pydantic import BaseModel from pydantic import BaseModel
from api.mcp.download import (
MAX_WAIT_SECONDS,
describe,
fetch_media,
inline_content,
share_media,
)
from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names
from dependencies.container import container from dependencies.container import container
from utils.jobs import enqueue from utils.files import resolve_mime
from utils.read import annotations, chats, media, peers, presence, social, watches from utils.read import annotations, chats, media, peers, presence, social, watches
from utils.read.accounts import self_user_id from utils.read.accounts import self_user_id, unknown_account_hint
from utils.read.models import DEFAULT_LIMIT, Page from utils.read.models import DEFAULT_LIMIT, Page
from utils.search.models import SearchFilters from utils.search.models import SearchFilters
from utils.search.repository import search_messages from utils.search.repository import search_messages
@@ -23,6 +34,9 @@ peer history) into Postgres and exposes it read-only over these tools.
Account scoping: Account scoping:
- Every tool takes `account_id`. It selects which archived Telegram account to - Every tool takes `account_id`. It selects which archived Telegram account to
read. Unless the user names a different one, always pass `account_id=1`. read. Unless the user names a different one, always pass `account_id=1`.
- `account_id` is beavergram's own small serial (1, 2, ...), never the archived
account's Telegram user id. Passing a Telegram id returns an error naming the
right `account_id`.
Identifiers: Identifiers:
- `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels, - `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels,
@@ -37,7 +51,7 @@ Typical flows:
`before_id` (first id you saw), to newer with `after_id` (last id you saw); `before_id` (first id you saw), to newer with `after_id` (last id you saw);
begin a full forward walk with `after_id=0`. Each transcript line starts begin a full forward walk with `after_id=0`. Each transcript line starts
with its `#message_id`, and photos show as `[photo #message_id]` — pass that with its `#message_id`, and photos show as `[photo #message_id]` — pass that
id to `get_media` (with `fetch=True` to download) or to annotation tools. id to `download_media` or to annotation tools.
Lines marked Lines marked
"📝 [your private note, NOT in Telegram]" are the user's own annotations "📝 [your private note, NOT in Telegram]" are the user's own annotations
attached locally via the web UI; they never existed in Telegram. Treat them attached locally via the web UI; they never existed in Telegram. Treat them
@@ -46,12 +60,21 @@ Typical flows:
transcripts; supports chat/sender/date filters and regex). transcripts; supports chat/sender/date filters and regex).
- Forensics: `get_deleted_messages` and `get_message_versions` recover content - Forensics: `get_deleted_messages` and `get_message_versions` recover content
removed or edited in Telegram but kept in the archive. removed or edited in Telegram but kept in the archive.
- Media: `get_media` returns metadata; pass `fetch=True` to enqueue a lazy - Files: `download_media` gets any message attachment (photo, video, voice,
download if the file isn't stored yet. document, ...). It pulls the file from Telegram if the archive lacks it and
returns a public link valid 24h that needs no auth (`curl -O <link>`), plus
the content itself for small images and text files. `get_media` returns
metadata only.
- Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and - Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and
`list_alerts` reads what those rules fired. `list_alerts` reads what those rules fired.
Writes: everything is read-only except `set_watch`, the only allowed write. Watched bots:
- a bot whose token the user added is just another archived account: it has its
own `account_id` and its chats, messages and media read through the same
tools as a user account. Nothing bot-specific to call.
Writes: everything is read-only except `set_watch` and the share links
`download_media` creates (visible and revocable in the web UI).
Paging: list/search tools accept `limit` and `offset`. Paging: list/search tools accept `limit` and `offset`.
""" """
@@ -63,6 +86,25 @@ async def _pool() -> asyncpg.Pool:
return await container.get(asyncpg.Pool) return await container.get(asyncpg.Pool)
class AccountGuard(Middleware):
"""Reject tool calls scoped to an account_id that does not exist."""
async def on_call_tool(
self,
context: MiddlewareContext[CallToolRequestParams],
call_next: CallNext[CallToolRequestParams, ToolResult],
) -> ToolResult:
account_id = (context.message.arguments or {}).get("account_id")
if isinstance(account_id, int) and not isinstance(account_id, bool):
hint = await unknown_account_hint(await _pool(), account_id)
if hint is not None:
raise ToolError(hint)
return await call_next(context)
mcp.add_middleware(AccountGuard())
def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]: def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]:
return [item.model_dump(mode="json") for item in items] return [item.model_dump(mode="json") for item in items]
@@ -208,21 +250,40 @@ async def get_message_versions(
@mcp.tool @mcp.tool
async def get_media( async def get_media(
account_id: int, chat_id: int, message_id: int, fetch: bool = False account_id: int, chat_id: int, message_id: int
) -> dict[str, Any] | None: ) -> dict[str, Any] | None:
"""Get media metadata for a message; set fetch=True to enqueue lazy download.""" """Get media metadata for a message; use download_media to get the file."""
item = await media.get_message_media(await _pool(), account_id, chat_id, message_id)
return item.model_dump(mode="json") if item else None
@mcp.tool(output_schema=None)
async def download_media(
account_id: int, chat_id: int, message_id: int, wait_seconds: int = 45
) -> list[Any]:
"""Download the file attached to a message: any photo, video, voice, document.
Pulls it from Telegram first if the archive doesn't have it yet, waiting up
to `wait_seconds`. Returns a public download link (valid 24h, no auth) and
inlines small images and text files.
"""
pool = await _pool() pool = await _pool()
item = await media.get_message_media(pool, account_id, chat_id, message_id) wait = min(max(wait_seconds, 0), MAX_WAIT_SECONDS)
if item is None: item = await fetch_media(pool, account_id, chat_id, message_id, wait)
return None if item is None or item.storage_key is None:
if fetch and not item.downloaded: return [
await enqueue( TextContent(
pool, type="text",
account_id, text="Still downloading from Telegram; call download_media again.",
"fetch_media",
{"chat_id": chat_id, "message_id": message_id},
) )
return item.model_dump(mode="json") ]
mime = resolve_mime(item.kind, item.mime, item.file_name)
share = await share_media(pool, item, mime)
storage = await container.get(ContentAddressedStorage)
return [
TextContent(type="text", text=describe(item, mime, share)),
*inline_content(storage, item.storage_key, mime),
]
@mcp.tool @mcp.tool
+6 -1
View File
@@ -65,11 +65,16 @@ class EventHub:
return return
account_id = event.get("account_id") account_id = event.get("account_id")
chat_id = event.get("chat_id") chat_id = event.get("chat_id")
scoped = event.get("kind") == "presence"
targets = [ targets = [
sub sub
for sub in self._subscribers for sub in self._subscribers
if sub.account_id == account_id if sub.account_id == account_id
and (sub.chat_id is None or sub.chat_id == chat_id) and (
sub.chat_id == chat_id
if scoped
else sub.chat_id is None or sub.chat_id == chat_id
)
] ]
if not targets: if not targets:
return return
+1 -1
View File
@@ -9,7 +9,7 @@ from pyrogram.errors import FloodWait, RPCError
from pyrogram.types import User from pyrogram.types import User
from api.login import LoginError, login_manager from api.login import LoginError, login_manager
from userbot import DEVICE_MODEL_LIMIT from userbot.modules.client import DEVICE_MODEL_LIMIT
from utils.read import accounts from utils.read import accounts
from utils.read.models import AccountView from utils.read.models import AccountView
+4 -1
View File
@@ -5,6 +5,7 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from utils.cache import IMMUTABLE_HEADERS, SHORT_HEADERS
from utils.jobs import enqueue from utils.jobs import enqueue
from utils.read.avatars import avatar_by_unique_id, avatar_history, current_avatar from utils.read.avatars import avatar_by_unique_id, avatar_history, current_avatar
from utils.read.models import AvatarHistoryView from utils.read.models import AvatarHistoryView
@@ -52,5 +53,7 @@ async def serve_avatar(
) )
raise HTTPException(status_code=409, detail="avatar not downloaded; fetching") raise HTTPException(status_code=409, detail="avatar not downloaded; fetching")
return FileResponse( return FileResponse(
storage.url(avatar.storage_key), media_type=avatar.mime or "image/jpeg" storage.url(avatar.storage_key),
media_type=avatar.mime or "image/jpeg",
headers=IMMUTABLE_HEADERS if unique_id is not None else SHORT_HEADERS,
) )
+30 -3
View File
@@ -8,6 +8,7 @@ from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel from pydantic import BaseModel
from utils.jobs import enqueue from utils.jobs import enqueue
from utils.read.models import DEFAULT_LIMIT, Page
router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute) router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute)
@@ -16,6 +17,12 @@ class BackfillRequest(BaseModel):
account_id: int account_id: int
chat_id: int chat_id: int
media: bool = False media: bool = False
full: bool = False
class StoriesBackfillRequest(BaseModel):
account_id: int
peer_id: int
class FetchMediaRequest(BaseModel): class FetchMediaRequest(BaseModel):
@@ -70,7 +77,17 @@ async def enqueue_backfill(
pool, pool,
body.account_id, body.account_id,
"backfill", "backfill",
{"chat_id": body.chat_id, "media": body.media}, {"chat_id": body.chat_id, "media": body.media, "full": body.full},
)
return EnqueueResponse(job_id=job_id)
@router.post("/stories/backfill", status_code=201)
async def enqueue_stories_backfill(
pool: FromDishka[asyncpg.Pool], body: StoriesBackfillRequest
) -> EnqueueResponse:
job_id = await enqueue(
pool, body.account_id, "backfill_stories", {"peer_id": body.peer_id}
) )
return EnqueueResponse(job_id=job_id) return EnqueueResponse(job_id=job_id)
@@ -114,16 +131,26 @@ async def list_jobs(
pool: FromDishka[asyncpg.Pool], pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()], account_id: Annotated[int, Query()],
status: Annotated[str | None, Query()] = None, status: Annotated[str | None, Query()] = None,
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[JobView]: ) -> list[JobView]:
page = Page(limit=limit, offset=offset)
if status is None: if status is None:
rows = await pool.fetch( rows = await pool.fetch(
"SELECT * FROM jobs WHERE account_id = $1 ORDER BY id DESC", account_id "SELECT * FROM jobs WHERE account_id = $1 "
"ORDER BY id DESC LIMIT $2 OFFSET $3",
account_id,
page.capped_limit,
page.offset,
) )
else: else:
rows = await pool.fetch( rows = await pool.fetch(
"SELECT * FROM jobs WHERE account_id = $1 AND status = $2 ORDER BY id DESC", "SELECT * FROM jobs WHERE account_id = $1 AND status = $2 "
"ORDER BY id DESC LIMIT $3 OFFSET $4",
account_id, account_id,
status, status,
page.capped_limit,
page.offset,
) )
return [_to_view(row) for row in rows] return [_to_view(row) for row in rows]
+42
View File
@@ -0,0 +1,42 @@
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from utils.read import bots
from utils.read.models import BotView
from watcher import tokens
router = APIRouter(prefix="/api/bots", tags=["bots"], route_class=DishkaRoute)
INVALID_TOKEN = "token must look like 123456:ABC-DEF..." # noqa: S105
class BotAdd(BaseModel):
token: str
@router.get("")
async def list_bots(pool: FromDishka[asyncpg.Pool]) -> list[BotView]:
return await bots.list_bots(pool)
@router.post("", status_code=201)
async def add_bot(pool: FromDishka[asyncpg.Pool], body: BotAdd) -> BotView:
bot_id = tokens.bot_id_of(body.token)
if bot_id is None:
raise HTTPException(status_code=422, detail=INVALID_TOKEN)
tokens.save(bot_id, body.token)
await bots.register_bot(pool, bot_id)
await bots.notify_bots_changed(pool)
bot = await bots.get_bot(pool, bot_id)
if bot is None:
raise HTTPException(status_code=500, detail="bot was not registered")
return bot
@router.delete("/{bot_id}", status_code=204)
async def remove_bot(pool: FromDishka[asyncpg.Pool], bot_id: int) -> None:
if not await bots.remove_bot(pool, bot_id):
raise HTTPException(status_code=404, detail="unknown bot")
await bots.notify_bots_changed(pool)
+27 -1
View File
@@ -6,9 +6,11 @@ from fastapi import APIRouter, Query
from pydantic import BaseModel from pydantic import BaseModel
from utils.jobs import enqueue from utils.jobs import enqueue
from utils.policy import repository
from utils.read import chats from utils.read import chats
from utils.read.models import ( from utils.read.models import (
DEFAULT_LIMIT, DEFAULT_LIMIT,
MAX_LIMIT,
ChatListItem, ChatListItem,
MessageVersionView, MessageVersionView,
MessageView, MessageView,
@@ -35,8 +37,32 @@ async def list_chats(
account_id: AccountId, account_id: AccountId,
limit: Limit = DEFAULT_LIMIT, limit: Limit = DEFAULT_LIMIT,
offset: Offset = 0, offset: Offset = 0,
folder_id: Annotated[int | None, Query()] = None,
search: Annotated[str | None, Query()] = None,
) -> list[ChatListItem]: ) -> list[ChatListItem]:
return await chats.list_chats(pool, account_id, Page(limit=limit, offset=offset)) folder = (
await repository.get_folder(pool, account_id, folder_id)
if folder_id is not None
else None
)
return await chats.list_chats(
pool, account_id, Page(limit=limit, offset=offset), folder=folder, search=search
)
@router.get("/chats/batch")
async def get_chats_batch(
pool: FromDishka[asyncpg.Pool], account_id: AccountId, ids: Annotated[str, Query()]
) -> list[ChatListItem]:
parsed = [int(part) for part in ids.split(",") if part.strip()]
return await chats.get_chats(pool, account_id, parsed[:MAX_LIMIT])
@router.get("/chats/{chat_id}")
async def get_chat(
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
) -> ChatListItem | None:
return await chats.get_chat(pool, account_id, chat_id)
@router.get("/chats/{chat_id}/messages") @router.get("/chats/{chat_id}/messages")
+6 -2
View File
@@ -13,7 +13,7 @@ from utils.policy import repository as policy_repository
from utils.policy.defaults import TRACKING from utils.policy.defaults import TRACKING
from utils.policy.models import ScopeType from utils.policy.models import ScopeType
from utils.read import discover from utils.read import discover
from utils.read.models import DiscoverItem from utils.read.models import MAX_LIMIT, DiscoverItem
router = APIRouter(prefix="/api", tags=["discover"], route_class=DishkaRoute) router = APIRouter(prefix="/api", tags=["discover"], route_class=DishkaRoute)
@@ -68,6 +68,7 @@ async def discover_peers(
) -> list[DiscoverItem]: ) -> list[DiscoverItem]:
if not query.strip(): if not query.strip():
return [] return []
limit = min(limit, MAX_LIMIT)
items = await discover.search(pool, account_id, query, limit) items = await discover.search(pool, account_id, query, limit)
if not remote: if not remote:
return items return items
@@ -103,7 +104,10 @@ async def track_chat(
await enqueue(pool, body.account_id, "enrich_chat", {"chat_id": chat_id}) await enqueue(pool, body.account_id, "enrich_chat", {"chat_id": chat_id})
if body.backfill: if body.backfill:
await enqueue( await enqueue(
pool, body.account_id, "backfill", {"chat_id": chat_id, "media": True} pool,
body.account_id,
"backfill",
{"chat_id": chat_id, "media": True, "full": True},
) )
return await discover.get_item(pool, body.account_id, chat_id) return await discover.get_item(pool, body.account_id, chat_id)
+113
View File
@@ -0,0 +1,113 @@
from datetime import UTC, datetime
from typing import Annotated
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, Query, Request
from fastapi.responses import FileResponse, PlainTextResponse
from utils.files import (
content_disposition,
counts_as_download,
is_active_mime,
is_inline_mime,
resolve_mime,
)
from utils.read import shares
from utils.storage import ContentAddressedStorage
router = APIRouter(tags=["files"], route_class=DishkaRoute)
_GONE = "This link is no longer available.\n"
_MISSING = "Not found.\n"
_NO_STORE = {
"Cache-Control": "private, no-store",
"X-Robots-Tag": "noindex, nofollow",
"X-Content-Type-Options": "nosniff",
}
_SANDBOX = {"Content-Security-Policy": "sandbox"}
def _client_ip(request: Request) -> str | None:
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
return forwarded.split(",")[0].strip()
return request.client.host if request.client else None
def _expired(row: asyncpg.Record) -> bool:
expires_at = row["expires_at"]
return expires_at is not None and expires_at <= datetime.now(UTC)
def _exhausted(row: asyncpg.Record) -> bool:
limit = row["max_downloads"]
return limit is not None and row["download_count"] >= limit
async def _serve(
request: Request,
pool: asyncpg.Pool,
storage: ContentAddressedStorage,
token: str,
*,
dl: bool,
) -> FileResponse | PlainTextResponse:
row = await shares.share_by_token(pool, token)
if row is None:
return PlainTextResponse(_MISSING, status_code=404, headers=_NO_STORE)
method = request.method
agent = request.headers.get("user-agent")
ip = _client_ip(request)
counts = counts_as_download(method, agent, request.headers.get("range"))
if row["revoked_at"] is not None or _expired(row) or _exhausted(row):
await shares.record_hit(pool, row["id"], method, ip, agent, counted=False)
return PlainTextResponse(_GONE, status_code=410, headers=_NO_STORE)
if not storage.exists(row["storage_key"]):
return PlainTextResponse(_MISSING, status_code=404, headers=_NO_STORE)
if counts and not await shares.consume_download(pool, row["id"]):
await shares.record_hit(pool, row["id"], method, ip, agent, counted=False)
return PlainTextResponse(_GONE, status_code=410, headers=_NO_STORE)
await shares.record_hit(pool, row["id"], method, ip, agent, counted=counts)
mime = resolve_mime(row["kind"], row["mime"], row["file_name"])
attachment = dl or not is_inline_mime(mime)
return FileResponse(
storage.url(row["storage_key"]),
media_type=mime,
headers={
**_NO_STORE,
**(_SANDBOX if is_active_mime(mime) else {}),
"Content-Disposition": content_disposition(
row["file_name"], attachment=attachment
),
},
)
@router.api_route("/f/{token}", methods=["GET", "HEAD"], response_model=None)
async def serve_shared_file(
request: Request,
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
token: str,
dl: Annotated[bool, Query()] = False,
) -> FileResponse | PlainTextResponse:
return await _serve(request, pool, storage, token, dl=dl)
@router.api_route("/f/{token}/{name}", methods=["GET", "HEAD"], response_model=None)
async def serve_shared_file_named(
request: Request,
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
token: str,
name: str, # noqa: ARG001
dl: Annotated[bool, Query()] = False,
) -> FileResponse | PlainTextResponse:
return await _serve(request, pool, storage, token, dl=dl)
+28 -3
View File
@@ -5,6 +5,8 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS
from utils.files import content_disposition, media_file_name, resolve_mime
from utils.read.media import ( from utils.read.media import (
get_media, get_media,
get_media_version, get_media_version,
@@ -16,6 +18,15 @@ from utils.storage import ContentAddressedStorage
router = APIRouter(prefix="/api/media", tags=["media"], route_class=DishkaRoute) router = APIRouter(prefix="/api/media", tags=["media"], route_class=DishkaRoute)
Download = Annotated[bool, Query()]
def _attachment(headers: dict[str, str], file_name: str) -> dict[str, str]:
return {
**headers,
"Content-Disposition": content_disposition(file_name, attachment=True),
}
@router.get("/{media_id}/meta") @router.get("/{media_id}/meta")
async def media_meta(pool: FromDishka[asyncpg.Pool], media_id: int) -> MediaView: async def media_meta(pool: FromDishka[asyncpg.Pool], media_id: int) -> MediaView:
@@ -40,13 +51,20 @@ async def serve_media_version(
pool: FromDishka[asyncpg.Pool], pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage], storage: FromDishka[ContentAddressedStorage],
version_id: int, version_id: int,
download: Download = False,
) -> FileResponse: ) -> FileResponse:
version = await get_media_version(pool, version_id) version = await get_media_version(pool, version_id)
if version is None: if version is None or not storage.exists(version.storage_key):
raise HTTPException(status_code=404, detail="media version not found") raise HTTPException(status_code=404, detail="media version not found")
headers = IMMUTABLE_HEADERS
if download:
headers = _attachment(
headers, media_file_name(version.kind, version.mime, version_id)
)
return FileResponse( return FileResponse(
storage.url(version.storage_key), storage.url(version.storage_key),
media_type=version.mime or "application/octet-stream", media_type=resolve_mime(version.kind, version.mime),
headers=headers,
) )
@@ -68,6 +86,7 @@ async def serve_media(
pool: FromDishka[asyncpg.Pool], pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage], storage: FromDishka[ContentAddressedStorage],
media_id: int, media_id: int,
download: Download = False,
) -> FileResponse: ) -> FileResponse:
media = await get_media(pool, media_id) media = await get_media(pool, media_id)
if media is None: if media is None:
@@ -77,7 +96,13 @@ async def serve_media(
status_code=409, status_code=409,
detail="media not downloaded; enqueue fetch via POST /api/media/fetch", detail="media not downloaded; enqueue fetch via POST /api/media/fetch",
) )
if not storage.exists(media.storage_key):
raise HTTPException(status_code=404, detail="media file is missing")
headers = DAY_HEADERS
if download:
headers = _attachment(headers, media.file_name or f"media_{media_id}")
return FileResponse( return FileResponse(
storage.url(media.storage_key), storage.url(media.storage_key),
media_type=media.mime or "application/octet-stream", media_type=resolve_mime(media.kind, media.mime, media.file_name),
headers=headers,
) )
+29 -1
View File
@@ -5,13 +5,15 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel from pydantic import BaseModel
from utils.policy import repository from utils.policy import overview, repository
from utils.policy.models import ( from utils.policy.models import (
CaptureToggles, CaptureToggles,
ChatKind, ChatKind,
ChatMeta, ChatMeta,
PolicyCreate, PolicyCreate,
PolicyRecord, PolicyRecord,
PolicyRuleView,
PolicyScopeView,
) )
from utils.policy.resolver import resolve from utils.policy.resolver import resolve
@@ -50,6 +52,32 @@ async def effective_policy(
return resolve(chat, folders, policies) return resolve(chat, folders, policies)
@router.get("/overview")
async def policy_overview(
pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()]
) -> list[PolicyRuleView]:
return await overview.list_rules(pool, account_id)
@router.get("/scope")
async def policy_scope(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
chat_id: Annotated[int | None, Query()] = None,
folder_id: Annotated[int | None, Query()] = None,
) -> PolicyScopeView:
if chat_id is not None:
return await overview.chat_scope(pool, account_id, chat_id)
scope = (
await overview.folder_scope(pool, account_id, folder_id)
if folder_id is not None
else None
)
if scope is None:
raise HTTPException(status_code=404, detail="scope not found")
return scope
@router.post("", status_code=201) @router.post("", status_code=201)
async def create_policy( async def create_policy(
pool: FromDishka[asyncpg.Pool], body: PolicyCreate pool: FromDishka[asyncpg.Pool], body: PolicyCreate
+8 -2
View File
@@ -52,9 +52,15 @@ async def chat_links(
@router.get("/calendar") @router.get("/calendar")
async def chat_calendar( async def chat_calendar(
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId pool: FromDishka[asyncpg.Pool],
chat_id: int,
account_id: AccountId,
date_from: Annotated[datetime | None, Query()] = None,
date_to: Annotated[datetime | None, Query()] = None,
) -> list[DayCount]: ) -> list[DayCount]:
return await profile.daily_counts(pool, account_id, chat_id) return await profile.daily_counts(
pool, account_id, chat_id, date_from=date_from, date_to=date_to
)
@router.get("/message-at") @router.get("/message-at")
+233
View File
@@ -0,0 +1,233 @@
from typing import Annotated
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel
from utils.files import (
expiry_from_seconds,
media_file_name,
resolve_mime,
story_file_name,
)
from utils.read import chats, peers, shares
from utils.read.media import get_media, get_media_version
from utils.read.models import DEFAULT_LIMIT, FileShareHitView, FileShareView, Page
router = APIRouter(prefix="/api/shares", tags=["shares"], route_class=DishkaRoute)
class ShareCreate(BaseModel):
account_id: int
kind: str
media_id: int | None = None
version_id: int | None = None
peer_id: int | None = None
story_id: int | None = None
expires_in_seconds: int | None = None
max_downloads: int | None = None
class ShareUpdate(BaseModel):
expires_in_seconds: int | None = None
max_downloads: int | None = None
keep_expiry: bool = False
class Subject(BaseModel):
storage_key: str
file_name: str
mime: str | None
file_size: int | None
title: str | None
chat_id: int | None = None
message_id: int | None = None
peer_id: int | None = None
story_id: int | None = None
_NOT_STORED = "file is not downloaded yet"
async def _chat_title(pool: asyncpg.Pool, account_id: int, chat_id: int) -> str | None:
chat = await chats.get_chat(pool, account_id, chat_id)
return chat.title if chat else None
async def _media_subject(pool: asyncpg.Pool, media_id: int) -> Subject:
media = await get_media(pool, media_id)
if media is None:
raise HTTPException(status_code=404, detail="media not found")
if not media.downloaded or media.storage_key is None:
raise HTTPException(status_code=409, detail=_NOT_STORED)
file_name = media.file_name or f"media_{media.id}"
return Subject(
storage_key=media.storage_key,
file_name=file_name,
mime=resolve_mime(media.kind, media.mime, file_name),
file_size=media.file_size,
title=await _chat_title(pool, media.account_id, media.chat_id),
chat_id=media.chat_id,
message_id=media.message_id,
)
async def _version_subject(pool: asyncpg.Pool, version_id: int) -> Subject:
version = await get_media_version(pool, version_id)
if version is None:
raise HTTPException(status_code=404, detail="media version not found")
version_name = media_file_name(version.kind, version.mime, version_id)
return Subject(
storage_key=version.storage_key,
file_name=version_name,
mime=resolve_mime(version.kind, version.mime, version_name),
file_size=version.file_size,
title=None,
chat_id=None,
message_id=None,
)
async def _story_subject(
pool: asyncpg.Pool, account_id: int, peer_id: int, story_id: int
) -> Subject:
story = await peers.get_story(pool, account_id, peer_id, story_id)
if story is None:
raise HTTPException(status_code=404, detail="story not found")
if not story.downloaded or story.storage_key is None:
raise HTTPException(status_code=409, detail=_NOT_STORED)
return Subject(
storage_key=story.storage_key,
file_name=story_file_name(peer_id, story_id, story.media_kind),
mime="video/mp4" if story.media_kind == "video" else "image/jpeg",
file_size=None,
title=await _chat_title(pool, account_id, peer_id),
peer_id=peer_id,
story_id=story_id,
)
async def _resolve(pool: asyncpg.Pool, body: ShareCreate) -> Subject:
if body.kind == "media" and body.media_id is not None:
return await _media_subject(pool, body.media_id)
if body.kind == "media_version" and body.version_id is not None:
return await _version_subject(pool, body.version_id)
if body.kind == "story" and body.peer_id is not None and body.story_id is not None:
return await _story_subject(pool, body.account_id, body.peer_id, body.story_id)
raise HTTPException(status_code=422, detail="unsupported share subject")
@router.get("")
async def list_shares(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
active_only: Annotated[bool, Query()] = False,
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[FileShareView]:
return await shares.list_shares(
pool, account_id, Page(limit=limit, offset=offset), active_only=active_only
)
@router.post("", status_code=201)
async def create_share(
pool: FromDishka[asyncpg.Pool], body: ShareCreate
) -> FileShareView:
subject = await _resolve(pool, body)
existing = await shares.find_active_share(
pool, body.account_id, subject.storage_key
)
if existing is not None:
return existing
return await shares.create_share(
pool,
body.account_id,
body.kind,
subject.storage_key,
subject.file_name,
mime=subject.mime,
file_size=subject.file_size,
title=subject.title,
chat_id=subject.chat_id,
message_id=subject.message_id,
peer_id=subject.peer_id,
story_id=subject.story_id,
expires_at=expiry_from_seconds(body.expires_in_seconds),
max_downloads=body.max_downloads,
)
@router.get("/lookup")
async def lookup_share(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
kind: Annotated[str, Query()],
media_id: Annotated[int | None, Query()] = None,
version_id: Annotated[int | None, Query()] = None,
peer_id: Annotated[int | None, Query()] = None,
story_id: Annotated[int | None, Query()] = None,
) -> FileShareView | None:
body = ShareCreate(
account_id=account_id,
kind=kind,
media_id=media_id,
version_id=version_id,
peer_id=peer_id,
story_id=story_id,
)
subject = await _resolve(pool, body)
return await shares.find_active_share(pool, account_id, subject.storage_key)
@router.get("/{share_id}/hits")
async def share_hits(
pool: FromDishka[asyncpg.Pool],
share_id: int,
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[FileShareHitView]:
return await shares.list_hits(pool, share_id, Page(limit=limit, offset=offset))
@router.patch("/{share_id}")
async def update_share(
pool: FromDishka[asyncpg.Pool], share_id: int, body: ShareUpdate
) -> FileShareView:
current = await shares.get_share(pool, share_id)
if current is None:
raise HTTPException(status_code=404, detail="share not found")
expires_at = (
current.expires_at
if body.keep_expiry
else expiry_from_seconds(body.expires_in_seconds)
)
updated = await shares.update_share(
pool, share_id, expires_at=expires_at, max_downloads=body.max_downloads
)
if updated is None:
raise HTTPException(status_code=404, detail="share not found")
return updated
@router.post("/{share_id}/revoke")
async def revoke_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> FileShareView:
share = await shares.revoke_share(pool, share_id)
if share is None:
raise HTTPException(status_code=404, detail="share not found")
return share
@router.post("/{share_id}/reissue")
async def reissue_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> FileShareView:
share = await shares.rotate_token(pool, share_id)
if share is None:
raise HTTPException(status_code=404, detail="share not found")
return share
@router.delete("/{share_id}", status_code=204)
async def delete_share(pool: FromDishka[asyncpg.Pool], share_id: int) -> None:
if not await shares.delete_share(pool, share_id):
raise HTTPException(status_code=404, detail="share not found")
+166
View File
@@ -0,0 +1,166 @@
from typing import Annotated
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import FileResponse
from pydantic import BaseModel
from utils import trash
from utils.cache import NO_STORE_HEADERS
from utils.files import resolve_mime
from utils.read import storage as storage_read
from utils.read.models import (
DEFAULT_LIMIT,
Page,
StorageChatView,
StorageFileView,
StorageScopeView,
StorageSummaryView,
TrashItemView,
TrashOutcome,
)
from utils.storage import ContentAddressedStorage
router = APIRouter(prefix="/api/storage", tags=["storage"], route_class=DishkaRoute)
class TrashRequest(BaseModel):
account_id: int
media_ids: list[int]
class RestoreRequest(BaseModel):
account_id: int
keys: list[str] | None = None
class ScopeRequest(BaseModel):
account_id: int
chat_id: int | None = None
folder_id: int | None = None
class PurgeRequest(BaseModel):
account_id: int
keys: list[str] | None = None
@router.get("/summary")
async def summary(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
chat_id: Annotated[int | None, Query()] = None,
) -> StorageSummaryView:
return await storage_read.get_summary(pool, account_id, chat_id)
@router.get("/files")
async def files(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
group: Annotated[str | None, Query()] = None,
chat_id: Annotated[int | None, Query()] = None,
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[StorageFileView]:
page = Page(limit=limit, offset=offset)
return await storage_read.list_files(pool, account_id, group, page, chat_id)
@router.get("/chats")
async def chats(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
group: Annotated[str | None, Query()] = None,
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[StorageChatView]:
page = Page(limit=limit, offset=offset)
return await storage_read.list_chats(pool, account_id, group, page)
async def _chat_ids(
pool: asyncpg.Pool, account_id: int, chat_id: int | None, folder_id: int | None
) -> list[int]:
chat_ids = await storage_read.scope_chat_ids(pool, account_id, chat_id, folder_id)
if chat_ids is None:
raise HTTPException(status_code=404, detail="scope not found")
return chat_ids
@router.get("/scope")
async def scope(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
chat_id: Annotated[int | None, Query()] = None,
folder_id: Annotated[int | None, Query()] = None,
) -> StorageScopeView:
chat_ids = await _chat_ids(pool, account_id, chat_id, folder_id)
return await storage_read.get_scope(pool, account_id, chat_ids)
@router.get("/trash")
async def trash_items(
pool: FromDishka[asyncpg.Pool],
account_id: Annotated[int, Query()],
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
offset: Annotated[int, Query()] = 0,
) -> list[TrashItemView]:
page = Page(limit=limit, offset=offset)
return await storage_read.list_trash(pool, account_id, page)
@router.post("/trash/scope")
async def trash_scope(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
body: ScopeRequest,
) -> TrashOutcome:
chat_ids = await _chat_ids(pool, body.account_id, body.chat_id, body.folder_id)
blobs = await pool.fetch(storage_read.SCOPE_BLOBS, body.account_id, chat_ids)
free = [blob for blob in blobs if blob["free"]]
return await trash.trash_blobs(pool, storage, body.account_id, free)
@router.post("/trash")
async def move_to_trash(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
body: TrashRequest,
) -> TrashOutcome:
return await trash.move_to_trash(pool, storage, body.account_id, body.media_ids)
@router.post("/trash/restore")
async def restore(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
body: RestoreRequest,
) -> TrashOutcome:
return await trash.restore(pool, storage, body.account_id, body.keys)
@router.post("/trash/purge")
async def purge(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
body: PurgeRequest,
) -> TrashOutcome:
return await trash.purge(pool, storage, body.account_id, body.keys)
@router.get("/trash/{storage_key}/file")
async def serve_trashed(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
storage_key: str,
) -> FileResponse:
item = await storage_read.get_trash_item(pool, storage_key)
if item is None or not storage.trashed(storage_key):
raise HTTPException(status_code=404, detail="file is not in trash")
return FileResponse(
storage.trash_url(storage_key),
media_type=resolve_mime(item["kind"], item["mime"], item["file_name"]),
headers=NO_STORE_HEADERS,
)
+8
View File
@@ -5,6 +5,7 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from utils.files import content_disposition, story_file_name
from utils.read import peers from utils.read import peers
from utils.read.models import DEFAULT_LIMIT, Page, StoryView from utils.read.models import DEFAULT_LIMIT, Page, StoryView
from utils.storage import ContentAddressedStorage from utils.storage import ContentAddressedStorage
@@ -36,13 +37,20 @@ async def serve_story_media(
peer_id: int, peer_id: int,
story_id: int, story_id: int,
account_id: AccountId, account_id: AccountId,
download: Annotated[bool, Query()] = False,
) -> FileResponse: ) -> FileResponse:
story = await peers.get_story(pool, account_id, peer_id, story_id) story = await peers.get_story(pool, account_id, peer_id, story_id)
if story is None: if story is None:
raise HTTPException(status_code=404, detail="story not found") raise HTTPException(status_code=404, detail="story not found")
if not story.downloaded or story.storage_key is None: if not story.downloaded or story.storage_key is None:
raise HTTPException(status_code=409, detail="story media not downloaded") raise HTTPException(status_code=409, detail="story media not downloaded")
headers = {}
if download:
headers["Content-Disposition"] = content_disposition(
story_file_name(peer_id, story_id, story.media_kind), attachment=True
)
return FileResponse( return FileResponse(
storage.url(story.storage_key), storage.url(story.storage_key),
media_type=_STORY_MIME.get(story.media_kind or "", "application/octet-stream"), media_type=_STORY_MIME.get(story.media_kind or "", "application/octet-stream"),
headers=headers,
) )
-3
View File
@@ -1,3 +0,0 @@
from userbot.modules.client import DEVICE_MODEL, DEVICE_MODEL_LIMIT, PyroClient
__all__ = ["DEVICE_MODEL", "DEVICE_MODEL_LIMIT", "PyroClient"]
+2 -2
View File
@@ -1,8 +1,8 @@
from pyrogram.types import Message from pyrogram.types import Message
from userbot import PyroClient
from userbot.modules.capture import repository from userbot.modules.capture import repository
from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD
from userbot.modules.client import PyroClient
from utils.events import notify_bg_event from utils.events import notify_bg_event
@@ -30,4 +30,4 @@ async def on_deleted_messages(client: PyroClient, messages: list[Message]) -> No
) )
handlers = on_deleted_messages.handlers handlers = on_deleted_messages.handlers # ty: ignore[unresolved-attribute]
+11 -4
View File
@@ -1,17 +1,24 @@
from pyrogram.types import Message from pyrogram.types import Message
from userbot import PyroClient
from userbot.modules.capture import repository from userbot.modules.capture import repository
from userbot.modules.capture.chat_meta import meta_from_chat from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.capture.message import sender_id from userbot.modules.capture.message import sender_id
from userbot.modules.client import PyroClient
from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
from userbot.modules.raw_json import raw_json
from utils.events import notify_bg_event from utils.events import notify_bg_event
@PyroClient.on_edited_message() @PyroClient.on_edited_message()
async def on_edited_message(client: PyroClient, message: Message) -> None: async def on_edited_message(client: PyroClient, message: Message) -> None:
ctx = client.capture ctx = client.capture
if ctx is None or message.empty or message.chat is None or message.date is None: if (
ctx is None
or message.empty
or message.scheduled
or message.chat is None
or message.date is None
):
return return
chat = message.chat chat = message.chat
chat_id = chat.id or 0 chat_id = chat.id or 0
@@ -27,7 +34,7 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
message.date, message.date,
sender_id(message), sender_id(message),
message.text or message.caption, message.text or message.caption,
str(message), raw_json(message),
message.edit_date, message.edit_date,
media_unique_id(message), media_unique_id(message),
has_media=message.media is not None, has_media=message.media is not None,
@@ -42,4 +49,4 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
) )
handlers = on_edited_message.handlers handlers = on_edited_message.handlers # ty: ignore[unresolved-attribute]
+9 -3
View File
@@ -1,8 +1,8 @@
from pyrogram.types import Message from pyrogram.types import Message
from userbot import PyroClient
from userbot.modules.capture import capture_message from userbot.modules.capture import capture_message
from userbot.modules.capture.chat_meta import meta_from_chat from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.client import PyroClient
from userbot.modules.stt import is_transcribable from userbot.modules.stt import is_transcribable
from userbot.modules.stt.gate import safe_transcribe from userbot.modules.stt.gate import safe_transcribe
from utils.events import notify_bg_event from utils.events import notify_bg_event
@@ -11,7 +11,13 @@ from utils.events import notify_bg_event
@PyroClient.on_message() @PyroClient.on_message()
async def on_message(client: PyroClient, message: Message) -> None: async def on_message(client: PyroClient, message: Message) -> None:
ctx = client.capture ctx = client.capture
if ctx is None or message.empty or message.chat is None or message.date is None: if (
ctx is None
or message.empty
or message.scheduled
or message.chat is None
or message.date is None
):
return return
meta = meta_from_chat(message.chat, ctx.contacts.ids) meta = meta_from_chat(message.chat, ctx.contacts.ids)
await ctx.watches.on_text(meta.chat_id, message.id, message.text or message.caption) await ctx.watches.on_text(meta.chat_id, message.id, message.text or message.caption)
@@ -30,4 +36,4 @@ async def on_message(client: PyroClient, message: Message) -> None:
await safe_transcribe(client, ctx, meta.chat_id, message.id) await safe_transcribe(client, ctx, meta.chat_id, message.id)
handlers = on_message.handlers handlers = on_message.handlers # ty: ignore[unresolved-attribute]
+4 -3
View File
@@ -2,8 +2,9 @@ from datetime import UTC, datetime
from pyrogram.types import User from pyrogram.types import User
from userbot import PyroClient from userbot.modules.client import PyroClient
from userbot.modules.presence import repository from userbot.modules.presence import repository
from userbot.modules.raw_json import raw_json
from utils.events import notify_bg_event from utils.events import notify_bg_event
@@ -20,10 +21,10 @@ async def on_user_status(client: PyroClient, user: User) -> None:
user.status.name.lower(), user.status.name.lower(),
user.last_online_date, user.last_online_date,
user.next_offline_date, user.next_offline_date,
str(user.raw), raw_json(user.raw),
) )
await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online") await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online")
await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id) await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id)
handlers = on_user_status.handlers handlers = on_user_status.handlers # ty: ignore[unresolved-attribute]
+1 -1
View File
@@ -1,6 +1,6 @@
from pyrogram import raw from pyrogram import raw
from userbot import PyroClient from userbot.modules.client import PyroClient
HANDLES = (raw.types.UpdateContactsReset, raw.types.UpdateUser) HANDLES = (raw.types.UpdateContactsReset, raw.types.UpdateUser)
@@ -1,6 +1,6 @@
from pyrogram import raw from pyrogram import raw
from userbot import PyroClient from userbot.modules.client import PyroClient
HANDLES = ( HANDLES = (
raw.types.UpdateDialogFilter, raw.types.UpdateDialogFilter,
@@ -2,7 +2,6 @@ from collections.abc import Awaitable, Callable
from pyrogram import raw from pyrogram import raw
from userbot import PyroClient
from userbot.handlers.raw import ( from userbot.handlers.raw import (
contacts, contacts,
dialog_filters, dialog_filters,
@@ -12,6 +11,7 @@ from userbot.handlers.raw import (
read_receipts, read_receipts,
transcribed, transcribed,
) )
from userbot.modules.client import PyroClient
RawHandler = Callable[..., Awaitable[None]] RawHandler = Callable[..., Awaitable[None]]
@@ -24,8 +24,9 @@ _MODULES = (
read_receipts, read_receipts,
transcribed, transcribed,
) )
_REGISTRY: dict[type, RawHandler] = { _REGISTRY: dict[type, list[RawHandler]] = {
update_type: module.handle for module in _MODULES for update_type in module.HANDLES update_type: [module.handle for module in _MODULES if update_type in module.HANDLES]
for update_type in {kind for module in _MODULES for kind in module.HANDLES}
} }
@@ -33,9 +34,8 @@ _REGISTRY: dict[type, RawHandler] = {
async def dispatch( async def dispatch(
client: PyroClient, update: raw.base.Update, users: dict, chats: dict client: PyroClient, update: raw.base.Update, users: dict, chats: dict
) -> None: ) -> None:
handler = _REGISTRY.get(type(update)) for handler in _REGISTRY.get(type(update), ()):
if handler is not None:
await handler(client, update, users, chats) await handler(client, update, users, chats)
handlers = dispatch.handlers handlers = dispatch.handlers # ty: ignore[unresolved-attribute]
+9 -6
View File
@@ -2,11 +2,12 @@ from dataclasses import replace
from pyrogram import raw from pyrogram import raw
from userbot import PyroClient
from userbot.modules.avatars import capture_avatar from userbot.modules.avatars import capture_avatar
from userbot.modules.client import PyroClient
from userbot.modules.profiles import active_username, snapshot_from_user from userbot.modules.profiles import active_username, snapshot_from_user
from userbot.modules.profiles.parse import ProfileFields from userbot.modules.profiles.parse import ProfileFields
from userbot.modules.profiles.repository import get_peer, write_profile from userbot.modules.profiles.repository import get_peer, write_profile
from userbot.modules.raw_json import raw_json
HANDLES = (raw.types.UpdateUserName, raw.types.UpdateUser, raw.types.UpdateUserPhone) HANDLES = (raw.types.UpdateUserName, raw.types.UpdateUser, raw.types.UpdateUserPhone)
@@ -20,11 +21,13 @@ async def _handle_user(
raw_user = users.get(update.user_id) raw_user = users.get(update.user_id)
if raw_user is None: if raw_user is None:
return return
fields, photo_file_id, photo_unique_id = snapshot_from_user(client, raw_user) fields, photo_file_id, photo_unique_id = await snapshot_from_user(client, raw_user)
current = await get_peer(ctx.pool, ctx.account_id, update.user_id) current = await get_peer(ctx.pool, ctx.account_id, update.user_id)
if current == fields: if current == fields:
return return
await write_profile(ctx.pool, ctx.account_id, update.user_id, fields, str(raw_user)) await write_profile(
ctx.pool, ctx.account_id, update.user_id, fields, raw_json(raw_user)
)
changed_photo = current is None or current.photo_unique_id != photo_unique_id changed_photo = current is None or current.photo_unique_id != photo_unique_id
if photo_file_id and photo_unique_id and changed_photo: if photo_file_id and photo_unique_id and changed_photo:
await capture_avatar( await capture_avatar(
@@ -34,7 +37,7 @@ async def _handle_user(
"peer", "peer",
photo_file_id, photo_file_id,
photo_unique_id, photo_unique_id,
str(raw_user), raw_json(raw_user),
) )
@@ -65,9 +68,9 @@ async def handle(
"last_name": update.last_name or None, "last_name": update.last_name or None,
"username": active_username(update.usernames), "username": active_username(update.usernames),
}, },
str(update), raw_json(update),
) )
elif isinstance(update, raw.types.UpdateUserPhone): elif isinstance(update, raw.types.UpdateUserPhone):
await _merge_partial( await _merge_partial(
client, update.user_id, {"phone": update.phone or None}, str(update) client, update.user_id, {"phone": update.phone or None}, raw_json(update)
) )
@@ -1,8 +1,8 @@
from pyrogram import raw, utils from pyrogram import raw, utils
from userbot import PyroClient
from userbot.modules.capture import repository from userbot.modules.capture import repository
from userbot.modules.capture.chat_meta import meta_from_peer from userbot.modules.capture.chat_meta import meta_from_peer
from userbot.modules.client import PyroClient
from utils.events import notify_bg_event from utils.events import notify_bg_event
HANDLES = (raw.types.UpdateMessageReactions,) HANDLES = (raw.types.UpdateMessageReactions,)
@@ -2,8 +2,9 @@ from datetime import UTC, datetime
from pyrogram import raw, utils from pyrogram import raw, utils
from userbot import PyroClient
from userbot.modules.capture.chat_meta import meta_from_chat_id from userbot.modules.capture.chat_meta import meta_from_chat_id
from userbot.modules.client import PyroClient
from userbot.modules.raw_json import raw_json
from userbot.modules.read_receipts import repository as receipts from userbot.modules.read_receipts import repository as receipts
from userbot.modules.stt import repository from userbot.modules.stt import repository
from userbot.modules.stt.gate import safe_transcribe from userbot.modules.stt.gate import safe_transcribe
@@ -41,7 +42,7 @@ async def handle(
cand_chat_id, cand_chat_id,
datetime.now(UTC), datetime.now(UTC),
message_id, message_id,
str(update), raw_json(update),
) )
continue continue
if not untranscribed: if not untranscribed:
@@ -2,7 +2,8 @@ from datetime import UTC, datetime
from pyrogram import raw, utils from pyrogram import raw, utils
from userbot import PyroClient from userbot.modules.client import PyroClient
from userbot.modules.raw_json import raw_json
from userbot.modules.read_receipts import repository from userbot.modules.read_receipts import repository
from utils.events import notify_bg_event from utils.events import notify_bg_event
@@ -23,7 +24,7 @@ async def handle(
chat_id, chat_id,
datetime.now(UTC), datetime.now(UTC),
update.max_id, update.max_id,
str(update), raw_json(update),
) )
await notify_bg_event( await notify_bg_event(
ctx.pool, "receipt", ctx.account_id, chat_id=chat_id, message_id=update.max_id ctx.pool, "receipt", ctx.account_id, chat_id=chat_id, message_id=update.max_id
@@ -1,7 +1,7 @@
from pyrogram import raw from pyrogram import raw
from userbot import PyroClient
from userbot.modules.capture.chat_meta import meta_from_peer from userbot.modules.capture.chat_meta import meta_from_peer
from userbot.modules.client import PyroClient
from userbot.modules.stt import repository from userbot.modules.stt import repository
HANDLES = (raw.types.UpdateTranscribedAudio,) HANDLES = (raw.types.UpdateTranscribedAudio,)
+5 -43
View File
@@ -1,52 +1,14 @@
from io import BytesIO
from pyrogram.types import Story from pyrogram.types import Story
from userbot import PyroClient from userbot.modules.client import PyroClient
from userbot.modules.stories import repository from userbot.modules.stories.service import save_story
def _peer_id(story: Story) -> int:
if story.chat is not None:
return story.chat.id or 0
if story.from_user is not None:
return story.from_user.id or 0
return 0
@PyroClient.on_story() @PyroClient.on_story()
async def on_story(client: PyroClient, story: Story) -> None: async def on_story(client: PyroClient, story: Story) -> None:
ctx = client.capture if client.capture is None:
if ctx is None:
return return
media_kind = story.media.name.lower() if story.media else None await save_story(client, client.capture, story)
storage_key: str | None = None
file_size: int | None = None
downloaded = False
if not story.deleted and story.media is not None:
buffer = await client.download_media(story, in_memory=True)
if isinstance(buffer, BytesIO):
data = buffer.getvalue()
storage_key = ctx.storage.put(data)
file_size = len(data)
downloaded = True
await repository.upsert_story(
ctx.pool,
ctx.account_id,
_peer_id(story),
story.id,
story.date,
story.expire_date,
story.caption,
media_kind,
storage_key,
file_size,
story.views,
str(story.raw),
pinned=bool(story.pinned),
deleted=bool(story.deleted),
downloaded=downloaded,
)
handlers = on_story.handlers handlers = on_story.handlers # ty: ignore[unresolved-attribute]
@@ -1,9 +1,9 @@
from io import BytesIO
from pyrogram import Client from pyrogram import Client
from userbot.modules.avatars import repository from userbot.modules.avatars import repository
from userbot.modules.capture.context import CaptureContext from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import download_bytes
from utils.logging import logger
async def capture_avatar( # noqa: PLR0913 async def capture_avatar( # noqa: PLR0913
@@ -20,9 +20,12 @@ async def capture_avatar( # noqa: PLR0913
storage_key: str | None = None storage_key: str | None = None
file_size: int | None = None file_size: int | None = None
downloaded = False downloaded = False
buffer = await client.download_media(file_id, in_memory=True) try:
if isinstance(buffer, BytesIO): data = await download_bytes(client, file_id)
data = buffer.getvalue() except TimeoutError:
logger.warning(f"[yellow]Avatar download timed out for {owner_id}.[/]")
data = None
if data is not None:
storage_key = ctx.storage.put(data) storage_key = ctx.storage.put(data)
file_size = len(data) file_size = len(data)
downloaded = True downloaded = True
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
from userbot.modules.capture.chat_meta import chat_kind from userbot.modules.capture.chat_meta import chat_kind
from userbot.modules.profiles.parse import ProfileFields, snapshot_from_high_level from userbot.modules.profiles.parse import ProfileFields, snapshot_from_high_level
from userbot.modules.profiles.repository import get_peer, write_profile from userbot.modules.profiles.repository import get_peer, write_profile
from userbot.modules.raw_json import raw_json
from utils.policy.models import ChatKind from utils.policy.models import ChatKind
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -40,9 +41,7 @@ class ChatMetaCache:
chat_id: int, chat_id: int,
meta: tuple[str | None, str | None], meta: tuple[str | None, str | None],
) -> bool: ) -> bool:
from userbot.modules.groups.repository import ( # noqa: PLC0415 from userbot.modules.groups.repository import get_latest_chat_meta # noqa: PLC0415
get_latest_chat_meta,
)
if chat_id not in self._cache: if chat_id not in self._cache:
self._cache[chat_id] = await get_latest_chat_meta(pool, account_id, chat_id) self._cache[chat_id] = await get_latest_chat_meta(pool, account_id, chat_id)
@@ -59,7 +58,7 @@ async def _capture_peer(message: Message, ctx: CaptureContext) -> None:
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user) fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
if not await ctx.peer_identity.changed(ctx.pool, ctx.account_id, user.id, fields): if not await ctx.peer_identity.changed(ctx.pool, ctx.account_id, user.id, fields):
return return
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user)) await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
from userbot.modules.avatars import note_avatar # noqa: PLC0415 from userbot.modules.avatars import note_avatar # noqa: PLC0415
@@ -95,7 +94,7 @@ async def _capture_chat(message: Message, ctx: CaptureContext) -> None:
photo_unique_id, photo_unique_id,
None, None,
message.date, message.date,
str(message), raw_json(message),
) )
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
from userbot.modules.avatars import note_avatar # noqa: PLC0415 from userbot.modules.avatars import note_avatar # noqa: PLC0415
@@ -6,6 +6,7 @@ from userbot.modules.capture.context import CaptureContext
from userbot.modules.links import extract_links from userbot.modules.links import extract_links
from userbot.modules.links import repository as links_repository from userbot.modules.links import repository as links_repository
from userbot.modules.media import capture_media, self_destruct_ttl from userbot.modules.media import capture_media, self_destruct_ttl
from userbot.modules.raw_json import raw_json
from utils.policy.models import CaptureToggles from utils.policy.models import CaptureToggles
@@ -47,7 +48,7 @@ async def capture_message(
message.date, message.date,
sender_id(message), sender_id(message),
message.text or message.caption, message.text or message.caption,
str(message), raw_json(message),
has_media=message.media is not None, has_media=message.media is not None,
is_self_destruct=self_destruct_ttl(message) is not None, is_self_destruct=self_destruct_ttl(message) is not None,
) )
@@ -68,5 +69,5 @@ async def capture_message(
links = extract_links(message) links = extract_links(message)
if links: if links:
await links_repository.insert_links( await links_repository.insert_links(
ctx.pool, ctx.account_id, chat_id, message.id, links, str(message) ctx.pool, ctx.account_id, chat_id, message.id, links, raw_json(message)
) )
@@ -30,6 +30,16 @@ ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
edited_at = now() edited_at = now()
""" """
_FIND_DELETED = """
SELECT chat_id, message_id, date FROM messages
WHERE account_id = $1 AND message_id = ANY($2::bigint[]) AND deleted_at IS NULL
"""
_MARK_DELETED = """
UPDATE messages SET deleted_at = now()
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
"""
_INSERT_VERSION = """ _INSERT_VERSION = """
INSERT INTO message_versions INSERT INTO message_versions
(account_id, chat_id, message_id, observed_at, edit_date, text, raw) (account_id, chat_id, message_id, observed_at, edit_date, text, raw)
@@ -78,8 +88,8 @@ ON CONFLICT (account_id, chat_id, message_id, storage_key) DO NOTHING
_INSERT_MEDIA = """ _INSERT_MEDIA = """
INSERT INTO media INSERT INTO media
(account_id, chat_id, message_id, kind, storage_key, file_size, mime, (account_id, chat_id, message_id, kind, storage_key, file_size, mime,
ttl_seconds, downloaded, unique_id) ttl_seconds, downloaded, unique_id, evicted_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, CASE WHEN $11 THEN now() END)
ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
kind = EXCLUDED.kind, kind = EXCLUDED.kind,
storage_key = EXCLUDED.storage_key, storage_key = EXCLUDED.storage_key,
@@ -87,7 +97,27 @@ ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
mime = EXCLUDED.mime, mime = EXCLUDED.mime,
ttl_seconds = EXCLUDED.ttl_seconds, ttl_seconds = EXCLUDED.ttl_seconds,
downloaded = EXCLUDED.downloaded, downloaded = EXCLUDED.downloaded,
unique_id = EXCLUDED.unique_id unique_id = EXCLUDED.unique_id,
evicted_at = CASE
WHEN EXCLUDED.downloaded THEN NULL
WHEN media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
THEN EXCLUDED.evicted_at
ELSE coalesce(media.evicted_at, EXCLUDED.evicted_at) END,
evicted_key = CASE WHEN EXCLUDED.downloaded
OR media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
THEN NULL ELSE media.evicted_key END
"""
_IS_EVICTED = """
SELECT EXISTS (
SELECT 1 FROM media WHERE unique_id = $1 AND evicted_at IS NOT NULL
)
"""
_UNEVICT = """
UPDATE media SET storage_key = $2, file_size = $3, downloaded = true,
evicted_at = NULL, evicted_key = NULL
WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL)
""" """
@@ -118,30 +148,44 @@ async def upsert_message( # noqa: PLR0913
) )
async def max_message_id(
pool: asyncpg.Pool, account_id: int, chat_id: int
) -> int | None:
return await pool.fetchval(
"SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
account_id,
chat_id,
)
async def _mark_deleted(
pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> None:
for row in rows:
await pool.execute(
_MARK_DELETED, account_id, row["chat_id"], row["message_id"], row["date"]
)
async def mark_deleted_box( async def mark_deleted_box(
pool: asyncpg.Pool, account_id: int, message_ids: list[int] pool: asyncpg.Pool, account_id: int, message_ids: list[int]
) -> None: ) -> None:
await pool.execute( rows = await pool.fetch(
"UPDATE messages SET deleted_at = now() " f"{_FIND_DELETED} AND chat_id > $3",
"WHERE account_id = $1 AND message_id = ANY($2::bigint[]) "
"AND chat_id > $3 AND deleted_at IS NULL",
account_id, account_id,
message_ids, message_ids,
CHANNEL_ID_THRESHOLD, CHANNEL_ID_THRESHOLD,
) )
await _mark_deleted(pool, account_id, rows)
async def mark_deleted_channel( async def mark_deleted_channel(
pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int] pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int]
) -> None: ) -> None:
await pool.execute( rows = await pool.fetch(
"UPDATE messages SET deleted_at = now() " f"{_FIND_DELETED} AND chat_id = $3", account_id, message_ids, chat_id
"WHERE account_id = $1 AND chat_id = $2 AND message_id = ANY($3::bigint[]) "
"AND deleted_at IS NULL",
account_id,
chat_id,
message_ids,
) )
await _mark_deleted(pool, account_id, rows)
async def add_version( # noqa: PLR0913 async def add_version( # noqa: PLR0913
@@ -205,6 +249,7 @@ async def insert_media( # noqa: PLR0913
unique_id: str | None, unique_id: str | None,
*, *,
downloaded: bool, downloaded: bool,
evicted: bool = False,
) -> None: ) -> None:
async with pool.acquire() as conn, conn.transaction(): async with pool.acquire() as conn, conn.transaction():
await conn.execute( await conn.execute(
@@ -219,6 +264,7 @@ async def insert_media( # noqa: PLR0913
ttl_seconds, ttl_seconds,
downloaded, downloaded,
unique_id, unique_id,
evicted,
) )
if storage_key is not None: if storage_key is not None:
await conn.execute( await conn.execute(
@@ -235,6 +281,22 @@ async def insert_media( # noqa: PLR0913
) )
async def is_evicted(pool: asyncpg.Pool, unique_id: str | None) -> bool:
return unique_id is not None and await pool.fetchval(_IS_EVICTED, unique_id)
async def unevict(
pool: asyncpg.Pool, unique_id: str | None, storage_key: str, file_size: int | None
) -> bool:
async with pool.acquire() as conn, conn.transaction():
await conn.execute(_UNEVICT, unique_id, storage_key, file_size)
dropped = await conn.fetchval(
"DELETE FROM media_trash WHERE storage_key = $1 RETURNING storage_key",
storage_key,
)
return dropped is not None
async def insert_callbacks( async def insert_callbacks(
pool: asyncpg.Pool, pool: asyncpg.Pool,
account_id: int, account_id: int,
+10 -3
View File
@@ -2,6 +2,8 @@ from typing import TYPE_CHECKING
from pyrogram import Client, enums from pyrogram import Client, enums
from utils.telegram import API_HASH, API_ID
if TYPE_CHECKING: if TYPE_CHECKING:
from userbot.modules.capture import CaptureContext from userbot.modules.capture import CaptureContext
@@ -11,19 +13,23 @@ DEVICE_MODEL_LIMIT = 32
class PyroClient(Client): class PyroClient(Client):
def __init__( def __init__( # noqa: PLR0913
self, self,
name: str, name: str,
*, *,
workdir: str = "sessions", workdir: str = "sessions",
device_model: str | None = None, device_model: str | None = None,
load_handlers: bool = True, load_handlers: bool = True,
bot_token: str | None = None,
skip_updates: bool = True,
) -> None: ) -> None:
super().__init__( super().__init__(
name, name,
workdir=workdir, workdir=workdir,
api_id=2040, api_id=API_ID,
api_hash="b18441a1ff607e10a989891a5462e627", api_hash=API_HASH,
bot_token=bot_token,
skip_updates=skip_updates,
device_model=device_model or DEVICE_MODEL, device_model=device_model or DEVICE_MODEL,
system_version="Windows 11 x64", system_version="Windows 11 x64",
app_version="7.0.8 x64", app_version="7.0.8 x64",
@@ -31,6 +37,7 @@ class PyroClient(Client):
lang_code="en", lang_code="en",
system_lang_code="en-US", system_lang_code="en-US",
client_platform=enums.ClientPlatform.DESKTOP, client_platform=enums.ClientPlatform.DESKTOP,
max_concurrent_transmissions=4,
) )
self.capture: CaptureContext | None = None self.capture: CaptureContext | None = None
+56
View File
@@ -0,0 +1,56 @@
import asyncio
import contextlib
from io import BytesIO
from pyrogram import Client
from pyrogram.types import (
Animation,
Audio,
Document,
Message,
Photo,
Sticker,
Story,
Video,
VideoNote,
Voice,
)
from utils.logging import logger
type Downloadable = (
str
| Message
| Story
| Audio
| Document
| Photo
| Sticker
| Animation
| Video
| Voice
| VideoNote
)
STOP_TIMEOUT = 5.0
_reset_lock = asyncio.Lock()
async def _reset_media_sessions(client: Client) -> None:
async with _reset_lock:
sessions = list(client.media_sessions.values())
client.media_sessions.clear()
for session in sessions:
with contextlib.suppress(Exception):
await asyncio.wait_for(session.stop(), STOP_TIMEOUT)
logger.warning(f"[yellow]Dropped {len(sessions)} stale media session(s).[/]")
async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
try:
buffer = await client.download_media(target, in_memory=True)
except (TimeoutError, OSError):
await _reset_media_sessions(client)
buffer = await client.download_media(target, in_memory=True)
return buffer.getvalue() if isinstance(buffer, BytesIO) else None
+13 -6
View File
@@ -5,6 +5,7 @@ from userbot.modules.avatars import capture_avatar
from userbot.modules.capture.context import CaptureContext from userbot.modules.capture.context import CaptureContext
from userbot.modules.capture.message import sender_id from userbot.modules.capture.message import sender_id
from userbot.modules.groups import repository from userbot.modules.groups import repository
from userbot.modules.raw_json import raw_json
async def capture_service( async def capture_service(
@@ -28,7 +29,7 @@ async def capture_service(
None, None,
actor, actor,
ts, ts,
str(message), raw_json(message),
) )
elif service is enums.MessageServiceType.NEW_CHAT_PHOTO: elif service is enums.MessageServiceType.NEW_CHAT_PHOTO:
photo = message.new_chat_photo photo = message.new_chat_photo
@@ -43,11 +44,17 @@ async def capture_service(
unique_id, unique_id,
actor, actor,
ts, ts,
str(message), raw_json(message),
) )
if photo is not None and unique_id is not None: if photo is not None and unique_id is not None:
await capture_avatar( await capture_avatar(
client, ctx, chat_id, "chat", photo.file_id, unique_id, str(message) client,
ctx,
chat_id,
"chat",
photo.file_id,
unique_id,
raw_json(message),
) )
elif service is enums.MessageServiceType.DELETE_CHAT_PHOTO: elif service is enums.MessageServiceType.DELETE_CHAT_PHOTO:
await repository.insert_chat_history( await repository.insert_chat_history(
@@ -60,7 +67,7 @@ async def capture_service(
None, None,
actor, actor,
ts, ts,
str(message), raw_json(message),
) )
elif service is enums.MessageServiceType.NEW_CHAT_MEMBERS: elif service is enums.MessageServiceType.NEW_CHAT_MEMBERS:
for member in message.new_chat_members or []: for member in message.new_chat_members or []:
@@ -73,7 +80,7 @@ async def capture_service(
"join", "join",
actor, actor,
ts, ts,
str(message), raw_json(message),
) )
elif service is enums.MessageServiceType.LEFT_CHAT_MEMBER: elif service is enums.MessageServiceType.LEFT_CHAT_MEMBER:
member = message.left_chat_member member = message.left_chat_member
@@ -87,5 +94,5 @@ async def capture_service(
"leave", "leave",
actor, actor,
ts, ts,
str(message), raw_json(message),
) )
@@ -69,6 +69,9 @@ class JobConsumer:
requeued = await repository.requeue_running(self.pool, self.account_id) requeued = await repository.requeue_running(self.pool, self.account_id)
if requeued: if requeued:
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]") logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
pruned = await repository.prune_finished(self.pool, self.account_id)
if pruned:
logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
conn = await asyncpg.connect(dsn=env.db.connection_url) conn = await asyncpg.connect(dsn=env.db.connection_url)
await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set()) await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]") logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
@@ -1,5 +1,6 @@
from userbot.modules.jobs.handlers import ( from userbot.modules.jobs.handlers import (
backfill, backfill,
backfill_stories,
enrich_chat, enrich_chat,
fetch_avatar, fetch_avatar,
fetch_custom_emoji, fetch_custom_emoji,
@@ -12,6 +13,7 @@ from userbot.modules.jobs.handlers import (
__all__ = [ __all__ = [
"backfill", "backfill",
"backfill_stories",
"enrich_chat", "enrich_chat",
"fetch_avatar", "fetch_avatar",
"fetch_custom_emoji", "fetch_custom_emoji",
@@ -1,17 +1,49 @@
from pyrogram.errors import PeerIdInvalid from pyrogram import Client
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
from pyrogram.types import Message
from userbot.modules.capture import capture_message from userbot.modules.capture import capture_message
from userbot.modules.capture import repository as capture_repo
from userbot.modules.capture.chat_meta import meta_from_chat from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.capture.context import CaptureContext
from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register from userbot.modules.jobs.registry import register
from userbot.modules.stt import repository as stt_repo from userbot.modules.stt import repository as stt_repo
from userbot.modules.stt import should_transcribe_on_backfill from userbot.modules.stt import should_transcribe_on_backfill
from userbot.modules.stt.gate import safe_transcribe from userbot.modules.stt.gate import safe_transcribe
from utils.logging import logger
from utils.policy.models import CaptureToggles from utils.policy.models import CaptureToggles
SAVE_EVERY = 100 SAVE_EVERY = 100
async def resolve_min_id(ctx: JobContext, chat_id: int) -> int:
cursor = ctx.job.cursor or {}
if "min_id" in cursor:
return int(cursor["min_id"])
if ctx.job.params.get("full"):
return 0
newest = await capture_repo.max_message_id(ctx.pool, ctx.account_id, chat_id)
return newest + 1 if newest else 0
async def maybe_transcribe(
client: Client,
capture: CaptureContext,
chat_id: int,
message: Message,
self_id: int | None,
) -> None:
if not (should_transcribe_on_backfill(message, self_id) and message.chat):
return
meta = meta_from_chat(message.chat, capture.contacts.ids)
already = await stt_repo.is_transcribed(
capture.pool, capture.account_id, chat_id, message.id
)
if capture.resolve(meta).stt and not already:
await safe_transcribe(client, capture, chat_id, message.id)
@register("backfill") @register("backfill")
async def backfill(ctx: JobContext) -> None: async def backfill(ctx: JobContext) -> None:
client = ctx.client client = ctx.client
@@ -26,28 +58,34 @@ async def backfill(ctx: JobContext) -> None:
media=bool(ctx.job.params.get("media")), media=bool(ctx.job.params.get("media")),
self_destruct_media=False, self_destruct_media=False,
) )
max_id = (ctx.job.cursor or {}).get("max_id", 0) max_id = int((ctx.job.cursor or {}).get("max_id", 0))
min_id = await resolve_min_id(ctx, chat_id)
await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
processed = ctx.job.progress.get("processed", 0) processed = ctx.job.progress.get("processed", 0)
kwargs = {"max_id": max_id} if max_id else {} errors = ctx.job.progress.get("errors", 0)
self_id = client.me.id if client.me else None self_id = client.me.id if client.me else None
try: try:
async for message in client.get_chat_history(chat_id, **kwargs): async for message in client.get_chat_history(
chat_id, max_id=max_id, min_id=min_id
):
try:
await capture_message(client, message, capture, toggles) await capture_message(client, message, capture, toggles)
if should_transcribe_on_backfill(message, self_id) and message.chat: await maybe_transcribe(client, capture, chat_id, message, self_id)
meta = meta_from_chat(message.chat, capture.contacts.ids) except (FloodWait, FloodPremiumWait, PeerIdInvalid):
already = await stt_repo.is_transcribed( raise
capture.pool, capture.account_id, chat_id, message.id except Exception:
) errors += 1
if capture.resolve(meta).stt and not already: logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
await safe_transcribe(client, capture, chat_id, message.id)
processed += 1 processed += 1
if processed % SAVE_EVERY == 0: if processed % SAVE_EVERY == 0:
next_max = message.id - 1 next_max = message.id - 1
await ctx.save_cursor({"max_id": next_max}) await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
await ctx.report_progress({"processed": processed, "max_id": next_max}) await ctx.report_progress(
{"processed": processed, "errors": errors, "max_id": next_max}
)
if await ctx.is_canceled(): if await ctx.is_canceled():
return return
except PeerIdInvalid: except PeerIdInvalid:
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"}) await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
return return
await ctx.report_progress({"processed": processed, "done": True}) await ctx.report_progress({"processed": processed, "errors": errors, "done": True})
@@ -0,0 +1,74 @@
from collections.abc import AsyncIterator, Callable
from pyrogram import Client
from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
from pyrogram.types import Story
from userbot.modules.capture.context import CaptureContext
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
from userbot.modules.stories.service import save_story
from utils.logging import logger
SAVE_EVERY = 10
StorySource = Callable[[], AsyncIterator[Story]]
def _sources(client: Client, peer_id: int, *, own: bool) -> dict[str, StorySource]:
sources: dict[str, StorySource] = {
"active": lambda: client.get_chat_stories(peer_id),
"pinned": lambda: client.get_pinned_stories(peer_id),
}
if own:
sources["archived"] = lambda: client.get_archived_stories(peer_id)
return sources
async def _drain(
ctx: JobContext, capture: CaptureContext, name: str, source: StorySource
) -> int:
client = ctx.client
if client is None:
return 0
saved = 0
async for story in source():
try:
await save_story(client, capture, story)
except (FloodWait, FloodPremiumWait):
raise
except RPCError:
continue
except Exception:
logger.exception(f"Failed to save story {story.id} from {name}")
continue
saved += 1
if saved % SAVE_EVERY == 0:
await ctx.report_progress({"saved": saved, "source": name})
if await ctx.is_canceled():
break
return saved
@register("backfill_stories")
async def backfill_stories(ctx: JobContext) -> None:
client = ctx.client
if client is None:
return
capture = getattr(client, "capture", None)
if capture is None:
return
peer_id = ctx.job.params["peer_id"]
own = client.me is not None and client.me.id == peer_id
saved = 0
errors: dict[str, str] = {}
for name, source in _sources(client, peer_id, own=own).items():
try:
saved += await _drain(ctx, capture, name, source)
except (FloodWait, FloodPremiumWait):
raise
except RPCError as exc:
errors[name] = type(exc).__name__
if await ctx.is_canceled():
break
await ctx.report_progress({"saved": saved, "done": True, "errors": errors})
@@ -11,6 +11,7 @@ from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register from userbot.modules.jobs.registry import register
from userbot.modules.profiles.parse import snapshot_from_high_level from userbot.modules.profiles.parse import snapshot_from_high_level
from userbot.modules.profiles.repository import write_profile from userbot.modules.profiles.repository import write_profile
from userbot.modules.raw_json import raw_json
MEMBER_CAP = 200 MEMBER_CAP = 200
@@ -26,7 +27,7 @@ LIMIT 100
async def _save_user(ctx: CaptureContext, user: User) -> None: async def _save_user(ctx: CaptureContext, user: User) -> None:
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user) fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user)) await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
ctx.pool, ctx.account_id, user.id, "peer", photo_unique_id, photo_file_id ctx.pool, ctx.account_id, user.id, "peer", photo_unique_id, photo_file_id
@@ -35,6 +36,8 @@ async def _save_user(ctx: CaptureContext, user: User) -> None:
async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -> None: async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -> None:
chat = await client.get_chat(chat_id) chat = await client.get_chat(chat_id)
if chat is None:
return
photo = chat.photo photo = chat.photo
photo_unique_id = photo.big_photo_unique_id if photo else None photo_unique_id = photo.big_photo_unique_id if photo else None
photo_file_id = photo.big_file_id if photo else None photo_file_id = photo.big_file_id if photo else None
@@ -48,7 +51,7 @@ async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -
photo_unique_id, photo_unique_id,
None, None,
datetime.now(UTC), datetime.now(UTC),
str(chat), raw_json(chat),
) )
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
@@ -65,15 +68,27 @@ async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) ->
return return
async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
try:
users = await client.get_users(ids)
except BadRequest:
users = []
for sender_id in ids:
try:
users.append(await client.get_users(sender_id))
except BadRequest:
continue
if isinstance(users, User):
return [users]
return [user for user in users if isinstance(user, User)]
async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None: async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id) rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
ids = [row["sender_id"] for row in rows] ids = [row["sender_id"] for row in rows]
for sender_id in ids: if not ids:
try: return
user = await client.get_users(sender_id) for user in await _fetch_users(client, ids):
except BadRequest:
continue
if isinstance(user, User):
await _save_user(ctx, user) await _save_user(ctx, user)
@@ -1,9 +1,48 @@
from io import BytesIO from pyrogram import Client
from pyrogram.errors import (
ChannelInvalid,
ChannelPrivate,
FileIdInvalid,
FileReferenceExpired,
FileReferenceInvalid,
PeerIdInvalid,
)
from pyrogram.types import ChatPhoto
from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded
from userbot.modules.download import download_bytes
from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register from userbot.modules.jobs.registry import register
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
UNKNOWN_OWNER = (PeerIdInvalid, ChannelInvalid, ChannelPrivate)
async def _fresh_file_id(client: Client, owner_id: int, unique_id: str) -> str | None:
try:
chat = await client.get_chat(owner_id)
except UNKNOWN_OWNER:
return None
photo = chat.photo if chat else None
if photo and photo.big_photo_unique_id == unique_id:
return photo.big_file_id
async for item in client.get_chat_photos(owner_id):
if isinstance(item, ChatPhoto) and item.big_photo_unique_id == unique_id:
return item.big_file_id
return None
async def _download(
client: Client, owner_id: int, unique_id: str, file_id: str
) -> bytes | None:
try:
return await download_bytes(client, file_id)
except STALE_FILE_ID:
fresh = await _fresh_file_id(client, owner_id, unique_id)
if fresh is None:
return None
return await download_bytes(client, fresh)
@register("fetch_avatar") @register("fetch_avatar")
async def fetch_avatar(ctx: JobContext) -> None: async def fetch_avatar(ctx: JobContext) -> None:
@@ -24,10 +63,9 @@ async def fetch_avatar(ctx: JobContext) -> None:
file_id, downloaded = found file_id, downloaded = found
if downloaded or file_id is None: if downloaded or file_id is None:
return return
buffer = await client.download_media(file_id, in_memory=True) data = await _download(client, owner_id, unique_id, file_id)
if not isinstance(buffer, BytesIO): if data is None:
return return
data = buffer.getvalue()
storage_key = capture.storage.put(data) storage_key = capture.storage.put(data)
await mark_avatar_downloaded( await mark_avatar_downloaded(
ctx.pool, ctx.pool,
@@ -1,8 +1,7 @@
from io import BytesIO
from pyrogram.types import Sticker from pyrogram.types import Sticker
from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded
from userbot.modules.download import download_bytes
from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register from userbot.modules.jobs.registry import register
@@ -30,10 +29,9 @@ async def fetch_custom_emoji(ctx: JobContext) -> None:
if not stickers: if not stickers:
return return
sticker = stickers[0] sticker = stickers[0]
buffer = await client.download_media(sticker.file_id, in_memory=True) data = await download_bytes(client, sticker.file_id)
if not isinstance(buffer, BytesIO): if data is None:
return return
data = buffer.getvalue()
storage_key = capture.storage.put(data) storage_key = capture.storage.put(data)
await upsert_downloaded( await upsert_downloaded(
ctx.pool, ctx.pool,
@@ -18,6 +18,12 @@ async def fetch_media(ctx: JobContext) -> None:
if isinstance(message, list): if isinstance(message, list):
message = message[0] if message else None message = message[0] if message else None
if message is None or message.empty: if message is None or message.empty:
return msg = f"message {chat_id}/{message_id} is unavailable"
raise RuntimeError(msg)
toggles = CaptureToggles(media=True, self_destruct_media=True) toggles = CaptureToggles(media=True, self_destruct_media=True)
await capture_media(client, message, capture, chat_id, message_id, toggles) ok = await capture_media(
client, message, capture, chat_id, message_id, toggles, force=True
)
if not ok:
msg = f"media download failed for {chat_id}/{message_id}"
raise RuntimeError(msg)
@@ -46,7 +46,7 @@ async def _save_found(
source = _source(peer, users, chats) source = _source(peer, users, chats)
if source is None: if source is None:
return None return None
chat = Chat._parse_chat(client, source) # noqa: SLF001 chat = await Chat._parse_chat(client, source) # noqa: SLF001
if chat is None or chat.id is None: if chat is None or chat.id is None:
return None return None
await save_chat(ctx, chat) await save_chat(ctx, chat)
@@ -5,6 +5,7 @@ from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register from userbot.modules.jobs.registry import register
from userbot.modules.profiles.parse import snapshot_from_high_level from userbot.modules.profiles.parse import snapshot_from_high_level
from userbot.modules.profiles.repository import write_profile from userbot.modules.profiles.repository import write_profile
from userbot.modules.raw_json import raw_json
@register("sync_contacts") @register("sync_contacts")
@@ -21,7 +22,7 @@ async def sync_contacts(ctx: JobContext) -> None:
if not isinstance(user, User): if not isinstance(user, User):
continue continue
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user) fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user)) await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
ctx.pool, ctx.pool,
@@ -9,6 +9,8 @@ from userbot.modules.jobs.registry import register
from userbot.modules.profiles.parse import snapshot_from_high_level from userbot.modules.profiles.parse import snapshot_from_high_level
from userbot.modules.profiles.repository import write_profile from userbot.modules.profiles.repository import write_profile
from userbot.modules.profiles.snapshots import save_group, save_private from userbot.modules.profiles.snapshots import save_group, save_private
from userbot.modules.raw_json import raw_json
from utils.logging import logger
SAVE_EVERY = 100 SAVE_EVERY = 100
USERS_BATCH = 200 USERS_BATCH = 200
@@ -31,7 +33,9 @@ async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) ->
if not isinstance(user, User): if not isinstance(user, User):
continue continue
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user) fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user)) await write_profile(
ctx.pool, ctx.account_id, user.id, fields, raw_json(user)
)
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
ctx.pool, ctx.pool,
@@ -66,6 +70,8 @@ async def sync_dialogs(ctx: JobContext) -> None:
await save_group(capture, chat) await save_group(capture, chat)
except (BadRequest, Forbidden): except (BadRequest, Forbidden):
pass pass
except Exception:
logger.exception(f"Failed to snapshot dialog {chat_id}")
await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id) await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
processed += 1 processed += 1
if processed % SAVE_EVERY == 0: if processed % SAVE_EVERY == 0:
@@ -48,6 +48,16 @@ async def requeue_running(pool: asyncpg.Pool, account_id: int) -> int:
return int(result.split()[-1]) return int(result.split()[-1])
async def prune_finished(pool: asyncpg.Pool, account_id: int) -> int:
result = await pool.execute(
"DELETE FROM jobs WHERE account_id = $1 "
"AND status IN ('done', 'canceled', 'failed') "
"AND finished_at < now() - interval '30 days'",
account_id,
)
return int(result.split()[-1])
async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None: async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None:
await pool.execute( await pool.execute(
"UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1", "UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1",
@@ -1,13 +1,23 @@
from io import BytesIO
from typing import Any from typing import Any
from pyrogram import Client from pyrogram import Client
from pyrogram.errors import (
FileIdInvalid,
FileReferenceExpired,
FileReferenceInvalid,
FloodPremiumWait,
FloodWait,
)
from pyrogram.types import Message from pyrogram.types import Message
from userbot.modules.capture import repository from userbot.modules.capture import repository
from userbot.modules.capture.context import CaptureContext from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import Downloadable, download_bytes
from utils.logging import logger
from utils.policy.models import CaptureToggles from utils.policy.models import CaptureToggles
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
_MEDIA_ATTRS = ( _MEDIA_ATTRS = (
"photo", "photo",
"video", "video",
@@ -47,6 +57,48 @@ def media_unique_id(message: Message) -> str | None:
return getattr(obj, "file_unique_id", None) if obj is not None else None return getattr(obj, "file_unique_id", None) if obj is not None else None
def _download_target(message: Message) -> Downloadable | None:
kind, obj = media_object(message)
if obj is None:
return None
return message if getattr(message, kind or "", None) is obj else obj
async def _fresh_target(
client: Client, chat_id: int, message_id: int
) -> Downloadable | None:
message = await client.get_messages(chat_id, message_id)
if not isinstance(message, Message) or message.empty:
return None
return _download_target(message)
async def _download(
client: Client, target: Downloadable | None, chat_id: int, message_id: int
) -> bytes | None:
if target is None:
return None
try:
return await download_bytes(client, target)
except TimeoutError:
logger.warning(
f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
)
return None
except STALE_FILE_ID:
fresh = await _fresh_target(client, chat_id, message_id)
if fresh is None:
return None
try:
return await download_bytes(client, fresh)
except (TimeoutError, *STALE_FILE_ID):
logger.warning(
f"[yellow]Media download failed after refetch "
f"for {chat_id}/{message_id}.[/]"
)
return None
async def capture_media( # noqa: PLR0913 async def capture_media( # noqa: PLR0913
client: Client, client: Client,
message: Message, message: Message,
@@ -54,10 +106,12 @@ async def capture_media( # noqa: PLR0913
chat_id: int, chat_id: int,
message_id: int, message_id: int,
toggles: CaptureToggles, toggles: CaptureToggles,
) -> None: *,
force: bool = False,
) -> bool:
kind, obj = media_object(message) kind, obj = media_object(message)
if obj is None: if obj is None:
return return True
unique_id = getattr(obj, "file_unique_id", None) unique_id = getattr(obj, "file_unique_id", None)
ttl = getattr(obj, "ttl_seconds", None) ttl = getattr(obj, "ttl_seconds", None)
want = toggles.self_destruct_media if ttl else toggles.media want = toggles.self_destruct_media if ttl else toggles.media
@@ -65,6 +119,8 @@ async def capture_media( # noqa: PLR0913
mime = getattr(obj, "mime_type", None) mime = getattr(obj, "mime_type", None)
storage_key: str | None = None storage_key: str | None = None
downloaded = False downloaded = False
evicted = False
flood: FloodWait | FloodPremiumWait | None = None
if want: if want:
existing = await repository.current_media( existing = await repository.current_media(
ctx.pool, ctx.account_id, chat_id, message_id ctx.pool, ctx.account_id, chat_id, message_id
@@ -78,14 +134,19 @@ async def capture_media( # noqa: PLR0913
storage_key = existing["storage_key"] storage_key = existing["storage_key"]
file_size = existing["file_size"] file_size = existing["file_size"]
downloaded = True downloaded = True
else: elif force or not await repository.is_evicted(ctx.pool, unique_id):
target = message if getattr(message, kind or "", None) is obj else obj target = _download_target(message)
buffer = await client.download_media(target, in_memory=True) try:
if isinstance(buffer, BytesIO): data = await _download(client, target, chat_id, message_id)
data = buffer.getvalue() except (FloodWait, FloodPremiumWait) as exc:
flood = exc
data = None
if data is not None:
storage_key = ctx.storage.put(data) storage_key = ctx.storage.put(data)
file_size = len(data) file_size = len(data)
downloaded = True downloaded = True
else:
evicted = True
await repository.insert_media( await repository.insert_media(
ctx.pool, ctx.pool,
ctx.account_id, ctx.account_id,
@@ -98,4 +159,14 @@ async def capture_media( # noqa: PLR0913
ttl, ttl,
unique_id, unique_id,
downloaded=downloaded, downloaded=downloaded,
evicted=evicted,
) )
if (
force
and storage_key is not None
and await repository.unevict(ctx.pool, unique_id, storage_key, file_size)
):
ctx.storage.purge(storage_key)
if flood is not None:
raise flood
return downloaded or not want
@@ -21,10 +21,10 @@ def active_username(usernames: list[raw.types.Username]) -> str | None:
return usernames[0].username if usernames else None return usernames[0].username if usernames else None
def snapshot_from_user( async def snapshot_from_user(
client: Client, raw_user: raw.types.User client: Client, raw_user: raw.types.User
) -> tuple[ProfileFields, str | None, str | None]: ) -> tuple[ProfileFields, str | None, str | None]:
user = User._parse(client, raw_user) # noqa: SLF001 user = await User._parse(client, raw_user) # noqa: SLF001
photo = user.photo if user else None photo = user.photo if user else None
photo_unique_id = photo.big_photo_unique_id if photo else None photo_unique_id = photo.big_photo_unique_id if photo else None
photo_file_id = photo.big_file_id if photo else None photo_file_id = photo.big_file_id if photo else None
@@ -7,12 +7,13 @@ from userbot.modules.capture.context import CaptureContext
from userbot.modules.groups.repository import insert_chat_history from userbot.modules.groups.repository import insert_chat_history
from userbot.modules.profiles.parse import snapshot_from_chat from userbot.modules.profiles.parse import snapshot_from_chat
from userbot.modules.profiles.repository import write_profile from userbot.modules.profiles.repository import write_profile
from userbot.modules.raw_json import raw_json
async def save_private(ctx: CaptureContext, chat: Chat) -> bool: async def save_private(ctx: CaptureContext, chat: Chat) -> bool:
chat_id = chat.id or 0 chat_id = chat.id or 0
fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat) fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat)
await write_profile(ctx.pool, ctx.account_id, chat_id, fields, str(chat)) await write_profile(ctx.pool, ctx.account_id, chat_id, fields, raw_json(chat))
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id
@@ -35,7 +36,7 @@ async def save_group(ctx: CaptureContext, chat: Chat) -> None:
photo_unique_id, photo_unique_id,
None, None,
datetime.now(UTC), datetime.now(UTC),
str(chat), raw_json(chat),
) )
if photo_file_id and photo_unique_id: if photo_file_id and photo_unique_id:
await note_avatar( await note_avatar(
+23
View File
@@ -0,0 +1,23 @@
import json
from typing import Any
from pyrogram.raw.core import TLObject
from pyrogram.types import Contact, Object
EMPTY_JSON = "{}"
PHONE_KEY = "phone_number"
def _default(obj: Any) -> Any: # noqa: ANN401
if isinstance(obj, TLObject):
return TLObject.default(obj)
data = Object.default(obj)
if isinstance(obj, Contact) and isinstance(data, dict) and PHONE_KEY in data:
data[PHONE_KEY] = obj.phone_number
return data
def raw_json(obj: Object | TLObject | None) -> str:
if obj is None:
return EMPTY_JSON
return json.dumps(obj, indent=4, default=_default, ensure_ascii=False)
@@ -22,6 +22,20 @@ ON CONFLICT (account_id, peer_id, story_id) DO UPDATE SET
""" """
async def is_downloaded(
pool: asyncpg.Pool, account_id: int, peer_id: int, story_id: int
) -> bool:
return bool(
await pool.fetchval(
"SELECT downloaded FROM stories "
"WHERE account_id = $1 AND peer_id = $2 AND story_id = $3",
account_id,
peer_id,
story_id,
)
)
async def upsert_story( # noqa: PLR0913 async def upsert_story( # noqa: PLR0913
pool: asyncpg.Pool, pool: asyncpg.Pool,
account_id: int, account_id: int,
@@ -0,0 +1,48 @@
from pyrogram import Client
from pyrogram.types import Story
from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import download_bytes
from userbot.modules.raw_json import raw_json
from userbot.modules.stories import repository
def story_peer_id(story: Story) -> int:
if story.chat is not None:
return story.chat.id or 0
if story.from_user is not None:
return story.from_user.id or 0
return 0
async def save_story(client: Client, capture: CaptureContext, story: Story) -> None:
peer_id = story_peer_id(story)
storage_key: str | None = None
file_size: int | None = None
downloaded = False
stored = await repository.is_downloaded(
capture.pool, capture.account_id, peer_id, story.id
)
if not (stored or story.deleted or story.media is None):
data = await download_bytes(client, story)
if data is not None:
storage_key = capture.storage.put(data)
file_size = len(data)
downloaded = True
await repository.upsert_story(
capture.pool,
capture.account_id,
peer_id,
story.id,
story.date,
story.expire_date,
story.caption,
story.media.name.lower() if story.media else None,
storage_key,
file_size,
story.views,
raw_json(story.raw),
pinned=bool(story.pinned),
deleted=bool(story.deleted),
downloaded=downloaded,
)
+19 -16
View File
@@ -1,5 +1,6 @@
import asyncio import asyncio
import contextlib import contextlib
import signal
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@@ -9,8 +10,8 @@ import asyncpg
import uvloop import uvloop
from dependencies.container import container from dependencies.container import container
from userbot import PyroClient
from userbot.modules.capture import CaptureContext, build_capture_context from userbot.modules.capture import CaptureContext, build_capture_context
from userbot.modules.client import PyroClient
from userbot.modules.jobs import JobConsumer from userbot.modules.jobs import JobConsumer
from utils.env import env from utils.env import env
from utils.jobs import enqueue from utils.jobs import enqueue
@@ -26,6 +27,8 @@ from utils.storage import ContentAddressedStorage
setup_logging() setup_logging()
SESSION_FLUSH_INTERVAL = 60
@dataclass @dataclass
class RunningAccount: class RunningAccount:
@@ -46,18 +49,6 @@ async def _cancel(task: asyncio.Task) -> None:
await task await task
async def _enqueue_once(pool: asyncpg.Pool, account_id: int, kind: str) -> None:
existing = await pool.fetchval(
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = $2 "
"AND status IN ('pending', 'running') LIMIT 1",
account_id,
kind,
)
if existing is None:
await enqueue(pool, account_id, kind, {})
logger.info(f"[green]Queued {kind}.[/]")
class AccountRegistry: class AccountRegistry:
def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None: def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
self._pool = pool self._pool = pool
@@ -118,8 +109,8 @@ class AccountRegistry:
client, asyncio.create_task(consumer.run()), device_model client, asyncio.create_task(consumer.run()), device_model
) )
logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})") logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
await _enqueue_once(self._pool, account_id, "sync_dialogs") await enqueue(self._pool, account_id, "sync_dialogs", {})
await _enqueue_once(self._pool, account_id, "sync_contacts") await enqueue(self._pool, account_id, "sync_contacts", {})
async def _stop(self, session_name: str) -> None: async def _stop(self, session_name: str) -> None:
account = self._running.pop(session_name, None) account = self._running.pop(session_name, None)
@@ -149,6 +140,14 @@ class AccountRegistry:
logger.info(f"[yellow]Account logged out:[/] {path.stem}") logger.info(f"[yellow]Account logged out:[/] {path.stem}")
async def _flush_sessions(registry: AccountRegistry) -> None:
while True:
await asyncio.sleep(SESSION_FLUSH_INTERVAL)
for client in registry.clients:
with contextlib.suppress(Exception):
await client.storage.save()
async def _listen_changes( async def _listen_changes(
registry: AccountRegistry, tasks: set[asyncio.Task] registry: AccountRegistry, tasks: set[asyncio.Task]
) -> asyncpg.Connection: ) -> asyncpg.Connection:
@@ -193,14 +192,18 @@ async def runner() -> None:
registry = AccountRegistry(pool, storage) registry = AccountRegistry(pool, storage)
tasks: set[asyncio.Task] = set() tasks: set[asyncio.Task] = set()
listen_conn: asyncpg.Connection | None = None listen_conn: asyncpg.Connection | None = None
stop = asyncio.Event()
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
flush = asyncio.create_task(_flush_sessions(registry))
try: try:
await registry.sync() await registry.sync()
if not registry.clients: if not registry.clients:
logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]") logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
listen_conn = await _listen_changes(registry, tasks) listen_conn = await _listen_changes(registry, tasks)
logger.info("[green]Userbot running.[/]") logger.info("[green]Userbot running.[/]")
await asyncio.Event().wait() await stop.wait()
finally: finally:
await _cancel(flush)
if listen_conn is not None: if listen_conn is not None:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
await listen_conn.close() await listen_conn.close()
+4
View File
@@ -0,0 +1,4 @@
IMMUTABLE_HEADERS = {"Cache-Control": "public, max-age=31536000, immutable"}
DAY_HEADERS = {"Cache-Control": "public, max-age=86400"}
SHORT_HEADERS = {"Cache-Control": "public, max-age=300"}
NO_STORE_HEADERS = {"Cache-Control": "no-cache"}
+50
View File
@@ -493,3 +493,53 @@ class Dialog(SQLModel, table=True):
onupdate=func.now(), 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
+8
View File
@@ -31,6 +31,12 @@ class TelegramSettings(BaseSettings):
sessions_dir: str = "sessions" sessions_dir: str = "sessions"
class BotsSettings(BaseSettings):
sessions_dir: str = "sessions/bots"
backfill_per_chat: int = 0
backfill_chats: int = 0
class ApiSettings(BaseSettings): class ApiSettings(BaseSettings):
host: str = "0.0.0.0" # noqa: S104 host: str = "0.0.0.0" # noqa: S104
port: int = 8080 port: int = 8080
@@ -50,12 +56,14 @@ class LogSettings(BaseSettings):
level: str = "INFO" level: str = "INFO"
level_external: str = "WARNING" level_external: str = "WARNING"
show_time: bool = False show_time: bool = False
show_locals: bool = False
console_width: int = 150 console_width: int = 150
class Settings(BaseSettings): class Settings(BaseSettings):
db: DatabaseSettings = Field(default_factory=DatabaseSettings) db: DatabaseSettings = Field(default_factory=DatabaseSettings)
tg: TelegramSettings = Field(default_factory=TelegramSettings) tg: TelegramSettings = Field(default_factory=TelegramSettings)
bots: BotsSettings = Field(default_factory=BotsSettings)
api: ApiSettings = Field(default_factory=ApiSettings) api: ApiSettings = Field(default_factory=ApiSettings)
auth: AuthSettings = Field(default_factory=AuthSettings) auth: AuthSettings = Field(default_factory=AuthSettings)
storage: StorageSettings = Field(default_factory=StorageSettings) storage: StorageSettings = Field(default_factory=StorageSettings)
+242
View File
@@ -0,0 +1,242 @@
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)
+6
View File
@@ -0,0 +1,6 @@
import asyncio
from importlib import import_module
async def preload(module: str) -> None:
await asyncio.to_thread(import_module, module)
+11 -1
View File
@@ -9,12 +9,22 @@ JOBS_CHANGED_CHANNEL = "jobs_changed"
async def enqueue( async def enqueue(
pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any] pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
) -> int: ) -> int:
params_json = json.dumps(params)
existing = await pool.fetchval(
"SELECT id FROM jobs WHERE account_id = $1 AND kind = $2 "
"AND params = $3::jsonb AND status IN ('pending', 'running') LIMIT 1",
account_id,
kind,
params_json,
)
if existing is not None:
return existing
job_id = await pool.fetchval( job_id = await pool.fetchval(
"INSERT INTO jobs (account_id, kind, params) " "INSERT INTO jobs (account_id, kind, params) "
"VALUES ($1, $2, $3::jsonb) RETURNING id", "VALUES ($1, $2, $3::jsonb) RETURNING id",
account_id, account_id,
kind, kind,
json.dumps(params), params_json,
) )
await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}") await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}")
return job_id return job_id
+2 -2
View File
@@ -20,13 +20,13 @@ def setup_logging() -> None:
markup=True, markup=True,
rich_tracebacks=True, rich_tracebacks=True,
enable_link_path=False, enable_link_path=False,
tracebacks_show_locals=True, tracebacks_show_locals=env.log.show_locals,
omit_repeated_times=False, omit_repeated_times=False,
show_time=env.log.show_time, show_time=env.log.show_time,
) )
], ],
) )
install(console=console, show_locals=True) install(console=console, show_locals=env.log.show_locals)
logger = logging.getLogger("beavergram") logger = logging.getLogger("beavergram")
+18
View File
@@ -79,6 +79,24 @@ class FolderSpec:
is_chatlist: bool = False is_chatlist: bool = False
class PolicyParent(BaseModel):
scope_type: ScopeType
scope_id: int | None
title: str | None = None
toggles: CaptureToggles
class PolicyRuleView(PolicyRecord):
title: str | None = None
parent: PolicyParent | None = None
class PolicyScopeView(BaseModel):
title: str | None
rule: PolicyRecord | None
parent: PolicyParent | None
@dataclass(slots=True) @dataclass(slots=True)
class PolicySet: class PolicySet:
chat: dict[int, CaptureToggles] = field(default_factory=dict) chat: dict[int, CaptureToggles] = field(default_factory=dict)
+103
View File
@@ -0,0 +1,103 @@
import asyncpg
from utils.policy import repository
from utils.policy.models import (
ChatKind,
ChatMeta,
FolderSpec,
PolicyParent,
PolicyRecord,
PolicyRuleView,
PolicyScopeView,
PolicySet,
ScopeType,
)
from utils.policy.resolver import resolve_parent
from utils.read import chats
from utils.read.models import ChatListItem
def _meta(chat_id: int, chat: ChatListItem | None) -> ChatMeta:
if chat_id > 0:
kind = ChatKind.DM
elif chat is not None and chat.is_broadcast:
kind = ChatKind.CHANNEL
else:
kind = ChatKind.GROUP
return ChatMeta(
chat_id=chat_id,
kind=kind,
is_bot=bool(chat and chat.is_bot),
is_contact=chat.is_contact if chat else None,
)
def _without_chats(policies: PolicySet) -> PolicySet:
return PolicySet(folder=policies.folder, defaults=policies.defaults)
def _chat_parent(
chat_id: int,
chat: ChatListItem | None,
folders: list[FolderSpec],
policies: PolicySet,
) -> PolicyParent:
return resolve_parent(_meta(chat_id, chat), folders, _without_chats(policies))
async def list_rules(pool: asyncpg.Pool, account_id: int) -> list[PolicyRuleView]:
records = await repository.list_policies(pool, account_id)
folders = await repository.list_folders(pool, account_id)
policies = await repository.load_policy_set(pool, account_id)
chat_ids = [
record.scope_id
for record in records
if record.scope_type is ScopeType.CHAT and record.scope_id is not None
]
known = {
chat.chat_id: chat for chat in await chats.get_chats(pool, account_id, chat_ids)
}
titles = {folder.folder_id: folder.title for folder in folders}
rules: list[PolicyRuleView] = []
for record in records:
rule = PolicyRuleView(**record.model_dump())
if record.scope_type is ScopeType.FOLDER and record.scope_id is not None:
rule.title = titles.get(record.scope_id)
if record.scope_type is ScopeType.CHAT and record.scope_id is not None:
chat = known.get(record.scope_id)
rule.title = chat.title if chat else None
rule.parent = _chat_parent(record.scope_id, chat, folders, policies)
rules.append(rule)
return rules
async def _rule(
pool: asyncpg.Pool, account_id: int, scope_type: ScopeType, scope_id: int
) -> PolicyRecord | None:
return await repository.find_policy(pool, account_id, scope_type, scope_id)
async def chat_scope(
pool: asyncpg.Pool, account_id: int, chat_id: int
) -> PolicyScopeView:
folders = await repository.list_folders(pool, account_id)
policies = await repository.load_policy_set(pool, account_id)
chat = await chats.get_chat(pool, account_id, chat_id)
return PolicyScopeView(
title=chat.title if chat else None,
rule=await _rule(pool, account_id, ScopeType.CHAT, chat_id),
parent=_chat_parent(chat_id, chat, folders, policies),
)
async def folder_scope(
pool: asyncpg.Pool, account_id: int, folder_id: int
) -> PolicyScopeView | None:
folder = await repository.get_folder(pool, account_id, folder_id)
if folder is None:
return None
return PolicyScopeView(
title=folder.title,
rule=await _rule(pool, account_id, ScopeType.FOLDER, folder_id),
parent=None,
)
Loaded 100 of 293 files, more files were not shown because too many files have changed in this diff. Show more