Compare commits
90
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ec5979437d | ||
|
|
59f171a62d | ||
|
|
e3559da2a7 | ||
|
|
dcd6786ca5 | ||
|
|
32d58193f3 | ||
|
|
59d0e2c8ad | ||
|
|
850d39ad6f | ||
|
|
b8a38092eb | ||
|
|
2836dcb3a9 | ||
|
|
01cf01b629 | ||
|
|
7f124059be | ||
|
|
5eacabd101 | ||
|
|
0362777a51 | ||
|
|
dfc726976b | ||
|
|
cb01d7dff8 | ||
|
|
f646f29ff9 | ||
|
|
945adcf879 | ||
|
|
d213dfdf64 | ||
|
|
566ee09431 | ||
|
|
09c5f78061 | ||
|
|
897421e7c4 | ||
|
|
3223f62dea | ||
|
|
c9c088ed73 | ||
|
|
580927da0a | ||
|
|
06ae3e928f | ||
|
|
55634984fd | ||
|
|
f2c245707a | ||
|
|
c20e132872 | ||
|
|
4d16105e92 | ||
|
|
b1ffebe858 | ||
|
|
2512a7988d | ||
|
|
56146e4e8d | ||
|
|
5fb8f8a534 | ||
|
|
98ab8a579b | ||
|
|
a5d10146fd | ||
|
|
5f1f18b343 | ||
|
|
8eb31ed4d3 | ||
|
|
e67dbdb325 | ||
|
|
869f8291fb | ||
|
|
283eab1eb0 | ||
|
|
7f4d500302 | ||
|
|
cf48d2db3e | ||
|
|
114058e6ae | ||
|
|
eb8187ff6e | ||
|
|
53dc805fc3 | ||
|
|
ed8489616e | ||
|
|
b60815d042 | ||
|
|
468f72ed11 | ||
|
|
59d3f214e8 | ||
|
|
34dcbaac2b | ||
|
|
c8e5cff871 | ||
|
|
fc4758aa0b | ||
|
|
b0ed041c2e | ||
|
|
4bd031f21e | ||
|
|
089a31db8d | ||
|
|
81ebbdc93e | ||
|
|
a25f79543a | ||
|
|
c61ecfe8f2 | ||
|
|
3d0b15bd2c | ||
|
|
0e54786672 | ||
|
|
5778611a38 | ||
|
|
b138533d88 | ||
|
|
0b6cbac945 | ||
|
|
ff80141785 | ||
|
|
d206f7e5ca | ||
|
|
4ebe6de60b | ||
|
|
02fb567c1f | ||
|
|
ccbc167005 | ||
|
|
5ec0212af5 | ||
|
|
ab59ef4525 | ||
|
|
676680319b | ||
|
|
6a6e63ff75 | ||
|
|
1fa4f18a9b | ||
|
|
9c12628304 | ||
|
|
cf277c56cf | ||
|
|
8f7d47476c | ||
|
|
c832513e5c | ||
|
|
0399145791 | ||
|
|
0b8e90159f | ||
|
|
995322b17a | ||
|
|
6ae28cb321 | ||
|
|
e833de245a | ||
|
|
ef739838a5 | ||
|
|
3e08698b62 | ||
|
|
683b9a31a3 | ||
|
|
004fd56f2c | ||
|
|
12cc7d57e3 | ||
|
|
1e143c7573 | ||
|
|
ee63f8b783 | ||
|
|
6b6edc9a0d |
No files matched your search
@@ -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"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
_commit: v0.1.4
|
||||||
|
_src_path: https://git.kotikot.com/templates/caddy.git
|
||||||
|
project_slug: beavergram
|
||||||
|
routes:
|
||||||
|
- api
|
||||||
@@ -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
|
||||||
@@ -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
@@ -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
@@ -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/
|
||||||
@@ -1,12 +1,12 @@
|
|||||||
{
|
{
|
||||||
"mcpServers": {
|
"mcpServers": {
|
||||||
"svelte": {
|
"context7": {
|
||||||
"type": "http",
|
"type": "http",
|
||||||
"url": "https://mcp.svelte.dev/mcp"
|
"url": "https://mcp.context7.com/mcp"
|
||||||
},
|
},
|
||||||
"context7": {
|
"svelte": {
|
||||||
"type": "http",
|
"type": "http",
|
||||||
"url": "https://mcp.context7.com/mcp"
|
"url": "https://mcp.svelte.dev/mcp"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
+18
-22
@@ -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 ]
|
|
||||||
@@ -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
|
||||||
```
|
```
|
||||||
@@ -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 --profile db --profile migrate run --rm migrator $(filter-out $@,$(MAKECMDGOALS))
|
|
||||||
|
|
||||||
session-create:
|
|
||||||
cd backend && uv run python scripts/session/create.py
|
|
||||||
|
|
||||||
%:
|
|
||||||
@:
|
|
||||||
@@ -1,2 +1,46 @@
|
|||||||
# beavergram
|
# beavergram
|
||||||
[](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
|
||||||
|
```
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
.venv
|
||||||
|
.env
|
||||||
|
__pycache__
|
||||||
|
*.pyc
|
||||||
|
.ruff_cache
|
||||||
|
.pytest_cache
|
||||||
|
sessions
|
||||||
+27
-9
@@ -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
|
||||||
@@ -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
@@ -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"]
|
||||||
@@ -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
|
||||||
@@ -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:
|
||||||
@@ -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,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,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")
|
||||||
@@ -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"]
|
||||||
|
|||||||
@@ -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())
|
||||||
@@ -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())
|
||||||
@@ -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,19 +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.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="/")
|
||||||
|
|
||||||
|
|
||||||
@@ -49,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()
|
||||||
|
|
||||||
@@ -71,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)
|
||||||
@@ -88,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)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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."
|
||||||
|
)
|
||||||
@@ -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},
|
]
|
||||||
)
|
mime = resolve_mime(item.kind, item.mime, item.file_name)
|
||||||
return item.model_dump(mode="json")
|
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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -10,6 +10,7 @@ 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,
|
||||||
@@ -49,6 +50,14 @@ async def list_chats(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@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}")
|
@router.get("/chats/{chat_id}")
|
||||||
async def get_chat(
|
async def get_chat(
|
||||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -6,6 +6,7 @@ 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.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,
|
||||||
@@ -17,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:
|
||||||
@@ -41,14 +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=IMMUTABLE_HEADERS,
|
headers=headers,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -70,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:
|
||||||
@@ -79,8 +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=DAY_HEADERS,
|
headers=headers,
|
||||||
)
|
)
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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")
|
||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -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,
|
||||||
)
|
)
|
||||||
@@ -1,3 +0,0 @@
|
|||||||
from userbot.modules.client import DEVICE_MODEL, DEVICE_MODEL_LIMIT, PyroClient
|
|
||||||
|
|
||||||
__all__ = ["DEVICE_MODEL", "DEVICE_MODEL_LIMIT", "PyroClient"]
|
|
||||||
@@ -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]
|
||||||
@@ -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]
|
||||||
@@ -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]
|
||||||
@@ -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,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]
|
||||||
@@ -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,)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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(
|
||||||
await capture_message(client, message, capture, toggles)
|
chat_id, max_id=max_id, min_id=min_id
|
||||||
if should_transcribe_on_backfill(message, self_id) and message.chat:
|
):
|
||||||
meta = meta_from_chat(message.chat, capture.contacts.ids)
|
try:
|
||||||
already = await stt_repo.is_transcribed(
|
await capture_message(client, message, capture, toggles)
|
||||||
capture.pool, capture.account_id, chat_id, message.id
|
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
||||||
)
|
except (FloodWait, FloodPremiumWait, PeerIdInvalid):
|
||||||
if capture.resolve(meta).stt and not already:
|
raise
|
||||||
await safe_transcribe(client, capture, chat_id, message.id)
|
except Exception:
|
||||||
|
errors += 1
|
||||||
|
logger.exception(f"Backfill capture failed for {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,16 +68,28 @@ 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:
|
await _save_user(ctx, user)
|
||||||
continue
|
|
||||||
if isinstance(user, User):
|
|
||||||
await _save_user(ctx, user)
|
|
||||||
|
|
||||||
|
|
||||||
@register("enrich_chat")
|
@register("enrich_chat")
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -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()
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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)
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
import asyncio
|
||||||
|
from importlib import import_module
|
||||||
|
|
||||||
|
|
||||||
|
async def preload(module: str) -> None:
|
||||||
|
await asyncio.to_thread(import_module, module)
|
||||||
@@ -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
|
||||||
@@ -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")
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -1,11 +1,14 @@
|
|||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
|
|
||||||
from utils.policy.models import (
|
from utils.policy.models import (
|
||||||
|
DEFAULT_SCOPE_BY_KIND,
|
||||||
CaptureToggles,
|
CaptureToggles,
|
||||||
ChatKind,
|
ChatKind,
|
||||||
ChatMeta,
|
ChatMeta,
|
||||||
FolderSpec,
|
FolderSpec,
|
||||||
|
PolicyParent,
|
||||||
PolicySet,
|
PolicySet,
|
||||||
|
ScopeType,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -33,14 +36,29 @@ def folder_contains(folder: FolderSpec, chat: ChatMeta) -> bool:
|
|||||||
return _category_match(folder, chat)
|
return _category_match(folder, chat)
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_parent(
|
||||||
|
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
|
||||||
|
) -> PolicyParent:
|
||||||
|
for folder in sorted(folders, key=lambda f: f.order_index):
|
||||||
|
toggles = policies.folder.get(folder.folder_id)
|
||||||
|
if toggles is not None and folder_contains(folder, chat):
|
||||||
|
return PolicyParent(
|
||||||
|
scope_type=ScopeType.FOLDER,
|
||||||
|
scope_id=folder.folder_id,
|
||||||
|
title=folder.title,
|
||||||
|
toggles=toggles,
|
||||||
|
)
|
||||||
|
return PolicyParent(
|
||||||
|
scope_type=DEFAULT_SCOPE_BY_KIND[chat.kind],
|
||||||
|
scope_id=None,
|
||||||
|
toggles=policies.defaults.get(chat.kind, CaptureToggles()),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def resolve(
|
def resolve(
|
||||||
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
|
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
|
||||||
) -> CaptureToggles:
|
) -> CaptureToggles:
|
||||||
chat_override = policies.chat.get(chat.chat_id)
|
chat_override = policies.chat.get(chat.chat_id)
|
||||||
if chat_override is not None:
|
if chat_override is not None:
|
||||||
return chat_override
|
return chat_override
|
||||||
for folder in sorted(folders, key=lambda f: f.order_index):
|
return resolve_parent(chat, folders, policies).toggles
|
||||||
toggles = policies.folder.get(folder.folder_id)
|
|
||||||
if toggles is not None and folder_contains(folder, chat):
|
|
||||||
return toggles
|
|
||||||
return policies.defaults.get(chat.kind, CaptureToggles())
|
|
||||||
@@ -7,17 +7,18 @@ from utils.read.models import AccountView
|
|||||||
|
|
||||||
ACCOUNTS_CHANGED_CHANNEL = "accounts_changed"
|
ACCOUNTS_CHANGED_CHANNEL = "accounts_changed"
|
||||||
|
|
||||||
_ACCOUNT_COLS = "account_id, label, phone, tg_user_id, is_active, device_model"
|
_ACCOUNT_COLS = "account_id, label, phone, tg_user_id, is_active, is_bot, device_model"
|
||||||
|
|
||||||
_UPSERT_ACCOUNT = """
|
_UPSERT_ACCOUNT = """
|
||||||
INSERT INTO accounts
|
INSERT INTO accounts
|
||||||
(tg_user_id, label, phone, session_name, is_active, raw, updated_at)
|
(tg_user_id, label, phone, session_name, is_active, is_bot, raw, updated_at)
|
||||||
VALUES ($1, $2, $3, $4, TRUE, $5::jsonb, now())
|
VALUES ($1, $2, $3, $4, TRUE, $6, $5::jsonb, now())
|
||||||
ON CONFLICT (tg_user_id) DO UPDATE SET
|
ON CONFLICT (tg_user_id) DO UPDATE SET
|
||||||
label = EXCLUDED.label,
|
label = EXCLUDED.label,
|
||||||
phone = EXCLUDED.phone,
|
phone = EXCLUDED.phone,
|
||||||
session_name = EXCLUDED.session_name,
|
session_name = EXCLUDED.session_name,
|
||||||
is_active = TRUE,
|
is_active = TRUE,
|
||||||
|
is_bot = EXCLUDED.is_bot,
|
||||||
raw = EXCLUDED.raw,
|
raw = EXCLUDED.raw,
|
||||||
updated_at = now()
|
updated_at = now()
|
||||||
RETURNING account_id
|
RETURNING account_id
|
||||||
@@ -30,10 +31,48 @@ RETURNING {_ACCOUNT_COLS}
|
|||||||
""" # noqa: S608
|
""" # noqa: S608
|
||||||
|
|
||||||
|
|
||||||
|
_self_ids: dict[int, int] = {}
|
||||||
|
|
||||||
|
|
||||||
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
|
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
|
||||||
return await pool.fetchval(
|
cached = _self_ids.get(account_id)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
tg_user_id = await pool.fetchval(
|
||||||
"SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id
|
"SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id
|
||||||
)
|
)
|
||||||
|
if tg_user_id is not None:
|
||||||
|
_self_ids[account_id] = tg_user_id
|
||||||
|
return tg_user_id
|
||||||
|
|
||||||
|
|
||||||
|
_known_ids: set[int] = set()
|
||||||
|
|
||||||
|
|
||||||
|
def _account_label(row: asyncpg.Record) -> str:
|
||||||
|
named = f", {row['label']}" if row["label"] else ""
|
||||||
|
return f"account_id={row['account_id']} (Telegram id {row['tg_user_id']}{named})"
|
||||||
|
|
||||||
|
|
||||||
|
async def unknown_account_hint(pool: asyncpg.Pool, account_id: int) -> str | None:
|
||||||
|
"""Explain an account_id that matches no account, or None if it is valid."""
|
||||||
|
if account_id in _known_ids:
|
||||||
|
return None
|
||||||
|
rows = await pool.fetch(
|
||||||
|
"SELECT account_id, tg_user_id, label FROM accounts ORDER BY account_id"
|
||||||
|
)
|
||||||
|
_known_ids.update(row["account_id"] for row in rows)
|
||||||
|
if account_id in _known_ids:
|
||||||
|
return None
|
||||||
|
known = "; ".join(_account_label(row) for row in rows) or "none"
|
||||||
|
mistaken = next((row for row in rows if row["tg_user_id"] == account_id), None)
|
||||||
|
if mistaken is not None:
|
||||||
|
return (
|
||||||
|
f"{account_id} is a Telegram user id, not an account_id. In beavergram "
|
||||||
|
f"that account is account_id={mistaken['account_id']}. "
|
||||||
|
f"Retry with account_id={mistaken['account_id']}."
|
||||||
|
)
|
||||||
|
return f"No account with account_id={account_id}. Archived accounts: {known}."
|
||||||
|
|
||||||
|
|
||||||
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
|
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
|
||||||
@@ -51,7 +90,9 @@ async def get_account(pool: asyncpg.Pool, account_id: int) -> AccountView | None
|
|||||||
return AccountView(**dict(row)) if row else None
|
return AccountView(**dict(row)) if row else None
|
||||||
|
|
||||||
|
|
||||||
async def sync_account(pool: asyncpg.Pool, me: User, session_name: str) -> int:
|
async def sync_account(
|
||||||
|
pool: asyncpg.Pool, me: User, session_name: str, *, is_bot: bool = False
|
||||||
|
) -> int:
|
||||||
raw = json.dumps(
|
raw = json.dumps(
|
||||||
{
|
{
|
||||||
"id": me.id,
|
"id": me.id,
|
||||||
@@ -63,7 +104,7 @@ async def sync_account(pool: asyncpg.Pool, me: User, session_name: str) -> int:
|
|||||||
)
|
)
|
||||||
label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username
|
label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username
|
||||||
return await pool.fetchval(
|
return await pool.fetchval(
|
||||||
_UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw
|
_UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw, is_bot
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ async def message_volume(
|
|||||||
|
|
||||||
|
|
||||||
async def response_stats(
|
async def response_stats(
|
||||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
pool: asyncpg.Pool, account_id: int, chat_id: int, *, days: int = 365
|
||||||
) -> ResponseStats:
|
) -> ResponseStats:
|
||||||
self_id = await self_user_id(pool, account_id)
|
self_id = await self_user_id(pool, account_id)
|
||||||
rows = await pool.fetch(
|
rows = await pool.fetch(
|
||||||
@@ -43,6 +43,7 @@ async def response_stats(
|
|||||||
"lag(date) OVER w AS prev_date "
|
"lag(date) OVER w AS prev_date "
|
||||||
"FROM messages "
|
"FROM messages "
|
||||||
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
|
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
|
||||||
|
"AND date >= $4 "
|
||||||
"WINDOW w AS (ORDER BY date, message_id)), "
|
"WINDOW w AS (ORDER BY date, message_id)), "
|
||||||
"resp AS ("
|
"resp AS ("
|
||||||
"SELECT (sender_id = $3) AS is_mine, "
|
"SELECT (sender_id = $3) AS is_mine, "
|
||||||
@@ -55,6 +56,7 @@ async def response_stats(
|
|||||||
account_id,
|
account_id,
|
||||||
chat_id,
|
chat_id,
|
||||||
self_id,
|
self_id,
|
||||||
|
datetime.now(UTC) - timedelta(days=days),
|
||||||
)
|
)
|
||||||
stats = ResponseStats(
|
stats = ResponseStats(
|
||||||
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
|
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ WHERE account_id = $1 AND owner_id = $2 AND unique_id = $3
|
|||||||
_AVATAR_HISTORY = """
|
_AVATAR_HISTORY = """
|
||||||
SELECT unique_id, first_seen_at, downloaded FROM avatars
|
SELECT unique_id, first_seen_at, downloaded FROM avatars
|
||||||
WHERE account_id = $1 AND owner_id = $2
|
WHERE account_id = $1 AND owner_id = $2
|
||||||
ORDER BY first_seen_at DESC
|
ORDER BY first_seen_at DESC LIMIT 200
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,157 @@
|
|||||||
|
"""Registry of bot tokens the watcher holds.
|
||||||
|
|
||||||
|
Only session bookkeeping lives here: what the archive itself keeps about a
|
||||||
|
bot is an ordinary account row and ordinary messages, read through
|
||||||
|
``utils.read.accounts`` and ``utils.read.chats`` like any other account.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncpg
|
||||||
|
|
||||||
|
from utils.read.models import BotView
|
||||||
|
|
||||||
|
BOTS_CHANGED_CHANNEL = "watched_bots_changed"
|
||||||
|
|
||||||
|
STATUS_PENDING = "pending"
|
||||||
|
STATUS_OK = "ok"
|
||||||
|
STATUS_UNAUTHORIZED = "unauthorized"
|
||||||
|
STATUS_ERROR = "error"
|
||||||
|
STATUS_REMOVED = "removed"
|
||||||
|
|
||||||
|
_BOT_COLS = """
|
||||||
|
SELECT bot_id, account_id, username, title, enabled, status, status_detail,
|
||||||
|
backfilled_at, checked_at, added_at
|
||||||
|
FROM watched_bots
|
||||||
|
"""
|
||||||
|
|
||||||
|
_REGISTER = """
|
||||||
|
INSERT INTO watched_bots (bot_id, status) VALUES ($1, 'pending')
|
||||||
|
ON CONFLICT (bot_id) DO UPDATE SET
|
||||||
|
enabled = true, status = 'pending', status_detail = NULL, updated_at = now()
|
||||||
|
"""
|
||||||
|
|
||||||
|
_REMOVE = """
|
||||||
|
UPDATE watched_bots
|
||||||
|
SET enabled = false, status = 'removed', status_detail = NULL, updated_at = now()
|
||||||
|
WHERE bot_id = $1 AND enabled
|
||||||
|
"""
|
||||||
|
|
||||||
|
_BACKLOG_CURSOR = """
|
||||||
|
SELECT backfill_pts, backfill_qts, backfill_date
|
||||||
|
FROM watched_bots
|
||||||
|
WHERE bot_id = $1 AND backfill_pts IS NOT NULL
|
||||||
|
"""
|
||||||
|
|
||||||
|
_SAVE_BACKLOG_CURSOR = """
|
||||||
|
UPDATE watched_bots SET
|
||||||
|
backfill_pts = $2, backfill_qts = $3, backfill_date = $4, updated_at = now()
|
||||||
|
WHERE bot_id = $1
|
||||||
|
"""
|
||||||
|
|
||||||
|
_MARK_BACKFILLED = """
|
||||||
|
UPDATE watched_bots SET
|
||||||
|
backfilled_at = now(),
|
||||||
|
backfill_pts = NULL,
|
||||||
|
backfill_qts = NULL,
|
||||||
|
backfill_date = NULL,
|
||||||
|
updated_at = now()
|
||||||
|
WHERE bot_id = $1
|
||||||
|
"""
|
||||||
|
|
||||||
|
_SET_STATUS = """
|
||||||
|
UPDATE watched_bots SET
|
||||||
|
status = $2,
|
||||||
|
status_detail = $3,
|
||||||
|
username = coalesce($4, username),
|
||||||
|
title = coalesce($5, title),
|
||||||
|
checked_at = now(),
|
||||||
|
updated_at = now()
|
||||||
|
WHERE bot_id = $1
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
async def notify_bots_changed(pool: asyncpg.Pool) -> None:
|
||||||
|
await pool.execute("SELECT pg_notify($1, '')", BOTS_CHANGED_CHANNEL)
|
||||||
|
|
||||||
|
|
||||||
|
async def list_bots(pool: asyncpg.Pool) -> list[BotView]:
|
||||||
|
rows = await pool.fetch(f"{_BOT_COLS} ORDER BY added_at")
|
||||||
|
return [BotView(**dict(row)) for row in rows]
|
||||||
|
|
||||||
|
|
||||||
|
async def get_bot(pool: asyncpg.Pool, bot_id: int) -> BotView | None:
|
||||||
|
row = await pool.fetchrow(f"{_BOT_COLS} WHERE bot_id = $1", bot_id)
|
||||||
|
return BotView(**dict(row)) if row else None
|
||||||
|
|
||||||
|
|
||||||
|
async def enabled_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||||
|
rows = await pool.fetch(
|
||||||
|
"SELECT bot_id FROM watched_bots WHERE enabled AND status <> $1",
|
||||||
|
STATUS_UNAUTHORIZED,
|
||||||
|
)
|
||||||
|
return {row["bot_id"] for row in rows}
|
||||||
|
|
||||||
|
|
||||||
|
async def removed_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||||
|
rows = await pool.fetch(
|
||||||
|
"SELECT bot_id FROM watched_bots WHERE NOT enabled AND status = $1",
|
||||||
|
STATUS_REMOVED,
|
||||||
|
)
|
||||||
|
return {row["bot_id"] for row in rows}
|
||||||
|
|
||||||
|
|
||||||
|
async def known_bot_ids(pool: asyncpg.Pool) -> set[int]:
|
||||||
|
rows = await pool.fetch("SELECT bot_id FROM watched_bots")
|
||||||
|
return {row["bot_id"] for row in rows}
|
||||||
|
|
||||||
|
|
||||||
|
async def register_bot(pool: asyncpg.Pool, bot_id: int) -> None:
|
||||||
|
await pool.execute(_REGISTER, bot_id)
|
||||||
|
|
||||||
|
|
||||||
|
async def remove_bot(pool: asyncpg.Pool, bot_id: int) -> bool:
|
||||||
|
return await pool.execute(_REMOVE, bot_id) != "UPDATE 0"
|
||||||
|
|
||||||
|
|
||||||
|
async def link_account(pool: asyncpg.Pool, bot_id: int, account_id: int) -> None:
|
||||||
|
await pool.execute(
|
||||||
|
"UPDATE watched_bots SET account_id = $2, updated_at = now() WHERE bot_id = $1",
|
||||||
|
bot_id,
|
||||||
|
account_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def set_status( # noqa: PLR0913
|
||||||
|
pool: asyncpg.Pool,
|
||||||
|
bot_id: int,
|
||||||
|
status: str,
|
||||||
|
*,
|
||||||
|
detail: str | None = None,
|
||||||
|
username: str | None = None,
|
||||||
|
title: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
await pool.execute(_SET_STATUS, bot_id, status, detail, username, title)
|
||||||
|
|
||||||
|
|
||||||
|
async def needs_backfill(pool: asyncpg.Pool, bot_id: int) -> bool:
|
||||||
|
return not await pool.fetchval(
|
||||||
|
"SELECT backfilled_at IS NOT NULL FROM watched_bots WHERE bot_id = $1", bot_id
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def backlog_cursor(
|
||||||
|
pool: asyncpg.Pool, bot_id: int
|
||||||
|
) -> tuple[int, int, int] | None:
|
||||||
|
row = await pool.fetchrow(_BACKLOG_CURSOR, bot_id)
|
||||||
|
if row is None:
|
||||||
|
return None
|
||||||
|
return row["backfill_pts"], row["backfill_qts"], row["backfill_date"]
|
||||||
|
|
||||||
|
|
||||||
|
async def save_backlog_cursor(
|
||||||
|
pool: asyncpg.Pool, bot_id: int, pts: int, qts: int, date: int
|
||||||
|
) -> None:
|
||||||
|
await pool.execute(_SAVE_BACKLOG_CURSOR, bot_id, pts, qts, date)
|
||||||
|
|
||||||
|
|
||||||
|
async def mark_backfilled(pool: asyncpg.Pool, bot_id: int) -> None:
|
||||||
|
await pool.execute(_MARK_BACKFILLED, bot_id)
|
||||||
Loaded 100 of 287 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user