Compare commits

..
101 Commits
Author SHA1 Message Date
YoursFunny 65a9554173 bump version to 1.0.7 2026-08-06 18:49:08 +08:00
YoursFunny e755785147 deploy: use named volumes for certs/html/acme, keep data bind mount
certs, html and acme are shared across nginx-proxy, acme-companion and
the acme-ip service, so they become named volumes (dropping the nginx-
prefix). The vhost.d mount is removed: acme-companion only needs it with
ACME_HTTP_CHALLENGE_LOCATION=true or letsencrypt_user_data, and users
who customize nginx config can mount their own. ./data stays a bind
mount so the SQLite state stays directly backup-able. README updated.
2026-08-06 17:22:05 +08:00
YoursFunny 5d5b6d56e7 chore: strip compose comments, document env vars in README
compose files now carry a single comment: the commented ACME_HOST line
(uncomment to enable domain cert issuance). Everything else moved to the
README env table and the deployment sections, which explain ACME_HOST,
DEFAULT_HOST and every other variable.
2026-08-06 17:07:50 +08:00
YoursFunny d0fdf1c5da deploy: default to ACME-issued certs, drop manual cert setup
Domain deployment now uses ACME_HOST (acme-companion auto issue/renew)
instead of the legacy LETSENCRYPT_HOST name. IP-only deployments can use
acme.sh to obtain Let's Encrypt IP certificates (shortlived profile,
~7-day validity, http-01 only) instead of self-signed manual certs:
documented the acme-ip service, first-issue command with install-cert
into nginx-certs and nginx reload via docker socket.

Removed the manual WEBHOOK_CERT flow and the ./cert volume from the
compose example.
2026-08-06 16:48:52 +08:00
YoursFunny 1db4ecfafa chore: use modern nginx-proxy label for acme-companion
acme-companion prefers com.github.nginx-proxy.nginx over the legacy
jrcs.letsencrypt_nginx_proxy_companion namespace (functions.sh checks
the new label first, falls back to the old one).
2026-08-06 11:31:01 +08:00
YoursFunny 47039bbe2d fix: make docker entrypoint restart-safe
docker compose restart reuses the container, so the overlay fs keeps the
user created on first boot; a second useradd fails with exit code 9 and
set -e kills the container on every restart (crash loop, bot never comes
back). Create the user only if missing, align UID otherwise, and never
let chown failure on bind-mounted volumes kill the container.

Also comment out nginx-proxy's empty environment block in the compose
example: an empty mapping fails validation on newer compose versions
(must be a mapping).
2026-08-06 11:10:52 +08:00
YoursFunny bc954e6e0b fix: handle oversized photos Telegram rejects with pure Rust processing 2026-08-06 10:17:29 +08:00
YoursFunny f7cb809e5a bump version to 1.0.6 2026-08-05 02:07:34 +08:00
YoursFunny 51b40cdb42 feat: cache sent media file ids for instant repeat sends 2026-08-05 02:06:30 +08:00
YoursFunny ab2306002a perf: handle batch-forwarded URLs concurrently with queue workers 2026-08-05 01:26:24 +08:00
YoursFunny a92b12f633 bump version to 1.0.5 2026-08-05 00:16:45 +08:00
YoursFunny 74e3b7593c docs: note fresh twitter query ids and TID scope in auth fallback 2026-08-05 00:07:30 +08:00
YoursFunny 2faccaac42 fix: cut tweet text by code points, not UTF-16 units 2026-08-04 23:45:54 +08:00
YoursFunny f4e60d8946 fix: trim CRLF from TWITTER_AUTH_TOKEN 2026-08-04 23:45:54 +08:00
YoursFunny 3006dcd98c feat: fetch NSFW tweets via authenticated twitter API fallback 2026-08-04 23:45:53 +08:00
YoursFunny 8e8acdd859 deploy: add container names and startup order to compose 2026-08-04 23:11:04 +08:00
YoursFunny 8b9dd963e1 bump version to 1.0.4 2026-08-04 22:10:38 +08:00
YoursFunny e83d48f1f7 fix: handle SIGTERM for graceful shutdown on docker stop 2026-08-04 22:09:45 +08:00
YoursFunny 7c55b26731 deploy: add nginx-proxy reverse proxy for webhook TLS 2026-08-04 21:50:19 +08:00
YoursFunny 950db48a13 ci: cache buildkit layers across runs 2026-08-04 18:52:28 +08:00
YoursFunny 32ea8ec6ca fix docker build: fetch ffmpeg from martin-riedl.de 2026-08-04 18:52:27 +08:00
YoursFunny 62dc033452 docs: add AGENTS.md with repository guidelines 2026-08-04 18:52:21 +08:00
YoursFunny 2801aaa39c bump version to 1.0.3 2026-08-04 16:24:16 +08:00
YoursFunny 93d47752cf fallback to smaller media when file too large 2026-08-04 16:22:49 +08:00
YoursFunny 6fd4edb3c5 webhook: drop duplicate set_webhook, ignore empty env vars 2026-08-04 15:58:32 +08:00
YoursFunny 7c5afce0b4 ci: build once on tagged commits, bump docker actions 2026-08-04 01:58:19 +08:00
YoursFunny caf41183a4 bump version to 1.0.2 2026-08-04 01:29:50 +08:00
YoursFunny b6ae88e869 twitter: request original image via name=orig 2026-08-04 01:28:21 +08:00
YoursFunny 3dff9a2473 fix docker build: compile real sources, lf entrypoint 2026-08-04 00:55:35 +08:00
YoursFunny 41f25052fe merge refactor-rs into master 2026-08-03 23:34:45 +08:00
YoursFunny fa95091c7c bump version to 1.0.0 2026-08-03 23:34:01 +08:00
YoursFunny 84ab146069 finish rust rewrite, add docker, drop python 2026-08-03 23:30:42 +08:00
YoursFunny e14363dd1d fix changes of updated async-pixiv module 2026-06-24 00:17:15 +08:00
YoursFunny 0cffd5c297 revert to python 3.12 due to incompatibility of async-pixiv module 2026-06-23 23:46:27 +08:00
YoursFunny b83bed5423 fix pixiv exception import 2026-06-23 23:24:39 +08:00
YoursFunny 0f3131aec7 bump deps version 2026-06-23 23:16:57 +08:00
YoursFunny 5483855de8 fix webhook set as false in compose yaml 2025-12-18 14:59:12 +08:00
YoursFunny 052163084b change to x.com in url template 2025-12-17 17:24:22 +08:00
YoursFunny cdee105c60 bump deps version 2025-12-17 17:23:20 +08:00
YoursFunny fae3dad43a fix start command handler 2025-08-21 10:30:10 +08:00
YoursFunny f4c53719d8 bump dep version and fix start command handler 2025-08-21 10:23:25 +08:00
YoursFunny cf44c3ddd7 minor changes and adding /start command 2025-08-21 09:59:37 +08:00
YoursFunny cbb15e6a33 fix tag check 2024-11-12 17:35:11 +08:00
YoursFunny afc31430e7 new bsky sensitive contents 2024-11-12 17:23:46 +08:00
YoursFunny 282cf46d58 minor fix 2024-11-12 17:12:47 +08:00
YoursFunny 6c53a1429d update to python 3.13 2024-11-10 19:02:18 +08:00
YoursFunny ca9064a1b7 bump version 2024-11-10 18:19:29 +08:00
YoursFunny d58fd7cd46 add clear edit message command 2024-11-10 18:17:51 +08:00
YoursFunny 9aa7aa1c36 fix external bsky media (disable) 2024-11-10 18:17:15 +08:00
YoursFunny 9c2606160a format 2024-10-19 22:41:41 +08:00
YoursFunny 3fe8a1cf3b new bsky support 2024-10-19 22:25:07 +08:00
YoursFunny 62a2820af3 add params support in fetch_json 2024-10-19 22:23:51 +08:00
YoursFunny 80b630d28d minor fixes 2024-10-19 22:22:27 +08:00
YoursFunny 4173972407 fix empty env bot_admin 2024-10-19 22:21:56 +08:00
YoursFunny 0d2648162c bump version 2024-10-19 22:19:44 +08:00
YoursFunny 0847ada3b7 update gitignore 2024-10-18 00:21:19 +08:00
YoursFunny 7cad25125f rust refactor 2024-10-18 00:17:51 +08:00
YoursFunny f0c287a148 fix await 2024-08-15 17:45:14 +08:00
YoursFunny 3658bfd2aa fix handler 2024-08-15 17:42:19 +08:00
YoursFunny ae1b04dee8 add extract url from message 2024-08-15 17:34:40 +08:00
YoursFunny 641f7218a4 minor fixes 2024-08-15 16:45:36 +08:00
YoursFunny 61dfa16011 fix edit message str format 2024-08-15 16:45:06 +08:00
YoursFunny 2986e80076 fix edit message 2024-08-14 01:52:01 +08:00
YoursFunny f112ecdf63 fix text 2024-08-14 01:30:08 +08:00
YoursFunny 79d50ae9a3 fix 2024-08-14 01:19:50 +08:00
YoursFunny 416331303d fix 2024-08-14 01:14:01 +08:00
YoursFunny 55771e01fe fix 2024-08-14 01:13:47 +08:00
YoursFunny afca276f49 minor fixes 2024-08-14 00:57:20 +08:00
YoursFunny 06d25854d2 fix error import 2024-08-14 00:37:01 +08:00
YoursFunny 515aa9711d new set template 2024-08-14 00:29:02 +08:00
YoursFunny 733adfc4a6 refactor using custom callback context 2024-08-13 23:20:24 +08:00
YoursFunny 03066853de refactor 2024-08-13 19:17:58 +08:00
YoursFunny 566c17a855 fix pixiv url 2024-07-19 19:34:18 +08:00
YoursFunny 87f0d16028 fix regex https pattern 2024-07-16 14:28:49 +08:00
YoursFunny 197993e522 bump version 2024-07-16 14:27:15 +08:00
YoursFunny 02abcd899e update pixiv regex 2024-07-16 14:21:43 +08:00
YoursFunny ebb48d2bb5 fix forward channel check admin user 2024-07-09 16:27:13 +08:00
YoursFunny be77a33e9d use large instead of origin for pixiv 2024-06-24 20:02:35 +08:00
YoursFunny 0439adefa0 add refresh token if expire 2024-06-24 15:57:30 +08:00
YoursFunny 81e5d2dbd1 add log 2024-06-24 15:47:12 +08:00
YoursFunny e7ccd70e56 fix deal with invalid url 2024-06-24 02:17:26 +08:00
YoursFunny 23128af96a fix regex 2024-06-24 02:14:50 +08:00
YoursFunny fe1724aca4 fix annotations 2024-06-24 02:06:25 +08:00
YoursFunny 0433650563 fix init pixiv 2024-06-24 02:03:28 +08:00
YoursFunny a617a7e1a0 fix requirement conflict 2024-06-24 01:48:45 +08:00
YoursFunny e35e3895ba fix requirement conflict 2024-06-24 01:45:23 +08:00
YoursFunny a4a9f2eaa0 add pixiv package requirement 2024-06-24 01:29:29 +08:00
YoursFunny 1c779e98a8 add pixiv support 2024-06-24 01:25:03 +08:00
YoursFunny 4a5aa296ce add pixiv refresh token setting 2024-06-23 22:14:58 +08:00
YoursFunny c1c1d42892 refactor, clear structure 2024-06-23 22:06:27 +08:00
YoursFunny cade370538 refactor 2024-06-23 21:52:12 +08:00
YoursFunny a6325d5bc8 fix check instance 2024-06-20 15:46:47 +08:00
YoursFunny 6c0e581135 fix regex match group 2024-06-20 15:40:18 +08:00
YoursFunny 494fc50446 fix missing aexit 2024-06-20 15:29:30 +08:00
YoursFunny a4e568667c fix type import error 2024-06-20 15:25:17 +08:00
YoursFunny 8231eeb20e refactor tweet, better type hint 2024-06-20 15:13:09 +08:00
YoursFunny 74f42a2df3 fix status code check 2024-06-20 02:51:20 +08:00
YoursFunny ae1f5a8c8c rename 2024-06-20 02:48:06 +08:00
YoursFunny 38adbf0cb8 use cached property 2024-06-20 02:40:35 +08:00
YoursFunny 55ae73ee6b use uuid for query 2024-06-20 02:21:14 +08:00
YoursFunny d6589fec5b fix annotations 2024-06-20 02:20:10 +08:00
40 changed files with 10444 additions and 599 deletions
+6 -1
View File
@@ -24,8 +24,13 @@
**/secrets.dev.yaml
**/values.dev.yaml
*.db
.python-version
LICENSE
README.md
data/
cert/
nginx-certs/
nginx-vhost.d/
nginx-html/
nginx-acme/
**/target/
.idea/
+2
View File
@@ -0,0 +1,2 @@
# Shell scripts must stay LF: CRLF breaks the shebang inside containers.
*.sh text eol=lf
+36 -5
View File
@@ -12,12 +12,35 @@ env:
DOCKERHUB_REPO: yoursfunny/telegram-twitter-media-bot
jobs:
# A tag push and a branch push to the same commit fire two workflow runs;
# build only once. Tag runs always build; master runs build only when the
# pushed commit is not already tagged (the tag run covers it).
should-build:
runs-on: ubuntu-latest
outputs:
build: ${{ steps.check.outputs.build }}
steps:
- uses: actions/checkout@v7
with:
fetch-depth: 0
- id: check
shell: bash
run: |
if [ "$GITHUB_REF_TYPE" = "branch" ] && git tag --points-at "$GITHUB_SHA" | grep -q .; then
echo "commit already tagged; the tag run builds the image"
echo "build=false" >> "$GITHUB_OUTPUT"
else
echo "build=true" >> "$GITHUB_OUTPUT"
fi
docker:
needs: should-build
if: needs.should-build.outputs.build == 'true'
runs-on: ubuntu-latest
steps:
- name: Docker meta
id: meta
uses: docker/metadata-action@v5
uses: docker/metadata-action@v6
with:
images: ${{ env.DOCKERHUB_REPO }}
tags: |
@@ -28,22 +51,30 @@ jobs:
type=sha
-
name: Set up QEMU
uses: docker/setup-qemu-action@v3
uses: docker/setup-qemu-action@v4
-
name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
uses: docker/setup-buildx-action@v4
-
name: Login to Docker Hub
uses: docker/login-action@v3
uses: docker/login-action@v4
with:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_TOKEN }}
# Buildkit cache via the GitHub Actions cache backend (uses the
# automatic GITHUB_TOKEN, no extra secrets). mode=max keeps every
# stage's layers so the cargo-deps and ffmpeg layers are restored
# instead of re-downloaded/recompiled. The scope must be pinned to a
# fixed string: the gha backend defaults to the current git ref, which
# would give every new tag a cold cache on release builds.
-
name: Build and push
uses: docker/build-push-action@v5
uses: docker/build-push-action@v7
with:
push: true
build-args: |
APP_NAME=${{ env.APP_NAME }}
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha,scope=tgxmb-build
cache-to: type=gha,mode=max,scope=tgxmb-build
+11 -1
View File
@@ -2,5 +2,15 @@
__pycache__/
cert/
data/
nginx-certs/
nginx-vhost.d/
nginx-html/
nginx-acme/
docker-compose.yml
x.py
.env
# Added by cargo
/target
+99
View File
@@ -0,0 +1,99 @@
# Repository Guidelines
## Project Overview
Telegram bot (teloxide) that turns post links from X/Twitter, Pixiv, and Bluesky into media messages (images, video, GIF) with the post's title, author, and tags. It supports batch media splitting, retry with persistence, inline queries, forward-channel rebinding with caption templates, and Pixiv ugoira→MP4 transcoding. README and user-facing strings are in Chinese. The project is a Rust port of a Python predecessor (see `queue.rs` comments referencing `utils/task_queue.py`).
Two-crate Cargo workspace (both v1.0.3, edition 2024, resolver 3):
- **`crates/x-media`** — library that fetches and normalizes media from the three sites. Pure, no Telegram knowledge.
- **`crates/xmedia-bot`** — the bot binary: teloxide dispatcher, SQLite-backed chat state, persistent task queue.
## Architecture & Data Flow
```
Telegram update → Dispatcher (polling or axum webhook) → dptree branches
├─ message → commands (any chat) / URL links (private chat only)
├─ inline_query → InlineQueryResult Photo/Video/Mpeg4Gif
└─ callback_query → "forward" (copy to channel) / "template|<name>" (apply caption template)
```
Message flow: `message_handler` extracts URLs (from `url`/`text_link` entities, text + caption, deduped) → `x_media::site::fetch(url)``Fetched` → builds a `Task``send::send_media_sequence` (media groups ≤ 9, caption on first item) or `send::send_animation`. On Telegram URL-fetch failure or size error (`send_batch_via_upload`): download via `x_media::site::download_media` to a temp file (≤ 10 MiB), sniff magic bytes (`sniff_ext`), upload via multipart; oversized items fall back to `fallback_url`. On failure: `enqueue_retry` persists resume-state `Task` into the SQLite queue → single worker leases (120 s lock TTL) → retry with exponential backoff (≤ 30 s, `MAX_RETRIES = 2`) → dead-letter → `notify_failure`. Success → `post_send_actions`: edit-before-forward prompt with inline buttons, or `copy_messages` to the bound forward channel.
The `x-media` library: `site::fetch(url)` dispatches (in order) twitter → bsky → pixiv via per-site regex `PATTERN` and returns `Ok(None)` for unmatched URLs. `Fetched { source_url, caption, title, media: Vec<Media>, sensitive, … }`; `caption_with(format)` substitutes `{url} {author} {author_url} {title} {tags}`.
## Key Directories
| Path | Purpose |
|---|---|
| `crates/x-media/src/` | Fetch library. `site/mod.rs` = dispatcher + `Fetched`/`FetchError`/`download_media`/`media_size`; `media.rs` = `Media` enum; `examples/fetch.rs` = end-to-end usage sample |
| `crates/x-media/src/site/<twitter\|pixiv\|bsky>/` | One directory per site: `mod.rs` (re-exports), `interface.rs` (PATTERN, `enabled()`, `fetch_from_url()`, site struct, `From<SiteStruct> for Fetched`), `model.rs` (serde DTOs). Pixiv adds `api.rs` (auth + transport); twitter adds `auth.rs` (logged-in GraphQL `TweetDetail` fallback for NSFW tweets, gated on `TWITTER_AUTH_TOKEN`) |
| `crates/xmedia-bot/src/main.rs` | Entry point: env/log init, queue worker start, pixiv validation, 300 s edit-expiry sweep, dptree handler tree, webhook vs polling dispatch |
| `crates/xmedia-bot/src/config.rs` | Manual env parsing into `Config` |
| `crates/xmedia-bot/src/handlers.rs` | `Command` enum (teloxide `BotCommands`), message/inline/callback handlers, URL extraction, global statics; per-URL work spawned with a `Semaphore(8)` cap (teloxide's per-chat workers are sequential — batch-forwards need concurrency) |
| `crates/xmedia-bot/src/state.rs` | `ChatStore`: parking_lot `Mutex<HashMap>` cache + SQLite write-through (`chat_state` table) |
| `crates/xmedia-bot/src/link_cache.rs` | `LinkCache`: SQLite-backed cache (`link_cache` table) of successfully sent posts — raw caption fields + Telegram `file_id`s; repeat links re-send locally (no fetch/upload), TTL + prune, invalidated on permanent send failure |
| `crates/xmedia-bot/src/queue.rs` | `PersistentTaskQueue`: SQLite-backed queue (`tasks` table), `QUEUE_WORKERS = 4` concurrent workers (lease via `BEGIN IMMEDIATE` + `locked_until` TTL), retry→dead-letter, `Notify::notify_waiters` wakeup, `busy_timeout` on all connections |
| `crates/xmedia-bot/src/send.rs` | Media senders, upload fallback, error classification, queue task handlers |
## Development Commands
```bash
export TELOXIDE_TOKEN=<token> # required; PIXIV_REFRESH_TOKEN optional (Pixiv disabled without it)
cargo run -p xmedia-bot # run the bot (polling by default)
cargo run -p x-media --example fetch -- <url> # test a link through the fetch library
cargo test --workspace # full test suite (no CI test step exists — run locally)
cargo build --release -p xmedia-bot # release build (Dockerfile does this)
cargo clippy --workspace --all-targets # lint (Clippy is the configured IDE linter)
cargo fmt --check # formatting
```
Docker: `docker build -t tgxmb .` then `docker run --rm -d --name tgxmb --env-file .env -v ./data:/app/data tgxmb`. Runtime requires **ffmpeg** (built into the image).
## Code Conventions & Common Patterns
- **No anyhow/thiserror.** Errors are hand-rolled enums with manual `Display`/`source()`/`From` impls: `QueueError` (`Retryable { delay_seconds, payload }` / `Permanent`), `SendError` (Retryable/Permanent), `FetchError` (`Http`/`Json`/`Pixiv`/`NotFound`/`Blocked`), `PixivError`, `Classification`. New errors should follow this pattern.
- **Global state via `std::sync::LazyLock` statics**, not DI: `CONFIG`, `CHAT_STORE`, `TASK_QUEUE` in `handlers.rs`; shared reqwest `CLIENT` in `x-media/src/site/mod.rs`. `Bot` is passed/cloned into handlers; queue workers rebuild `Bot::from_env()`.
- **Async**: tokio multi-thread runtime (`#[tokio::main]` default). All rusqlite I/O inside `tokio::task::spawn_blocking`. Long loops use `tokio::select!` with `tokio::sync::{watch, Notify}` stop/wake channels. No streams.
- **Blocking sync primitives**: `parking_lot::Mutex` for hot caches, `tokio::sync::Mutex` for async-shared state (pixiv token cache), `AtomicBool` for feature gates.
- **Site adapter convention** (no trait, no enum dispatch — follow the existing convention): each site module exports `PATTERN: LazyLock<Regex>`, `enabled() -> bool`, `fetch_from_url(url) -> Result<Fetched, FetchError>`; `site/mod.rs` re-exports the site struct and `fetch_once` adds one guarded if-branch. Adding a site = new `site/<name>/{mod.rs,interface.rs,model.rs}` + one branch in `fetch_once`.
- **Serde**: per-site `model.rs` are pure `Deserialize` DTOs mirroring API JSON; site structs in `interface.rs` have private fields, a `caption()` builder, and `impl From<SiteStruct> for Fetched`. Persisted payloads use internally-tagged enums (`#[serde(tag = "kind")]` / `type`).
- **Naming**: module-per-concern, snake_case files, `CamelCase` types, `snake_case` fns. `//!` module docs and `///` docs on non-obvious logic (syndication token, ugoira encoding, `display_text_range`).
- **Retries**: only `x-media::site::fetch` retries (3 attempts, `1 << attempt` backoff, HTTP errors only). Queue retries are explicit `QueueError::Retryable` with computed delay (`retry_delay_seconds`).
- Logging via `log` macros (`pretty_env_logger`, level from `RUST_LOG`).
## Important Files
| File | Why it matters |
|---|---|
| `crates/xmedia-bot/src/main.rs` | Startup sequence, webhook vs polling, graceful shutdown (SIGINT via teloxide ctrlc / SIGTERM via `stop_token` for docker, → sweep stop → admin msg → queue stop) |
| `crates/xmedia-bot/src/handlers.rs` | `CHAT_STORE`/`TASK_QUEUE`/`CONFIG` singletons (open `data/task_queue.db` **relative to CWD**); command dispatch; URL extraction; retry enqueue |
| `crates/xmedia-bot/src/send.rs` | Constants `MAX_MEDIA_GROUP = 9`; fallback chain; `classify_request_error`; download-and-reupload fallback triggered only by Telegram API errors (`is_media_fetch_failure` / `is_size_error`) |
| `crates/xmedia-bot/src/photo.rs` | Pure-Rust photo processing (no ffmpeg): `png` (image-png) decode/encode + `zune-jpeg` decode + `fast_image_resize` Lanczos3 downscale + `jpeg-encoder`. Photos over Telegram's limits (width + height > 10000 px → `PHOTO_INVALID_DIMENSIONS`; bytes > 10 MiB) are decoded, downscaled keeping the format, PNG bit depth > 24 (RGBA 32-bit / 16-bit per channel) reduced to 24-bit RGB with alpha flattened white (≤24-bit untouched, never upconverted), and transcoded to JPEG only if still over the cap; memory budget guarded, otherwise the item's smaller fallback URL |
| `crates/x-media/src/site/mod.rs` | Dispatcher, `Fetched`/`FetchError`, shared `CLIENT`, `download_media` (adds `Referer: https://www.pixiv.net/` for `pximg.net` hotlink protection) |
| `crates/x-media/src/site/pixiv/api.rs` | OAuth token exchange (hardcoded app client id/secret), access-token cache, ugoira zip→MP4 via ffmpeg in `spawn_blocking` |
| `Dockerfile` | Multi-stage: cached dep layer via stub sources + `touch *.rs` mtime hack, static ffmpeg from ffmpeg.martin-riedl.de (`FFMPEG_URL` arg, `unzip -t` integrity check), `debian:bookworm-slim` runtime, entrypoint |
| `docker-entrypoint.sh` | Privilege drop: `useradd` with `LOCAL_USER_ID` (default 9001) + `setpriv` (no gosu on bookworm-slim) |
| `docker-compose.yml.example` | Deployment env reference (real `docker-compose.yml` is gitignored). Ships nginx-proxy + acme-companion: webhook mode needs TLS termination in front (teloxide's axum listener is HTTP-only; `WEBHOOK_CERT` only feeds `set_webhook`), bot exposes `VIRTUAL_HOST`/`VIRTUAL_PORT` on the shared `proxy` network, no host port; container names `nginx-proxy`/`acme-companion`/`tgxmb`, start order via `depends_on` (proxy → acme → bot) |
| `.github/workflows/docker.yml` | CI: build+push to Docker Hub on tag `v*`/master; **no test step**; buildx gha cache (`cache-from`/`cache-to`, scope `tgxmb-build`, `mode=max`) so cargo deps + ffmpeg layers are restored across runs |
| `README.md` | Feature docs + command table (Chinese) |
## Runtime/Tooling Preferences
- **Rust, stable, edition 2024**, workspace resolver 3. No `rust-version`/MSRV pin, no `rust-toolchain.toml` — recent stable is assumed. No nightly features.
- Package manager: **Cargo** (workspace with path dep `x-media``xmedia-bot`). No `[workspace.package]`/shared deps — each crate lists deps independently.
- **Two reqwest versions coexist in the lock** (0.12.28 via teloxide, 0.13.3 in x-media) — don't unify casually.
- Config is **environment-variable driven** (dotenv loads `.env`, gitignored; no `.env.example` exists). Key vars: `TELOXIDE_TOKEN` (required), `PIXIV_REFRESH_TOKEN`, `TWITTER_AUTH_TOKEN` (optional; x.com `auth_token` cookie — enables the logged-in GraphQL fallback that fetches NSFW tweets syndication withholds), `BOT_ADMIN` (comma-separated ids), `EDIT_MESSAGE_TTL_SECONDS` (default 86400), `LINK_CACHE_TTL_SECONDS` (default 604800), `WEBHOOK`/`WEBHOOK_URL`/`WEBHOOK_LISTEN`/`WEBHOOK_PORT`/`WEBHOOK_CERT`/`WEBHOOK_SECRET_TOKEN` (webhook mode requires URL/listen/port, `.expect`ed; `WEBHOOK_CERT` is Telegram-facing self-signed validation only — TLS must be terminated by a reverse proxy), `RUST_LOG`, `TELOXIDE_PROXY`, `LOCAL_USER_ID` (entrypoint only).
- SQLite via `rusqlite` with `bundled` feature (no system libsqlite needed). DB file `data/task_queue.db` is CWD-relative — run from the workspace root, or `/app` in Docker. Mount `./data` and `./cert` volumes.
- `.gitattributes` enforces LF for `*.sh` (CRLF breaks shebangs in containers). `.gitignore`: `.env`, `data/`, `cert/`, `docker-compose.yml`, `/target`, `.idea/`.
- Docs are in Chinese; user-facing bot strings too. Keep that convention when editing captions/templates/docs.
## Testing & QA
- **~51 tests, all inline `#[cfg(test)] mod tests`** — no `tests/` integration directories. Framework: built-in Rust test + `#[tokio::test]` (dev-deps only in `x-media`: tokio macros/rt-multi-thread, dotenv).
- No mocking framework anywhere (no mockito/wiremock/mockall). Conventions: pure-function units (regex parsing, serde round-trips, chunking, retry math) tested synchronously; async tests use real dependencies — file-backed SQLite via `tempfile` (`queue.rs::new_queue()` helper), live network fetches.
- Live-network tests exist in `site/twitter/interface.rs` (3), `site/bsky/interface.rs` (2), `site/pixiv/interface.rs`/`api.rs` (env-gated on `PIXIV_REFRESH_TOKEN`/dotenv, skip by early return). Run the full suite with `cargo test --workspace`.
- Fixtures are inline `serde_json::json!` builder fns (`fixture()`, `thread_json()`, `illust_json()`), not files. The shared `CLIENT` sets `pool_max_idle_per_host(0)` under `#[cfg(test)]` to avoid cross-runtime `DispatchGone`.
- **CI runs no tests** — `.github/workflows/docker.yml` only builds/pushes the image; verification is a local responsibility.
- Untested and hard to test without a mock seam: `handlers.rs` (depends directly on teloxide `Bot`); `main.rs`, `config.rs`, `state.rs`; `media.rs`, `lib.rs`, all `model.rs`.
- No coverage tracking, no lint gate in CI.
Generated
+3611
View File
File diff suppressed because it is too large Load Diff
+3
View File
@@ -0,0 +1,3 @@
[workspace]
members = ["crates/x-media", "crates/xmedia-bot"]
resolver = "3"
+66 -19
View File
@@ -1,26 +1,73 @@
FROM python:3.12-slim-bullseye
# ---------- build stage ----------
# rust:1-bookworm (full, not slim) ships the C toolchain needed by
# rusqlite's bundled SQLite, plus wget/unzip for the ffmpeg download.
FROM rust:1-bookworm AS builder
ARG APP_NAME=telegram-twitter-media-bot
# Prebuilt static ffmpeg (glibc-linked, includes libx264) for ugoira MP4
# encoding. Served from https://ffmpeg.martin-riedl.de (Cloudflare CDN,
# built on Debian 12 — glibc-compatible with the bookworm-slim runtime).
# johnvansickle.com throttles datacenter IPs and served garbage from GitHub
# runners. `/redirect/latest/` floats to the newest release build; each build
# also ships a .sha256. Swap `amd64` for `arm64` when building arm64 images.
ARG FFMPEG_URL=https://ffmpeg.martin-riedl.de/redirect/latest/linux/amd64/release/ffmpeg.zip
WORKDIR /build
# 1. Rust dependencies first: only the manifests plus stub sources, so the
# expensive dependency fetch + compile lives in a layer invalidated only by
# manifest/lock changes.
COPY Cargo.toml Cargo.lock ./
COPY crates/x-media/Cargo.toml crates/x-media/Cargo.toml
COPY crates/xmedia-bot/Cargo.toml crates/xmedia-bot/Cargo.toml
RUN mkdir -p crates/x-media/src crates/xmedia-bot/src \
&& printf 'fn main() {}\n' > crates/xmedia-bot/src/main.rs \
&& : > crates/x-media/src/lib.rs \
&& cargo build --release -p xmedia-bot
# 2. Static ffmpeg next (cached unless FFMPEG_URL changes), so source edits
# never re-download it. The zip contains a single `ffmpeg` binary at the
# root. `unzip -t` verifies the archive before extraction so a bad
# download fails loudly here instead of a cryptic later error.
RUN wget -q -O /tmp/ffmpeg.zip "$FFMPEG_URL" \
&& unzip -tq /tmp/ffmpeg.zip \
&& unzip -q /tmp/ffmpeg.zip -d /usr/local/bin \
&& chmod +x /usr/local/bin/ffmpeg \
&& rm /tmp/ffmpeg.zip \
&& /usr/local/bin/ffmpeg -version >/dev/null
# 3. Real sources last: only our crates recompile on source changes. The
# COPY preserves host mtimes, which predate the stub artifacts from step 1;
# cargo's mtime-based freshness check would otherwise treat the stub build
# as up-to-date and never compile the real sources. `touch` forces cargo to
# see the real files as newer.
COPY crates/ ./crates/
RUN find crates -type f -name '*.rs' -exec touch {} + \
&& cargo build --release -p xmedia-bot
# ---------- runtime stage ----------
FROM debian:bookworm-slim
# ARG scope is per-stage: re-declare for the label below.
ARG APP_NAME=telegram-twitter-media-bot
LABEL maintainer="admin@yoursfunny.top"
LABEL org.opencontainers.image.title="${APP_NAME}"
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
RUN set -eux; \
apt-get update; \
apt-get install -y gosu; \
rm -rf /var/lib/apt/lists/*; \
# verify that the binary works
gosu nobody true
# Everything is copied in — no apt in the runtime stage. Privilege dropping is
# done by docker-entrypoint.sh with setpriv (util-linux, already in
# bookworm-slim), so no gosu needed.
COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
COPY --from=builder /usr/lib/x86_64-linux-gnu/libssl.so.3* /usr/lib/x86_64-linux-gnu/
COPY --from=builder /usr/lib/x86_64-linux-gnu/libcrypto.so.3* /usr/lib/x86_64-linux-gnu/
COPY --from=builder /usr/local/bin/ffmpeg /usr/local/bin/ffmpeg
WORKDIR /app
COPY --from=builder /build/target/release/xmedia-bot /usr/local/bin/xmedia-bot
COPY docker-entrypoint.sh /app/docker-entrypoint.sh
RUN chmod a+x /app/docker-entrypoint.sh
COPY requirements.txt /app
RUN python -m pip install --no-cache-dir --upgrade -r requirements.txt
COPY . /app
RUN chmod a+x docker-entrypoint.sh
# State lives in /app/data (SQLite task queue + chat state); mount a volume
# there to keep it across restarts.
ENTRYPOINT ["/app/docker-entrypoint.sh"]
CMD ["python", "main.py"]
CMD ["xmedia-bot"]
+117
View File
@@ -0,0 +1,117 @@
# TelegramXMediaBot
Telegram 机器人,将 X / Twitter、Pixiv、Bluesky 的帖子链接转换为媒体消息发送,附带帖子标题、作者与标签。
## 功能
- 私聊发送链接后自动抓取并发送图片、视频与 GIF,超量图片自动分批
- 纯文字帖提示无媒体;不支持的链接静默忽略
- 支持内联查询(`@机器人 <链接>`
- 可绑定转发频道自动转发;支持转发前编辑 caption 与自定义模板
- 发送失败自动重试并持久化,重试耗尽后通知用户
- Pixiv ugoira 动图自动转码为 MP4
- 链接结果本地缓存:成功发送后缓存 Telegram file id 与 caption 等,再次收到相同链接直接本地重发,不再请求源站、不保存媒体文件(`LINK_CACHE_TTL_SECONDS` 控制过期,默认 7 天)
## 快速开始
```bash
# 必填:BotFather 的 token;可选:PIXIV_REFRESH_TOKEN(未设置则禁用 Pixiv
export TELOXIDE_TOKEN=<token>
export PIXIV_REFRESH_TOKEN=<token>
cargo run -p xmedia-bot
```
Docker 部署(参考 `docker-compose.yml.example`):
```bash
docker build -t tgxmb .
docker run --rm -d --name tgxmb --env-file .env -v ./data:/app/data tgxmb
```
环境变量:`TELOXIDE_TOKEN`(必填)、`PIXIV_REFRESH_TOKEN``BOT_ADMIN``EDIT_MESSAGE_TTL_SECONDS``LINK_CACHE_TTL_SECONDS``RUST_LOG``WEBHOOK*``TWITTER_AUTH_TOKEN`(可选)。
NSFW 推文:公开的 syndication 接口不返回敏感内容。设置 `TWITTER_AUTH_TOKEN`(登录 x.com 后浏览器 Cookie 里的 `auth_token` 值)后,bot 会仅在遇到 NSFW 推文时以登录态获取媒体;未设置则提示无媒体。
### Webhook 部署(需要反向代理)
`docker-compose.yml.example` 内置了 [nginx-proxy](https://github.com/nginx-proxy/nginx-proxy) + [acme-companion](https://github.com/nginx-proxy/acme-companion) 反向代理编排,按部署环境二选一:
**有域名**
1. DNS A 记录指向服务器
2. compose 里设 `VIRTUAL_HOST``WEBHOOK_URL` 为域名,并取消注释 `ACME_HOST`(设为域名)
3. acme-companion 自动签发与续期证书,无需手动处理
**只有 IP**
Let's Encrypt 支持为公网 IP 签发证书(2026 年起可用,有效期约 7 天,须 `shortlived` profile)。用 [acme.sh](https://github.com/acmesh-official/acme.sh) 自动签发与续期,无需手动证书:
1. compose 里增加 acme-ip 服务(签发 + 每日检查自动续期):
```yaml
acme-ip:
image: neilpang/acme.sh
container_name: acme-ip
command: daemon
restart: always
volumes:
- certs:/acme.sh
- html:/usr/share/nginx/html
- /var/run/docker.sock:/var/run/docker.sock:ro
networks: [proxy]
```
2. 首次签发(把 `<SERVER_IP>` 换成服务器公网 IP,IPv6 同样支持,多个 `-d` 可并列):
```bash
docker compose exec acme-ip acme.sh --issue --server letsencrypt \
-d <SERVER_IP> --cert-profile shortlived --days 3 \
--webroot /usr/share/nginx/html \
--install-cert --cert-file /acme.sh/<SERVER_IP>.crt \
--key-file /acme.sh/<SERVER_IP>.key \
--reloadcmd "curl --unix-socket /var/run/docker.sock -X POST http://localhost/containers/nginx-proxy/kill?signal=HUP"
```
3. compose 里设 `VIRTUAL_HOST: '<SERVER_IP>'`、`WEBHOOK_URL: 'https://<SERVER_IP>/'`,无需 `WEBHOOK_CERT`。续期由 acme.sh daemon 自动完成(`--days 3` = 每 3 天续一次,证书 7 天有效有缓冲),续期成功后自动 HUP 通知 nginx-proxy 加载新证书。
限制:证书约 7 天有效;验证仅支持 http-01/tls-alpn-0180 端口必须公网可达);不支持 DNS-01、私有 IP 与 IP 段;同一 IP 集合每 168 小时限签发 5 张。建议先用 `--server letsencrypt_test` 试签,成功后再切正式服务器。
Telegram 只接受 443/80/88/8443 端口。
<details>
<summary>环境变量说明</summary>
| 变量 | 说明 |
|---|---|
| `TELOXIDE_TOKEN` | Bot token(必填) |
| `PIXIV_REFRESH_TOKEN` | Pixiv 刷新令牌;未设置则禁用 Pixiv |
| `BOT_ADMIN` | 管理员聊天 ID,逗号分隔;接收启动/停止通知 |
| `EDIT_MESSAGE_TTL_SECONDS` | 转发前编辑记录过期秒数,默认 86400 |
| `LINK_CACHE_TTL_SECONDS` | 链接结果缓存过期秒数,默认 604800(7 天) |
| `RUST_LOG` | 日志级别 |
| `LOCAL_USER_ID` | 容器内运行用户 UID,默认 9001 |
| `VIRTUAL_HOST` | 对外域名或 IPnginx-proxy 按此路由 |
| `VIRTUAL_PORT` | bot 容器内监听端口,nginx-proxy 的转发目标 |
| `ACME_HOST` | 域名部署:设为域名时由 acme-companion 自动签发/续期证书 |
| `DEFAULT_HOST` | nginx-proxy 将未知 Host 的请求路由到该 vhost(IP 访问时需要) |
| `DEFAULT_EMAIL` | acme-companion 证书通知邮箱 |
| `WEBHOOK` | `true` 启用 webhook 模式(默认轮询) |
| `WEBHOOK_LISTEN` / `WEBHOOK_PORT` | bot 容器内监听地址/端口 |
| `WEBHOOK_URL` | 对外公网 HTTPS 地址(`https://域名/` 或 `https://IP/` |
| `WEBHOOK_SECRET_TOKEN` | 更新校验令牌(`X-Telegram-Bot-Api-Secret-Token` |
</details>
## 命令
| 命令 | 说明 |
|---|---|
| `/set_forward_channel <频道>` | 设置转发频道 |
| `/remove_forward_channel` | 取消转发频道 |
| `/edit_before_forward` | 开关转发前编辑 |
| `/set_template <名称>` | 将回复的消息(含 `[]`)保存为模板 |
| `/set_format <站点> <格式>` | 自定义 caption 格式(占位符 `{url}` `{title}` `{tags}` 等) |
| `/bot_dict` | 查看聊天状态 |
链接处理仅限私聊;命令在任意聊天可用。
## 备注
- 数据持久化于 `data/task_queue.db`compose 部署使用 bind mount `./data`(保持目录形式便于备份)
- 运行环境需安装 ffmpeg(Docker 镜像已内置)
- 测试:`cargo test --workspace`
-36
View File
@@ -1,36 +0,0 @@
import logging
import os
import re
try:
import uvloop, asyncio
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
uvloop = None
BOT_TOKEN = os.getenv("BOT_TOKEN")
ADMIN = [int(i) for i in os.getenv("BOT_ADMIN").split(",")]
WEBHOOK = os.getenv("WEBHOOK", False)
if WEBHOOK:
WEBHOOK_LISTEN = os.getenv("WEBHOOK_LISTEN", "0.0.0.0")
WEBHOOK_PORT = int(os.getenv("WEBHOOK_PORT", 8443))
WEBHOOK_URL = os.getenv("WEBHOOK_URL")
WEBHOOK_KEY = os.getenv("WEBHOOK_KEY", "cert/private.key")
WEBHOOK_CERT = os.getenv("WEBHOOK_CERT", "cert/cert.pem")
WEBHOOK_SECRET_TOKEN = os.getenv("WEBHOOK_SECRET_TOKEN")
x_url_regex = re.compile(r"^(?:https?://)(?:www\.|mobile\.|)(?:x|twitter|fixvx|vxtwitter)\.com/(.+)/status/(\d+)")
x_media_regex = re.compile(r"^(?:https?://)(pbs|video)\.twimg\.com/(.*)")
x_tco_regex = re.compile(r"(?:https?://)t\.co/.+$", re.M)
message_url_regex = re.compile(r"\[.+]", re.S)
logging.basicConfig(
level=os.getenv("LOG_LEVEL", "WARNING"),
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
def get_logger(name: str) -> logging.Logger:
return logging.getLogger(name)
+22
View File
@@ -0,0 +1,22 @@
[package]
name = "x-media"
version = "1.0.7"
edition = "2024"
[dependencies]
reqwest = { version = "0.13", features = ["json", "query", "form"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
regex = "1.12"
html-escape = "0.2"
url = "2.5.2"
bytes = "1"
zip = "2"
tempfile = "3"
rand = "0.8"
log = "0.4"
tokio = { version = "1.40", features = ["time"] }
[dev-dependencies]
tokio = { version = "1.40", features = ["macros", "rt-multi-thread"] }
dotenv = "0.15"
+10
View File
@@ -0,0 +1,10 @@
use x_media::site;
#[tokio::main]
async fn main() {
let url = std::env::args()
.nth(1)
.expect("usage: cargo run -p x-media --example fetch -- <url>");
let result = site::fetch(&url).await;
println!("{result:#?}");
}
+2
View File
@@ -0,0 +1,2 @@
pub mod media;
pub mod site;
+55
View File
@@ -0,0 +1,55 @@
impl Media {
pub fn url(&self) -> &str {
match self {
Media::Illustration { url, .. } => url,
Media::Video { url, .. } => url,
Media::Animated { url, .. } => url,
}
}
pub fn thumbnail_url(&self) -> Option<&str> {
match self {
Media::Illustration { thumbnail_url, .. } => thumbnail_url.as_deref(),
Media::Video { thumbnail_url, .. } => Some(thumbnail_url),
Media::Animated { thumbnail_url, .. } => Some(thumbnail_url),
}
}
/// A smaller variant of this media's file (used as the fallback when the
/// primary URL or upload exceeds Telegram's size limits). None when no
/// smaller variant exists (videos, animated gifs).
pub fn smaller_url(&self) -> Option<&str> {
match self {
Media::Illustration {
url,
fallback_url,
thumbnail_url,
..
} => fallback_url
.as_deref()
.or(thumbnail_url.as_deref())
.filter(|smaller| *smaller != url),
Media::Video { .. } | Media::Animated { .. } => None,
}
}
}
#[derive(Debug)]
pub enum Media {
Illustration {
title: Option<String>,
url: String,
thumbnail_url: Option<String>,
fallback_url: Option<String>,
},
Video {
title: Option<String>,
url: String,
thumbnail_url: String,
},
Animated {
title: Option<String>,
url: String,
thumbnail_url: String,
},
}
+272
View File
@@ -0,0 +1,272 @@
use super::model;
use crate::media::Media;
use crate::site::{FetchError, Fetched};
use html_escape::encode_text;
use regex::Regex;
use std::sync::LazyLock;
pub static PATTERN: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"bsky\.app/profile/([\w.\-:]+)/post/([\w.\-~]+)").unwrap()
});
pub fn enabled() -> bool {
true
}
pub async fn fetch_from_url(url: &str) -> Result<Fetched, FetchError> {
let caps = PATTERN.captures(url).ok_or(FetchError::NotFound)?;
let handle = caps.get(1).map(|m| m.as_str()).ok_or(FetchError::NotFound)?;
let rkey = caps.get(2).map(|m| m.as_str()).ok_or(FetchError::NotFound)?;
Ok(fetch(handle, rkey).await?.into())
}
/// Fetches a post thread by handle or DID (`at://` URIs work for both).
pub async fn fetch(handle: &str, rkey: &str) -> Result<Post, FetchError> {
let response = crate::site::CLIENT
.get(API_URL)
.query(&[
("uri", format!("at://{handle}/app.bsky.feed.post/{rkey}")),
("depth", "0".to_string()),
])
.send()
.await?;
let text = response.text().await?;
Ok(Post::from_json(&text, rkey.to_string())?)
}
#[derive(Debug)]
pub struct Post {
id: String,
author: String,
author_id: String,
text: String,
media: Vec<Media>,
sensitive: bool,
}
impl Post {
fn url(&self) -> String {
format!("{}/post/{}", self.author_url(), self.id)
}
fn author_url(&self) -> String {
format!("https://bsky.app/profile/{}", self.author_id)
}
pub fn caption(&self) -> String {
format!(
"{url}\n<a href=\"{author_url}\">{author}</a>: {text}",
url = self.url(),
author_url = self.author_url(),
author = encode_text(&self.author),
text = encode_text(&self.text),
)
}
pub fn from_json(raw_json: &str, id: String) -> Result<Self, FetchError> {
let json: serde_json::Value = serde_json::from_str(raw_json).map_err(FetchError::Json)?;
let json: model::Info = serde_json::from_value(json).map_err(FetchError::Json)?;
match json.thread {
model::Thread::Post { post } => {
let text = post.record.text;
let author = post.author.display_name.unwrap_or_default();
let author_id = post.author.handle;
let mut media = vec![];
if let Some(embed) = post.embed {
match embed {
model::Media::Images { images } => {
media.extend(images.into_iter().map(|image| Media::Illustration {
title: None,
url: image.fullsize,
thumbnail_url: Some(image.thumb),
fallback_url: None,
}));
}
model::Media::Video {
playlist,
thumbnail,
} => {
media.push(Media::Video {
title: None,
url: playlist,
thumbnail_url: thumbnail,
});
}
model::Media::External => {}
}
}
let sensitive = post
.labels
.iter()
.any(|label| SENSITIVE_LABEL.contains(&label.val.as_str()));
Ok(Post {
id,
author,
author_id,
text,
media,
sensitive,
})
}
model::Thread::NotFound => Err(FetchError::NotFound),
model::Thread::Blocked => Err(FetchError::Blocked),
}
}
}
impl From<Post> for Fetched {
fn from(post: Post) -> Self {
let url = post.url();
let author_url = post.author_url();
let render_data = Some(crate::site::RenderData {
url: url.clone(),
author: encode_text(&post.author).into_owned(),
author_url: author_url.clone(),
title: encode_text(&post.text).into_owned(),
tags: String::new(),
});
Fetched {
source_url: url,
caption: post.caption(),
title: post.text.clone(),
media: post.media,
sensitive: post.sensitive,
render_data,
_keep_alive: None,
}
}
}
const API_URL: &str = "https://public.api.bsky.app/xrpc/app.bsky.feed.getPostThread";
const SENSITIVE_LABEL: [&str; 4] = ["sexual", "nudity", "porn", "graphic-media"];
#[cfg(test)]
mod tests {
use super::*;
fn thread_json(post_json: serde_json::Value) -> serde_json::Value {
serde_json::json!({ "thread": post_json })
}
#[test]
fn pattern_matches_handle_and_did() {
let cases = [
(
"https://bsky.app/profile/user.bsky.social/post/3laoveufjv224",
"user.bsky.social",
"3laoveufjv224",
),
(
"https://bsky.app/profile/did:plc:abc123def/post/3xxxx",
"did:plc:abc123def",
"3xxxx",
),
];
for (url, handle, rkey) in cases {
let caps = PATTERN.captures(url).unwrap_or_else(|| panic!("{url}"));
assert_eq!(caps.get(1).unwrap().as_str(), handle);
assert_eq!(caps.get(2).unwrap().as_str(), rkey);
}
}
#[test]
fn pattern_rejects_non_post_urls() {
for url in [
"https://bsky.app/profile/user.bsky.social",
"https://bsky.app/profile/user.bsky.social/posts",
"https://x.com/user/status/123",
] {
assert!(!PATTERN.is_match(url), "{url}");
}
}
#[test]
fn from_json_images_with_missing_defaults() {
let raw = thread_json(serde_json::json!({
"$type": "app.bsky.feed.defs#threadViewPost",
"post": {
"author": { "handle": "user.bsky.social" },
"record": { "$type": "app.bsky.feed.post", "text": "hello <world>" },
"embed": {
"$type": "app.bsky.embed.images#view",
"images": [
{ "thumb": "https://cdn.bsky.app/img/thumb", "fullsize": "https://cdn.bsky.app/img/full", "alt": "" }
]
}
}
}));
let post = Post::from_json(&raw.to_string(), "3xxxx".into()).unwrap();
let fetched: Fetched = post.into();
assert_eq!(fetched.source_url, "https://bsky.app/profile/user.bsky.social/post/3xxxx");
assert_eq!(fetched.title, "hello <world>");
assert_eq!(fetched.media.len(), 1);
assert!(!fetched.sensitive);
// display_name absent -> empty fallback
assert!(
fetched.caption.contains("</a>: hello &lt;world&gt;"),
"caption: {}",
fetched.caption
);
}
#[test]
fn from_json_sensitive_labels() {
let raw = thread_json(serde_json::json!({
"$type": "app.bsky.feed.defs#threadViewPost",
"post": {
"author": { "handle": "u.bsky.social", "displayName": "U" },
"record": { "$type": "app.bsky.feed.post", "text": "x" },
"labels": [{ "val": "porn" }]
}
}));
let post = Post::from_json(&raw.to_string(), "3xxxx".into()).unwrap();
assert!(post.sensitive);
}
#[test]
fn from_json_blocked_and_not_found() {
let blocked = thread_json(serde_json::json!({
"$type": "app.bsky.feed.defs#blockedPost",
"blocked": true
}));
assert!(matches!(
Post::from_json(&blocked.to_string(), "3xxxx".into()),
Err(FetchError::Blocked)
));
let not_found = thread_json(serde_json::json!({
"$type": "app.bsky.feed.defs#notFoundPost",
"notFound": true
}));
assert!(matches!(
Post::from_json(&not_found.to_string(), "3xxxx".into()),
Err(FetchError::NotFound)
));
}
#[tokio::test]
async fn live_fetch_with_photos() {
let fetched = fetch_from_url(
"https://bsky.app/profile/asagi0398.bsky.social/post/3mqkhrq5w6k2m",
)
.await
.unwrap();
assert_eq!(
fetched.source_url,
"https://bsky.app/profile/asagi0398.bsky.social/post/3mqkhrq5w6k2m"
);
assert!(!fetched.caption.is_empty());
}
#[tokio::test]
async fn live_fetch_smoke() {
let fetched = fetch_from_url("https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224")
.await
.unwrap();
assert_eq!(
fetched.source_url,
"https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224"
);
assert!(!fetched.caption.is_empty());
}
}
+4
View File
@@ -0,0 +1,4 @@
mod interface;
mod model;
pub use interface::{PATTERN, Post, enabled, fetch_from_url};
+60
View File
@@ -0,0 +1,60 @@
use serde::Deserialize;
#[derive(Deserialize, Debug)]
pub(crate) struct Info {
pub(crate) thread: Thread,
}
#[derive(Deserialize, Debug)]
#[serde(tag = "$type")]
pub(crate) enum Thread {
#[serde(rename = "app.bsky.feed.defs#threadViewPost")]
Post { post: Post },
#[serde(rename = "app.bsky.feed.defs#notFoundPost")]
NotFound,
#[serde(rename = "app.bsky.feed.defs#blockedPost")]
Blocked,
}
#[derive(Deserialize, Debug)]
pub(crate) struct Post {
pub(crate) author: Author,
pub(crate) record: PostRecord,
pub(crate) embed: Option<Media>,
#[serde(default)]
pub(crate) labels: Vec<Label>,
}
#[derive(Deserialize, Debug)]
pub(crate) struct Author {
pub(crate) handle: String,
#[serde(rename = "displayName", default)]
pub(crate) display_name: Option<String>,
}
#[derive(Deserialize, Debug)]
pub(crate) struct PostRecord {
pub(crate) text: String,
}
#[derive(Deserialize, Debug)]
#[serde(tag = "$type")]
pub(crate) enum Media {
#[serde(rename = "app.bsky.embed.images#view")]
Images { images: Vec<Image> },
#[serde(rename = "app.bsky.embed.video#view")]
Video { playlist: String, thumbnail: String },
#[serde(rename = "app.bsky.embed.external#view")]
External,
}
#[derive(Deserialize, Debug)]
pub(crate) struct Image {
pub(crate) thumb: String,
pub(crate) fullsize: String,
}
#[derive(Deserialize, Debug)]
pub(crate) struct Label {
pub(crate) val: String,
}
+363
View File
@@ -0,0 +1,363 @@
//! Site fetching dispatcher and unified result types.
//!
//! Dispatch order: twitter → bsky → pixiv. Each site module exports a
//! `PATTERN`, `enabled()` and `fetch_from_url()`; a future site plugs in by
//! adding one guarded entry in [`fetch_once`].
use std::fmt;
use std::sync::LazyLock;
use std::time::Duration;
pub mod bsky;
pub mod pixiv;
pub mod twitter;
pub use pixiv::PixivError;
/// The result of fetching a post: canonical URL, HTML caption, raw text,
/// media list and spoiler flag. Produced by [`fetch`].
#[derive(Debug)]
pub struct Fetched {
/// Canonical URL: x.com/{author}/status/{id} |
/// https://www.pixiv.net/artworks/{id} |
/// https://bsky.app/profile/{handle}/post/{rkey}
pub source_url: String,
/// The exact HTML produced by the site's caption().
pub caption: String,
/// Raw post text (tweet text / bsky text / pixiv title).
pub title: String,
pub media: Vec<crate::media::Media>,
/// Spoiler flag for all media of this post.
pub sensitive: bool,
/// Raw values (pre-escaped) for user-customizable caption formats.
pub(crate) render_data: Option<RenderData>,
/// Keeps temp files (e.g. an encoded ugoira MP4) alive until the caller
/// finishes uploading; not part of the public contract.
pub(crate) _keep_alive: Option<tempfile::TempDir>,
}
/// Pre-escaped values for `{url} {author} {author_url} {title} {tags}`
/// placeholders in user-supplied caption formats.
#[derive(Debug)]
pub(crate) struct RenderData {
pub url: String,
pub author: String,
pub author_url: String,
pub title: String,
pub tags: String,
}
impl Fetched {
/// The site this post came from (used for per-site format overrides).
pub fn site_name(&self) -> &'static str {
if self.source_url.contains("x.com") || self.source_url.contains("twitter.com") {
"twitter"
} else if self.source_url.contains("bsky.app") {
"bsky"
} else if self.source_url.contains("pixiv.net") {
"pixiv"
} else {
"unknown"
}
}
/// Renders a user-supplied caption format. The format string is
/// HTML-escaped in full, then the (already-escaped) placeholder values
/// are substituted — users can structure text but never inject raw HTML
/// or attributes. An empty/unknown format falls back to the built-in
/// caption.
pub fn caption_with(&self, format: &str) -> String {
match (&self.render_data, format.is_empty()) {
(Some(data), false) => caption_from_fields(
format,
"",
&data.url,
&data.author,
&data.author_url,
&data.title,
&data.tags,
),
_ => self.caption.clone(),
}
}
/// The pre-escaped placeholder values (author, author_url, title, tags)
/// a caller needs to rebuild a caption later, e.g. for a cached post
/// where the [`Fetched`] is no longer available.
pub fn render_fields(&self) -> Option<(&str, &str, &str, &str)> {
self.render_data.as_ref().map(|d| {
(
d.author.as_str(),
d.author_url.as_str(),
d.title.as_str(),
d.tags.as_str(),
)
})
}
}
/// Renders a user-supplied caption format from raw (already-escaped) field
/// values with the same escaping/substitution rules as
/// [`Fetched::caption_with`]. An empty format returns `built_in` unchanged.
pub fn caption_from_fields(
format: &str,
built_in: &str,
url: &str,
author: &str,
author_url: &str,
title: &str,
tags: &str,
) -> String {
if format.is_empty() {
return built_in.to_string();
}
let escaped = html_escape::encode_text(format).into_owned();
escaped
.replace("{url}", url)
.replace("{author}", author)
.replace("{author_url}", author_url)
.replace("{title}", title)
.replace("{tags}", tags)
}
/// Stable per-post cache key derived from any supported URL, so variant
/// domains (x.com / twitter.com / fxtwitter.com, mobile, `/photo/N`
/// suffixes) map to the same post. Returns `"twitter:<id>"`,
/// `"pixiv:<id>"` or `"bsky:<handle>/<rkey>"`.
pub fn cache_key(url: &str) -> Option<String> {
if let Some(caps) = twitter::PATTERN.captures(url) {
return Some(format!("twitter:{}", &caps[1]));
}
if let Some(caps) = pixiv::PATTERN.captures(url) {
return Some(format!("pixiv:{}", &caps[1]));
}
if let Some(caps) = bsky::PATTERN.captures(url) {
return Some(format!("bsky:{}/{}", &caps[1], &caps[2]));
}
None
}
#[derive(Debug)]
pub enum FetchError {
Http(reqwest::Error),
Json(serde_json::Error),
Pixiv(PixivError),
NotFound,
Blocked,
/// The post exists but its content is withheld (twitter NSFW /
/// age-restricted tweets come back as an empty `{}` from syndication).
Sensitive,
}
impl fmt::Display for FetchError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
FetchError::Http(e) => write!(f, "http error: {e}"),
FetchError::Json(e) => write!(f, "json error: {e}"),
FetchError::Pixiv(e) => write!(f, "pixiv error: {e}"),
FetchError::NotFound => write!(f, "not found"),
FetchError::Blocked => write!(f, "blocked"),
FetchError::Sensitive => write!(f, "content withheld (sensitive)"),
}
}
}
impl std::error::Error for FetchError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
FetchError::Http(e) => Some(e),
FetchError::Json(e) => Some(e),
FetchError::Pixiv(e) => Some(e),
FetchError::NotFound | FetchError::Blocked | FetchError::Sensitive => None,
}
}
}
impl From<reqwest::Error> for FetchError {
fn from(e: reqwest::Error) -> Self {
FetchError::Http(e)
}
}
impl From<serde_json::Error> for FetchError {
fn from(e: serde_json::Error) -> Self {
FetchError::Json(e)
}
}
impl From<PixivError> for FetchError {
fn from(e: PixivError) -> Self {
FetchError::Pixiv(e)
}
}
/// Shared HTTP client (browser User-Agent) for twitter/bsky fetches and
/// [`download_media`].
pub(crate) static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
let builder = reqwest::Client::builder().user_agent("Mozilla/5.0");
// Each `#[tokio::test]` runs on its own runtime; the connection pool is
// bound to the runtime that created it, so cross-runtime reuse of idle
// connections fails with DispatchGone. In test builds every request uses
// a fresh connection. Production runs on one runtime and keeps pooling.
#[cfg(test)]
let builder = builder.pool_max_idle_per_host(0);
builder.build().expect("failed to build HTTP client")
});
/// Fetches a post from its URL. Returns `Ok(None)` when no site pattern
/// matches (unsupported links are silently ignored by the bot).
///
/// Transient network failures are retried: 3 total attempts with 1s then 2s
/// delays. Non-Http errors (Json/NotFound/Blocked/Pixiv) are not retried.
pub async fn fetch(url: &str) -> Result<Option<Fetched>, FetchError> {
let mut last_http_error = None;
for attempt in 0..3u32 {
match fetch_once(url).await {
Ok(Some(fetched)) => {
log::info!(
"fetched {url}: site {} returned {} media",
fetched.site_name(),
fetched.media.len()
);
return Ok(Some(fetched));
}
Ok(None) => return Ok(None),
Err(FetchError::Http(e)) => {
last_http_error = Some(e);
if attempt < 2 {
tokio::time::sleep(Duration::from_secs(1 << attempt)).await;
}
}
Err(other) => return Err(other),
}
}
Err(FetchError::Http(
last_http_error.expect("retry loop always ran 3 attempts"),
))
}
async fn fetch_once(url: &str) -> Result<Option<Fetched>, FetchError> {
if twitter::enabled() && twitter::PATTERN.is_match(url) {
return Ok(Some(twitter::fetch_from_url(url).await?));
}
if bsky::enabled() && bsky::PATTERN.is_match(url) {
return Ok(Some(bsky::fetch_from_url(url).await?));
}
if pixiv::enabled() && pixiv::PATTERN.is_match(url) {
return Ok(Some(pixiv::fetch_from_url(url).await?));
}
Ok(None)
}
/// Downloads media bytes for the bot's upload fallback: when Telegram's own
/// fetch of a media URL is blocked (hotlink protection), the bot downloads
/// the file itself and uploads it via multipart. Site-appropriate headers:
/// pixiv image hosts need the `Referer` header.
/// Returns the Content-Length of a media URL, or `None` when the server does
/// not report one. Used to check whether a file fits Telegram's size limits
/// before downloading/uploading it.
pub async fn media_size(url: &str) -> Result<Option<u64>, FetchError> {
let mut request = CLIENT.get(url);
let lower = url.to_ascii_lowercase();
if lower.contains("pximg.net") {
request = request.header("Referer", "https://www.pixiv.net/");
}
let response = request.send().await?;
Ok(response.content_length())
}
pub async fn download_media(url: &str) -> Result<bytes::Bytes, FetchError> {
let mut request = CLIENT.get(url);
let lower = url.to_ascii_lowercase();
if lower.contains("pximg.net") {
request = request.header("Referer", "https://www.pixiv.net/");
}
let response = request.send().await?;
Ok(response.bytes().await?)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_key_normalizes_domain_variants() {
assert_eq!(
cache_key("https://x.com/user/status/1234567890/photo/1"),
Some("twitter:1234567890".into())
);
assert_eq!(
cache_key("https://mobile.twitter.com/user/status/1234567890"),
Some("twitter:1234567890".into())
);
assert_eq!(
cache_key("https://fxtwitter.com/user/status/1234567890"),
Some("twitter:1234567890".into())
);
assert_eq!(
cache_key("https://www.pixiv.net/artworks/123456"),
Some("pixiv:123456".into())
);
assert_eq!(
cache_key("https://bsky.app/profile/handle.example/post/3lorem"),
Some("bsky:handle.example/3lorem".into())
);
assert_eq!(cache_key("https://example.com/not-a-post"), None);
}
#[test]
fn caption_from_fields_substitutes_and_escapes() {
// The format string is escaped, the field values are substituted
// verbatim (callers pass the already-escaped render data).
let out = caption_from_fields(
"see {author} at {url} — {title}",
"",
"https://x.com/u/status/1",
"A &amp; B",
"https://x.com/u",
"hello <world>",
"",
);
assert_eq!(
out,
"see A &amp; B at https://x.com/u/status/1 — hello <world>"
);
// Empty format keeps the built-in caption untouched.
assert_eq!(
caption_from_fields("", "built-in", "u", "a", "au", "t", "g"),
"built-in"
);
}
#[tokio::test]
async fn unsupported_url_returns_none() {
let result = fetch("https://example.com/some/article").await;
assert!(matches!(result, Ok(None)), "got {result:?}");
}
#[tokio::test]
async fn unknown_scheme_returns_none() {
let result = fetch("not a url at all").await;
assert!(matches!(result, Ok(None)), "got {result:?}");
}
#[tokio::test]
async fn download_media_pixiv_original_with_referer() {
// Proves the Referer header is attached for i.pximg.net: a header-less
// GET to a pixiv original URL is rejected with 403.
if std::env::var("PIXIV_REFRESH_TOKEN").is_err() {
eprintln!("skipping: no PIXIV_REFRESH_TOKEN");
return;
}
let illustration = pixiv::fetch(126839080).await.unwrap();
let fetched: Fetched = illustration.into();
let url = match fetched.media.first() {
Some(crate::media::Media::Illustration { url, .. }) => url.clone(),
other => panic!("expected illustration media, got {other:?}"),
};
assert!(url.contains("i.pximg.net"));
let bytes = download_media(&url).await.unwrap();
assert!(!bytes.is_empty());
}
}
+394
View File
@@ -0,0 +1,394 @@
//! Native pixiv app-API client (replaces pixiv3-rs).
//!
//! Token exchange against `oauth.secure.pixiv.net` and illust detail against
//! `app-api.pixiv.net`, deserialized with the kept `model.rs` types.
use super::interface::Illustration;
use super::model::{IllustrationModel, TypeModel, UgoiraMetadataModel};
use crate::media::Media;
use crate::site::FetchError;
use std::env;
use std::fmt;
use std::io::{Cursor, Read};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::LazyLock;
use std::time::{Duration, SystemTime};
const AUTH_TOKEN_URL: &str = "https://oauth.secure.pixiv.net/auth/token";
const APP_API_URL: &str = "https://app-api.pixiv.net";
const CLIENT_ID: &str = "MOBrBDS8blbauoSck0ZfDbtuzpyT";
const CLIENT_SECRET: &str = "lsACyCD94FhDUtGTXi3QzcFE2uU1hqtDaKeqrdwj";
const AUTH_USER_AGENT: &str = "PixivAndroidApp/5.0.234 (Android 11; Pixel 5)";
const APP_USER_AGENT: &str = "PixivIOSApp/7.13.3 (iOS 14.6; iPhone13,2)";
/// Token refresh safe margin (seconds).
const TOKEN_REFRESH_SAFE_MARGIN: u64 = 300;
#[derive(Debug)]
pub enum PixivError {
/// No refresh token available (PIXIV_REFRESH_TOKEN unset).
NoAuth,
Http(reqwest::Error),
Json(serde_json::Error),
Api(String),
}
impl fmt::Display for PixivError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
PixivError::NoAuth => write!(f, "pixiv: no authentication"),
PixivError::Http(e) => write!(f, "pixiv http error: {e}"),
PixivError::Json(e) => write!(f, "pixiv json error: {e}"),
PixivError::Api(message) => write!(f, "pixiv api error: {message}"),
}
}
}
impl std::error::Error for PixivError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
PixivError::Http(e) => Some(e),
PixivError::Json(e) => Some(e),
_ => None,
}
}
}
impl From<reqwest::Error> for PixivError {
fn from(e: reqwest::Error) -> Self {
PixivError::Http(e)
}
}
impl From<serde_json::Error> for PixivError {
fn from(e: serde_json::Error) -> Self {
PixivError::Json(e)
}
}
/// Native pixiv app-API client.
pub struct PixivAPI {
refresh_token: String,
access_token: tokio::sync::Mutex<Option<(String, SystemTime)>>,
}
impl PixivAPI {
pub fn new(refresh_token: String) -> Self {
Self {
refresh_token,
access_token: tokio::sync::Mutex::new(None),
}
}
/// Returns a valid access token, exchanging the refresh token when none
/// is cached or it has expired.
pub async fn get_access_token(&self) -> Result<String, PixivError> {
let mut guard = self.access_token.lock().await;
if let Some((token, expires_at)) = guard.as_ref()
&& *expires_at > SystemTime::now()
{
return Ok(token.clone());
}
let response = crate::site::CLIENT
.post(AUTH_TOKEN_URL)
.form(&[
("client_id", CLIENT_ID),
("client_secret", CLIENT_SECRET),
("grant_type", "refresh_token"),
("include_policy", "true"),
("refresh_token", &self.refresh_token),
])
.header("User-Agent", AUTH_USER_AGENT)
.send()
.await?;
let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
let access_token = json
.get("access_token")
.and_then(|v| v.as_str())
.ok_or_else(|| {
let message = json
.get("error")
.and_then(|v| v.as_str())
.unwrap_or("invalid token response");
PixivError::Api(message.to_string())
})?
.to_string();
let expires_in = json
.get("expires_in")
.and_then(|v| v.as_u64())
.filter(|&sec| sec > 0)
.unwrap_or(3600);
let expires_at = SystemTime::now()
+ Duration::from_secs(expires_in.saturating_sub(TOKEN_REFRESH_SAFE_MARGIN));
*guard = Some((access_token.clone(), expires_at));
Ok(access_token)
}
/// Fetches illust detail from the app API.
pub async fn illust_detail(&self, illust_id: u64) -> Result<IllustrationModel, PixivError> {
let access_token = self.get_access_token().await?;
let response = crate::site::CLIENT
.get(format!("{APP_API_URL}/v1/illust/detail?illust_id={illust_id}"))
.header("app-os", "ios")
.header("app-os-version", "14.6")
.header("User-Agent", APP_USER_AGENT)
.bearer_auth(access_token)
.send()
.await?;
let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
if json.get("error").is_some() {
let message = json
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("illust detail failed");
return Err(PixivError::Api(message.to_string()));
}
let illust = json
.get("illust")
.ok_or_else(|| PixivError::Api("missing illust in response".to_string()))?;
Ok(serde_json::from_value(illust.clone())?)
}
pub async fn fetch(&self, illust_id: u64) -> Result<Illustration, FetchError> {
let model = self.illust_detail(illust_id).await?;
let mut illustration = Illustration::from_model(&model);
if matches!(&model.r#type, TypeModel::Ugoira) {
// Real ugoira support: download the frame zip and encode an MP4.
// Without ffmpeg (or on encode failure) the post stays
// unsupported (empty media, like Python).
match self.ugoira_video(illust_id).await {
Ok(Some((mp4_path, _keep_alive))) => {
illustration.media.push(Media::Video {
title: None,
url: mp4_path,
thumbnail_url: model.image_urls.medium.clone(),
});
illustration._keep_alive = Some(_keep_alive);
}
Ok(None) => {}
Err(e) => log::error!("ugoira encode failed for {illust_id}: {e}"),
}
}
Ok(illustration)
}
/// Fetches ugoira metadata (frame zip + frame delays) from the app API.
pub async fn ugoira_metadata(&self, illust_id: u64) -> Result<UgoiraMetadataModel, PixivError> {
let access_token = self.get_access_token().await?;
let response = crate::site::CLIENT
.get(format!("{APP_API_URL}/v1/ugoira/metadata?illust_id={illust_id}"))
.header("app-os", "ios")
.header("app-os-version", "14.6")
.header("User-Agent", APP_USER_AGENT)
.bearer_auth(access_token)
.send()
.await?;
let json: serde_json::Value = serde_json::from_str(&response.text().await?)?;
if json.get("error").is_some() {
let message = json
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("ugoira metadata failed");
return Err(PixivError::Api(message.to_string()));
}
let metadata = json
.get("ugoira_metadata")
.ok_or_else(|| PixivError::Api("missing ugoira_metadata".to_string()))?;
Ok(serde_json::from_value(metadata.clone())?)
}
/// Downloads the frame zip and encodes one MP4 via ffmpeg. Returns the
/// MP4 path plus the temp directory that must stay alive until the file
/// is uploaded.
async fn ugoira_video(
&self,
illust_id: u64,
) -> Result<Option<(String, tempfile::TempDir)>, PixivError> {
if !ffmpeg_available() {
log_once_ffmpeg_missing();
return Ok(None);
}
let metadata = self.ugoira_metadata(illust_id).await?;
if metadata.frames.is_empty() {
return Ok(None);
}
let zip_url = metadata
.zip_url
.clone()
.or_else(|| metadata.zip_urls.as_ref().map(|z| z.medium.clone()));
let Some(zip_url) = zip_url else {
return Ok(None);
};
let zip_bytes = crate::site::download_media(&zip_url).await.map_err(|e| match e {
FetchError::Http(e) => PixivError::Http(e),
other => PixivError::Api(format!("frame zip download failed: {other}")),
})?;
let frame_delays = metadata.frames.iter().map(|f| f.delay).collect::<Vec<_>>();
let result = tokio::task::spawn_blocking(
move || -> Result<(String, tempfile::TempDir), String> {
let frames_dir = tempfile::tempdir().map_err(|e| e.to_string())?;
let out_dir = tempfile::tempdir().map_err(|e| e.to_string())?;
// Extract frames to canonical zero-padded names; pixiv ugoira
// frames are uniformly jpg or png per artwork.
let mut archive = zip::ZipArchive::new(Cursor::new(zip_bytes))
.map_err(|e| format!("unzip: {e}"))?;
// pixiv ugoira frames are uniformly jpg or png per artwork; take
// the extension from the first entry.
let extension = if archive.len() > 0 {
let first_name = archive
.by_index(0)
.map_err(|e| e.to_string())?
.name()
.to_string();
first_name
.rsplit('.')
.next()
.unwrap_or("jpg")
.to_string()
} else {
"jpg".to_string()
};
let mut count = 0usize;
for i in 0..archive.len() {
let mut entry = archive.by_index(i).map_err(|e| e.to_string())?;
let mut bytes = Vec::new();
entry.read_to_end(&mut bytes).map_err(|e| e.to_string())?;
let path = frames_dir.path().join(format!("img_{count:05}.{extension}"));
std::fs::write(&path, bytes).map_err(|e| e.to_string())?;
count += 1;
}
if count == 0 {
return Err("empty frame zip".to_string());
}
// Constant rate from the median frame delay (ms).
let mut delays = frame_delays;
delays.sort_unstable();
let median = delays[delays.len() / 2].max(1);
let framerate = 1000.0 / median as f64;
let output = out_dir.path().join("ugoira.mp4");
let status = std::process::Command::new("ffmpeg")
.args([
"-y",
"-framerate",
&framerate.to_string(),
"-i",
&frames_dir.path().join(format!("img_%05d.{extension}")).to_string_lossy(),
// libx264 needs even dimensions; pixiv ugoira frames can
// be odd-sized (e.g. 277x405).
"-vf",
"scale=trunc(iw/2)*2:trunc(ih/2)*2",
"-c:v",
"libx264",
"-pix_fmt",
"yuv420p",
"-movflags",
"+faststart",
&output.to_string_lossy(),
])
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
.map_err(|e| format!("ffmpeg spawn failed: {e}"))?;
if !status.success() {
return Err(format!("ffmpeg exited with {status}"));
}
Ok((output.to_string_lossy().into_owned(), out_dir))
},
)
.await
.expect("ugoira encode worker panicked");
match result {
Ok(pair) => Ok(Some(pair)),
Err(message) => {
log::error!("ugoira encode failed for {illust_id}: {message}");
Ok(None)
}
}
}
}
static FFMPEG_AVAILABLE: LazyLock<bool> = LazyLock::new(|| {
std::process::Command::new("ffmpeg")
.arg("-version")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
.map(|s| s.success())
.unwrap_or(false)
});
static FFMPEG_MISSING_LOGGED: AtomicBool = AtomicBool::new(false);
fn ffmpeg_available() -> bool {
*FFMPEG_AVAILABLE
}
fn log_once_ffmpeg_missing() {
if !FFMPEG_MISSING_LOGGED.swap(true, Ordering::Relaxed) {
log::warn!("ffmpeg not found; pixiv ugoira posts stay unsupported");
}
}
/// pixiv3-rs replacement: `None` when `PIXIV_REFRESH_TOKEN` is unset.
static PIXIV_CLIENT: LazyLock<Option<PixivAPI>> = LazyLock::new(|| {
env::var("PIXIV_REFRESH_TOKEN").ok().map(PixivAPI::new)
});
/// Set at startup when the login validation fails; pixiv stays disabled until
/// the next process start.
static DISABLED: AtomicBool = AtomicBool::new(false);
pub fn enabled() -> bool {
!DISABLED.load(Ordering::Relaxed) && env::var("PIXIV_REFRESH_TOKEN").is_ok()
}
/// Permanently disables pixiv until the next process start.
pub fn disable() {
DISABLED.store(true, Ordering::Relaxed);
}
/// Forces the refresh-token → access-token exchange now, surfacing invalid
/// tokens and network errors. Called once at bot startup; on failure the bot
/// calls [`disable`].
pub async fn validate() -> Result<(), PixivError> {
match PIXIV_CLIENT.as_ref() {
None => Err(PixivError::NoAuth),
Some(client) => {
client.get_access_token().await?;
Ok(())
}
}
}
pub async fn fetch(illust_id: u64) -> Result<Illustration, FetchError> {
let client = PIXIV_CLIENT
.as_ref()
.filter(|_| enabled())
.ok_or(FetchError::Pixiv(PixivError::NoAuth))?;
client.fetch(illust_id).await
}
#[cfg(test)]
mod tests {
use super::*;
use dotenv::dotenv;
#[tokio::test]
async fn test_fetch() {
dotenv().ok();
let result = fetch(126839080).await;
assert!(result.is_ok());
println!("{:#?}", result);
}
#[tokio::test]
async fn validate_with_bogus_token_fails() {
dotenv().ok();
// A bogus token must surface as Api error (invalid_grant), not panic.
let client = PixivAPI::new("bogus_token_for_testing".to_string());
let result = client.get_access_token().await;
assert!(matches!(result, Err(PixivError::Api(_))), "got {result:?}");
}
}
+375
View File
@@ -0,0 +1,375 @@
use super::model::{IllustrationModel, TypeModel};
use crate::media::Media;
use crate::site::{FetchError, Fetched};
use html_escape::encode_text;
use regex::Regex;
use std::sync::LazyLock;
pub static PATTERN: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?:www\.)?pixiv\.net/(?:en/)?(?:(?:i|artworks)/|member_illust\.php\?(?:mode=[a-z_]*&)?illust_id=)(\d+)").unwrap()
});
pub fn enabled() -> bool {
super::api::enabled()
}
pub async fn fetch_from_url(url: &str) -> Result<Fetched, FetchError> {
let id = PATTERN
.captures(url)
.and_then(|caps| caps.get(1))
.map(|m| m.as_str())
.ok_or(FetchError::NotFound)?;
let id = id.parse::<u64>().map_err(|_| FetchError::NotFound)?;
Ok(super::api::fetch(id).await?.into())
}
#[derive(Debug)]
pub struct Illustration {
id: String,
title: String,
author: String,
author_id: String,
tags: Vec<String>,
pub(crate) media: Vec<Media>,
nsfw: bool,
/// Keeps a temp dir (ugoira MP4) alive until the send completes.
pub(crate) _keep_alive: Option<tempfile::TempDir>,
}
impl Illustration {
fn url(&self) -> String {
format!("https://www.pixiv.net/artworks/{}", self.id)
}
fn author_url(&self) -> String {
format!("https://www.pixiv.net/users/{}", self.author_id)
}
pub fn caption(&self) -> String {
format!(
"<a href=\"{url}\">{title}</a> / <a href=\"{author_url}\">{author}</a>\n{tags}",
url = self.url(),
title = encode_text(&self.title),
author_url = self.author_url(),
author = encode_text(&self.author),
tags = encode_text(
&self
.tags
.iter()
.map(|tag| format!("#{tag}"))
.collect::<Vec<_>>()
.join(" ")
),
)
}
pub fn from_model(model: &IllustrationModel) -> Self {
let id = model.id.to_string();
let title = model.title.clone();
let author = model.user.name.clone();
let author_id = model.user.id.to_string();
let mut tags: Vec<String> = model.tags.iter().map(|tag| tag.name.clone()).collect();
// illust_ai_type: 0 = undefined, 1 = not AI, 2 = AI-generated.
// Mark AI works with a leading #AI tag (rendered via the `#{tag}`
// caption format).
if model.illust_ai_type == 2 {
tags.insert(0, "AI".to_string());
}
let mut media = vec![];
if matches!(&model.r#type, TypeModel::Ugoira) {
// No static images for ugoira; the fetch path encodes an MP4 via
// ffmpeg and appends it as a Video item (api.rs). This fallback
// keeps media empty when encoding fails or ffmpeg is missing.
} else if model.page_count > 1 {
media.extend(model.meta_pages.iter().filter_map(|page| {
page.image_urls.original.clone().map(|original| Media::Illustration {
title: None,
url: original,
thumbnail_url: Some(page.image_urls.medium.clone()),
fallback_url: Some(page.image_urls.large.clone()),
})
}));
} else if let Some(original) = model
.meta_single_page
.original_image_url
.clone()
.or(model.image_urls.original.clone())
{
media.push(Media::Illustration {
title: None,
url: original,
thumbnail_url: Some(model.image_urls.medium.clone()),
fallback_url: Some(model.image_urls.large.clone()),
});
}
let nsfw = model.sanity_level > 5;
Self {
id,
title,
author,
author_id,
tags,
media,
nsfw,
_keep_alive: None,
}
}
}
impl From<Illustration> for Fetched {
fn from(illustration: Illustration) -> Self {
let url = illustration.url();
let author_url = illustration.author_url();
let tags = illustration
.tags
.iter()
.map(|tag| format!("#{tag}"))
.collect::<Vec<_>>()
.join(" ");
let render_data = Some(crate::site::RenderData {
url: url.clone(),
author: encode_text(&illustration.author).into_owned(),
author_url: author_url.clone(),
title: encode_text(&illustration.title).into_owned(),
tags: encode_text(&tags).into_owned(),
});
Fetched {
source_url: url,
caption: illustration.caption(),
title: illustration.title.clone(),
media: illustration.media,
sensitive: illustration.nsfw,
render_data,
_keep_alive: illustration._keep_alive,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use super::super::model::IllustrationModel;
fn illust_json(
type_: &str,
page_count: u8,
single_original: Option<&str>,
image_urls_original: Option<&str>,
pages: Vec<(Option<&str>, &str, &str)>,
ai_type: i32,
) -> serde_json::Value {
let meta_pages: Vec<serde_json::Value> = pages
.into_iter()
.map(|(original, medium, large)| {
serde_json::json!({
"image_urls": {
"medium": medium,
"large": large,
"original": original
}
})
})
.collect();
serde_json::json!({
"illust": {
"id": 123,
"title": "Art <title>",
"type": type_,
"image_urls": {
"medium": "medium.jpg",
"large": "large.jpg",
"original": image_urls_original
},
"user": { "id": 456, "name": "Artist" },
"tags": [{ "name": "tag1" }, { "name": "tag2" }],
"page_count": page_count,
"sanity_level": 6,
"illust_ai_type": ai_type,
"meta_single_page": { "original_image_url": single_original },
"meta_pages": meta_pages
}
})
}
fn parse(v: serde_json::Value) -> Illustration {
let model: IllustrationModel = serde_json::from_value(v["illust"].clone()).unwrap();
Illustration::from_model(&model)
}
#[test]
fn pattern_matches_all_forms() {
let cases = [
("https://www.pixiv.net/artworks/123456", "123456"),
("https://pixiv.net/artworks/123456", "123456"),
("https://www.pixiv.net/en/artworks/123456", "123456"),
("https://www.pixiv.net/i/123456", "123456"),
("https://www.pixiv.net/member_illust.php?mode=medium&illust_id=123456", "123456"),
("https://www.pixiv.net/en/member_illust.php?illust_id=123456", "123456"),
];
for (url, id) in cases {
let caps = PATTERN.captures(url).unwrap_or_else(|| panic!("{url}"));
assert_eq!(caps.get(1).unwrap().as_str(), id);
}
}
#[test]
fn pattern_rejects_non_artwork_urls() {
for url in [
"https://www.pixiv.net/users/123",
"https://x.com/user/status/123",
"https://bsky.app/profile/u/post/3xxxx",
] {
assert!(!PATTERN.is_match(url), "{url}");
}
}
#[test]
fn ugoira_yields_empty_media() {
let v = illust_json("ugoira", 1, Some("https://i.pximg.net/orig.jpg"), None, vec![], 0);
let illustration = parse(v);
let fetched: Fetched = illustration.into();
assert!(fetched.media.is_empty());
assert!(fetched.sensitive, "sanity_level 6 > 5");
assert_eq!(fetched.title, "Art <title>");
}
#[test]
fn single_page_with_single_original() {
let v = illust_json(
"illust",
1,
Some("https://i.pximg.net/single.jpg"),
None,
vec![],
0,
);
let fetched: Fetched = parse(v).into();
assert_eq!(fetched.media.len(), 1);
match &fetched.media[0] {
Media::Illustration { url, .. } => {
assert_eq!(url, "https://i.pximg.net/single.jpg")
}
other => panic!("expected Illustration, got {other:?}"),
}
}
#[test]
fn single_page_falls_back_to_image_urls_original() {
let v = illust_json(
"illust",
1,
None,
Some("https://i.pximg.net/fallback.jpg"),
vec![],
0,
);
let fetched: Fetched = parse(v).into();
assert_eq!(fetched.media.len(), 1);
match &fetched.media[0] {
Media::Illustration { url, .. } => {
assert_eq!(url, "https://i.pximg.net/fallback.jpg")
}
other => panic!("expected Illustration, got {other:?}"),
}
}
#[test]
fn single_page_without_any_original_is_empty() {
let v = illust_json("illust", 1, None, None, vec![], 0);
let fetched: Fetched = parse(v).into();
assert!(fetched.media.is_empty());
}
#[test]
fn multi_page_skips_pages_without_original() {
let v = illust_json(
"illust",
2,
None,
None,
vec![
(None, "m1.jpg", "l1.jpg"),
(Some("https://i.pximg.net/p2.jpg"), "m2.jpg", "l2.jpg"),
],
0,
);
let fetched: Fetched = parse(v).into();
assert_eq!(fetched.media.len(), 1);
match &fetched.media[0] {
Media::Illustration { url, thumbnail_url, fallback_url, .. } => {
assert_eq!(url, "https://i.pximg.net/p2.jpg");
assert_eq!(thumbnail_url.as_deref(), Some("m2.jpg"));
assert_eq!(fallback_url.as_deref(), Some("l2.jpg"));
}
other => panic!("expected Illustration, got {other:?}"),
}
}
#[test]
fn caption_with_escapes_format_and_substitutes() {
let v = illust_json("illust", 1, Some("https://i.pximg.net/o.jpg"), None, vec![], 0);
let fetched: Fetched = parse(v).into();
// Format string is escaped in full, then placeholders substituted.
let out = fetched.caption_with("{title} by {author} <script> {tags}");
assert!(
out.contains("Art &lt;title&gt; by Artist &lt;script&gt; #tag1 #tag2"),
"got: {out}"
);
assert!(!out.contains("<script>"), "no raw HTML injection: {out}");
// {url} and {author_url} carry the site's own URLs.
let out = fetched.caption_with("{url} {author_url}");
assert_eq!(
out,
"https://www.pixiv.net/artworks/123 https://www.pixiv.net/users/456"
);
// Empty format falls back to the built-in caption.
assert_eq!(fetched.caption_with(""), fetched.caption);
assert_eq!(fetched.site_name(), "pixiv");
}
#[test]
fn ai_work_gets_leading_ai_tag() {
// illust_ai_type == 2 is the only AI marker.
let v = illust_json("illust", 1, Some("https://i.pximg.net/o.jpg"), None, vec![], 2);
let fetched: Fetched = parse(v).into();
assert!(
fetched.caption.contains("#AI #tag1 #tag2"),
"caption: {}",
fetched.caption
);
// The {tags} placeholder reflects the tag array too.
assert!(fetched.caption_with("{tags}").starts_with("#AI "), "got: {}", fetched.caption_with("{tags}"));
}
#[test]
fn non_ai_work_has_no_ai_tag() {
// 1 = explicitly not AI, 0 = undefined: neither gets the #AI tag.
for ai_type in [0, 1] {
let v = illust_json("illust", 1, Some("https://i.pximg.net/o.jpg"), None, vec![], ai_type);
let fetched: Fetched = parse(v).into();
assert!(
!fetched.caption.contains("#AI"),
"ai_type={ai_type} got: {}",
fetched.caption
);
}
}
#[test]
fn caption_escapes_and_links() {
let v = illust_json("illust", 1, Some("https://i.pximg.net/o.jpg"), None, vec![], 0);
let fetched: Fetched = parse(v).into();
assert!(
fetched
.caption
.contains("<a href=\"https://www.pixiv.net/artworks/123\">Art &lt;title&gt;</a>"),
"caption: {}",
fetched.caption
);
assert!(fetched.caption.contains("#tag1 #tag2"));
assert_eq!(
fetched.source_url,
"https://www.pixiv.net/artworks/123"
);
}
}
+6
View File
@@ -0,0 +1,6 @@
mod api;
mod interface;
mod model;
pub use api::{PixivAPI, PixivError, disable, fetch, validate};
pub use interface::{PATTERN, Illustration, enabled, fetch_from_url};
+79
View File
@@ -0,0 +1,79 @@
// Model set for the native pixiv app-API client (app-api.pixiv.net).
use serde::Deserialize;
#[derive(Deserialize, Debug)]
pub struct IllustrationModel {
pub id: u64,
pub title: String,
pub r#type: TypeModel,
pub image_urls: ImageUrlsModel,
pub user: UserInfoModel,
pub tags: Vec<IllustrationTagModel>,
pub page_count: u8,
pub sanity_level: u8,
/// 0 = undefined (unlabeled), 1 = not AI, 2 = AI-generated.
pub illust_ai_type: i32,
pub meta_single_page: MetaSinglePageModel,
pub meta_pages: Vec<MetaPageModel>,
}
#[derive(Deserialize, Debug)]
pub enum TypeModel {
#[serde(rename = "illust")]
Illust,
#[serde(rename = "manga")]
Manga,
#[serde(rename = "ugoira")]
Ugoira,
}
#[derive(Deserialize, Debug)]
pub struct UserInfoModel {
pub id: u64,
pub name: String,
}
#[derive(Deserialize, Debug)]
pub struct ImageUrlsModel {
pub medium: String,
pub large: String,
#[serde(default)]
pub original: Option<String>,
}
#[derive(Deserialize, Debug)]
pub struct IllustrationTagModel {
pub name: String,
}
#[derive(Deserialize, Debug)]
pub struct MetaSinglePageModel {
#[serde(default)]
pub original_image_url: Option<String>,
}
#[derive(Deserialize, Debug)]
pub struct MetaPageModel {
pub image_urls: ImageUrlsModel,
}
#[derive(Deserialize, Debug)]
pub struct UgoiraMetadataModel {
/// Older API shape (`zip_url`); newer responses use `zip_urls.medium`.
#[serde(default)]
pub zip_url: Option<String>,
#[serde(default)]
pub zip_urls: Option<UgoiraZipUrlsModel>,
pub frames: Vec<UgoiraFrameModel>,
}
#[derive(Deserialize, Debug)]
pub struct UgoiraZipUrlsModel {
pub medium: String,
}
#[derive(Deserialize, Debug)]
pub struct UgoiraFrameModel {
pub delay: u32,
}
+376
View File
@@ -0,0 +1,376 @@
//! Authenticated fallback for tweets the public syndication endpoint refuses
//! to serve (NSFW / age-restricted tweets come back as an empty `{}`).
//!
//! Mirrors nazurin's web API client ([`web.py`]) and is used *only* when
//! syndication reports [`FetchError::Sensitive`]: the private GraphQL
//! `TweetDetail` endpoint, authenticated with a browser session cookie from
//! `TWITTER_AUTH_TOKEN` (the `auth_token` cookie value of a logged-in x.com
//! session). A fresh random `ct0` is generated per call; X checks that the
//! `x-csrf-token` header matches the cookie, not that it issued the value.
//!
//! [`web.py`]: https://github.com/y-young/nazurin/blob/master/nazurin/sites/twitter/api/web.py
//!
//! # Caveats
//! - X rotates the GraphQL query id when it rolls the web app; if requests
//! start failing, update [`TWEET_DETAIL_QUERY_ID`]. Fresh references from
//! the actively maintained FxEmbed/FxEmbed: TweetDetail
//! `R9IzzyzQBV87-DOWpcvDmw`, TweetResultByRestId `f2sagi1jweVHFkTUIHzmMQ`
//! (the latter is anonymous and surfaces NSFW tweets as
//! `reason: NsfwLoggedOut`).
//! - `x-client-transaction-id` is only required for `SearchTimeline`
//! (verified against FxEmbed's `proxy/allowlist.ts`) — TweetDetail works
//! without it; no need for the nazurin home-page/JS-bundle derivation.
use std::sync::LazyLock;
use serde_json::{json, Value};
use crate::site::FetchError;
use super::interface::Tweet;
/// `auth_token` cookie of a logged-in x.com session; enables the fallback.
/// Trimmed: a CRLF `.env` (Windows) leaves a trailing `\r` on the value,
/// which would make the Cookie header invalid.
static AUTH_TOKEN: LazyLock<Option<String>> = LazyLock::new(|| {
std::env::var("TWITTER_AUTH_TOKEN")
.ok()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
});
/// Public "logged in" client token used by the x.com web app.
const LOGGED_IN_BEARER: &str =
"Bearer AAAAAAAAAAAAAAAAAAAAANRILgAAAAAAnNwIzUejRCOuH5E6I8xnZz4puTs%3D1Zv7ttfk8LF81IUq16cHjhLTvJu4FA33AGWWjCpTnA";
/// `TweetDetail` query id (from nazurin; still valid as of 2026-08,
/// corroborated by the current FxEmbed build — see module caveats).
const TWEET_DETAIL_QUERY_ID: &str = "_8aYOgEDz35BrBcBal1-_w";
fn variables(id: &str) -> Value {
json!({
"focalTweetId": id,
"with_rux_injections": false,
"includePromotedContent": false,
"withCommunity": true,
"withQuickPromoteEligibilityTweetFields": false,
"withBirdwatchNotes": false,
"withVoice": true,
})
}
fn features() -> Value {
json!({
"rweb_video_screen_enabled": false,
"profile_label_improvements_pcf_label_in_post_enabled": true,
"rweb_tipjar_consumption_enabled": true,
"verified_phone_label_enabled": false,
"creator_subscriptions_tweet_preview_api_enabled": true,
"responsive_web_graphql_timeline_navigation_enabled": true,
"responsive_web_graphql_skip_user_profile_image_extensions_enabled": false,
"premium_content_api_read_enabled": false,
"communities_web_enable_tweet_community_results_fetch": true,
"c9s_tweet_anatomy_moderator_badge_enabled": true,
"responsive_web_grok_analyze_button_fetch_trends_enabled": false,
"responsive_web_grok_analyze_post_followups_enabled": true,
"responsive_web_jetfuel_frame": false,
"responsive_web_grok_share_attachment_enabled": true,
"articles_preview_enabled": true,
"responsive_web_edit_tweet_api_enabled": true,
"graphql_is_translatable_rweb_tweet_is_translatable_enabled": true,
"view_counts_everywhere_api_enabled": true,
"longform_notetweets_consumption_enabled": true,
"responsive_web_twitter_article_tweet_consumption_enabled": true,
"tweet_awards_web_tipping_enabled": false,
"responsive_web_grok_show_grok_translated_post": false,
"responsive_web_grok_analysis_button_from_backend": true,
"creator_subscriptions_quote_tweet_preview_enabled": false,
"freedom_of_speech_not_reach_fetch_enabled": true,
"standardized_nudges_misinfo": true,
"tweet_with_visibility_results_prefer_gql_limited_actions_policy_enabled": true,
"longform_notetweets_rich_text_read_enabled": true,
"longform_notetweets_inline_media_enabled": true,
"responsive_web_grok_image_annotation_enabled": true,
"responsive_web_enhance_cards_enabled": false,
})
}
/// Whether the authenticated fallback is available.
pub fn enabled() -> bool {
AUTH_TOKEN.is_some()
}
/// Fetches a tweet as the logged-in user via the private GraphQL API.
/// Returns the syndication-shaped [`Tweet`] (media included for NSFW posts).
pub async fn fetch(id: &str) -> Result<Tweet, FetchError> {
let token = AUTH_TOKEN
.as_deref()
.ok_or(FetchError::Sensitive)?;
// 16 random bytes as 32 hex chars: X rejects ct0 values of any other
// length with 403 code 353 ("matching csrf cookie and header").
let ct0: String = (0..16)
.map(|_| format!("{:02x}", rand::random::<u8>()))
.collect();
let response = crate::site::CLIENT
.get(format!(
"https://x.com/i/api/graphql/{TWEET_DETAIL_QUERY_ID}/TweetDetail"
))
.query(&[
("variables", variables(id).to_string()),
("features", features().to_string()),
])
.header("authorization", LOGGED_IN_BEARER)
.header("x-csrf-token", &ct0)
.header("x-twitter-auth-type", "OAuth2Session")
.header("cookie", format!("auth_token={token}; ct0={ct0}"))
.header("x-twitter-client-language", "en")
.header("x-twitter-active-user", "yes")
.header("referer", "https://x.com/")
.send()
.await?;
if !response.status().is_success() {
log::warn!("twitter auth fetch {id}: HTTP {}", response.status());
return Err(FetchError::NotFound);
}
let text = response.text().await?;
let json: Value = serde_json::from_str(&text)?;
let result = parse_tweet_result(&json, id)?;
let syndication_shape = to_syndication_shape(&result)
.ok_or_else(|| FetchError::Json(serde_json::Error::io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"missing tweet fields in GraphQL response",
))))?;
Tweet::from_syndication_json(&syndication_shape.to_string()).map_err(FetchError::Json)
}
/// Locates the tweet for `id` in a `TweetDetail` response and unwraps
/// visibility wrappers / retweets, mirroring nazurin's `_process_response`.
fn parse_tweet_result(json: &Value, id: &str) -> Result<Value, FetchError> {
if let Some(errors) = json.get("errors").and_then(|e| e.as_array()) {
let messages: Vec<&str> = errors
.iter()
.filter_map(|e| e.get("message").and_then(|m| m.as_str()))
.collect();
log::warn!("twitter auth fetch {id} failed: {}", messages.join("; "));
return Err(FetchError::NotFound);
}
let instructions = json
.pointer("/data/threaded_conversation_with_injections_v2/instructions")
.and_then(|v| v.as_array())
.ok_or(FetchError::NotFound)?;
for instruction in instructions {
if instruction.get("type").and_then(|t| t.as_str()) != Some("TimelineAddEntries") {
continue;
}
let entries = instruction
.get("entries")
.and_then(|e| e.as_array())
.ok_or(FetchError::NotFound)?;
let wanted = format!("tweet-{id}");
for entry in entries {
if entry.get("entryId").and_then(|i| i.as_str()) == Some(wanted.as_str()) {
let result = entry
.pointer("/content/itemContent/tweet_results/result")
.ok_or(FetchError::NotFound)?;
return normalize_tweet_result(result);
}
}
}
Err(FetchError::NotFound)
}
/// Unwraps TweetTombstone/TweetUnavailable errors, the
/// TweetWithVisibilityResults wrapper and retweets, returning the
/// `{core, legacy, ...}` tweet object.
fn normalize_tweet_result(result: &Value) -> Result<Value, FetchError> {
match result.get("__typename").and_then(|t| t.as_str()) {
Some("TweetTombstone") => {
let text = result
.pointer("/tombstone/text/text")
.and_then(|t| t.as_str())
.unwrap_or("tweet is unavailable");
log::warn!("twitter auth fetch: tombstone: {text}");
return Err(FetchError::NotFound);
}
Some("TweetUnavailable") => {
let reason = result
.get("reason")
.and_then(|r| r.as_str())
.unwrap_or("unknown");
log::warn!("twitter auth fetch: tweet unavailable: {reason}");
return Err(FetchError::NotFound);
}
_ => {}
}
// TweetWithVisibilityResults (e.g. limited replies) nests the real tweet.
let tweet = result.get("tweet").unwrap_or(result);
// A retweet's media lives on the original tweet.
if let Some(original) = tweet.pointer("/legacy/retweeted_status_result/result") {
return Ok(original.clone());
}
Ok(tweet.clone())
}
/// Maps a GraphQL `{core, legacy, ...}` tweet onto the syndication JSON
/// shape [`Tweet::from_syndication_json`] parses, so the existing text /
/// media handling (t.co expansion, `name=orig`, mp4 variant) is reused.
fn to_syndication_shape(tweet: &Value) -> Option<Value> {
let legacy = tweet.get("legacy")?;
let user = tweet.pointer("/core/user_results/result/legacy")?;
Some(json!({
"id_str": legacy.get("id_str"),
"text": legacy.get("full_text"),
"user": {
"name": user.get("name"),
"screen_name": user.get("screen_name"),
},
"possibly_sensitive": legacy.get("possibly_sensitive"),
"display_text_range": legacy.get("display_text_range"),
"entities": legacy.get("entities"),
"mediaDetails": legacy.pointer("/extended_entities/media"),
}))
}
#[cfg(test)]
mod tests {
use super::*;
fn tweet_result() -> Value {
json!({
"__typename": "Tweet",
"core": {
"user_results": {
"result": {
"legacy": { "name": "Display Name", "screen_name": "nsfw_author" }
}
}
},
"legacy": {
"id_str": "2083868672721039569",
"full_text": "nsfw content https://t.co/abc123",
"display_text_range": [0, 12],
"possibly_sensitive": true,
"entities": {
"urls": [
{ "url": "https://t.co/abc123", "expanded_url": "https://example.com/x" }
]
},
"extended_entities": {
"media": [
{
"type": "photo",
"media_url_https": "https://pbs.twimg.com/media/nsfw.jpg",
"original_info": { "width": 1200, "height": 800 }
},
{
"type": "video",
"media_url_https": "https://pbs.twimg.com/thumb.jpg",
"video_info": {
"variants": [
{ "content_type": "application/x-mpegURL", "url": "https://x.com/pl.m3u8" },
{ "content_type": "video/mp4", "url": "https://video.twimg.com/nsfw.mp4" }
]
}
}
]
}
}
})
}
fn conversation(tweet: Value) -> Value {
json!({
"data": {
"threaded_conversation_with_injections_v2": {
"instructions": [
{ "type": "TimelineAddEntries", "entries": [
{ "entryId": "tweet-2083868672721039569",
"content": { "itemContent": { "tweet_results": { "result": tweet } } } }
]}
]
}
}
})
}
#[test]
fn parses_graphql_tweet_into_fetched() {
let json = conversation(tweet_result());
let result = parse_tweet_result(&json, "2083868672721039569").unwrap();
let shape = to_syndication_shape(&result).unwrap();
let tweet = Tweet::from_syndication_json(&shape.to_string()).unwrap();
let fetched: crate::site::Fetched = tweet.into();
assert!(fetched.sensitive);
assert_eq!(fetched.media.len(), 2);
match &fetched.media[0] {
crate::media::Media::Illustration { url, .. } => {
assert_eq!(url, "https://pbs.twimg.com/media/nsfw.jpg?name=orig");
}
other => panic!("expected illustration, got {other:?}"),
}
match &fetched.media[1] {
crate::media::Media::Video { url, .. } => {
assert_eq!(url, "https://video.twimg.com/nsfw.mp4");
}
other => panic!("expected video, got {other:?}"),
}
assert_eq!(fetched.source_url, "https://x.com/nsfw_author/status/2083868672721039569");
// display_text_range cuts the trailing t.co link.
assert_eq!(fetched.title, "nsfw content");
}
#[test]
fn unwraps_retweet_to_original() {
let original = tweet_result();
let mut rt = tweet_result();
rt["legacy"]["retweeted_status_result"] = json!({ "result": original });
let json = conversation(rt);
let result = parse_tweet_result(&json, "2083868672721039569").unwrap();
assert!(result.pointer("/legacy/retweeted_status_result").is_none());
assert_eq!(result.pointer("/legacy/id_str").unwrap(), "2083868672721039569");
}
#[test]
fn error_response_maps_to_not_found() {
let json = json!({ "errors": [{ "message": "NsfwLoggedOut" }] });
assert!(matches!(
parse_tweet_result(&json, "1"),
Err(FetchError::NotFound)
));
}
#[test]
fn missing_entry_maps_to_not_found() {
let json = conversation(json!({ "__typename": "Tweet" }));
assert!(matches!(
parse_tweet_result(&json, "999"),
Err(FetchError::NotFound)
));
}
#[test]
fn tombstone_maps_to_not_found() {
let tombstone = json!({
"__typename": "TweetTombstone",
"tombstone": { "text": { "text": "Age-restricted adult content" } }
});
let json = conversation(tombstone);
assert!(matches!(
parse_tweet_result(&json, "2083868672721039569"),
Err(FetchError::NotFound)
));
}
#[test]
fn visibility_wrapper_unwraps() {
let inner = tweet_result();
let wrapped = json!({ "__typename": "TweetWithVisibilityResults", "tweet": inner });
let json = conversation(wrapped);
let result = parse_tweet_result(&json, "2083868672721039569").unwrap();
assert_eq!(result.get("__typename").unwrap(), "Tweet");
}
}
@@ -0,0 +1,552 @@
use super::model;
use crate::media::Media;
use crate::site::{FetchError, Fetched};
use html_escape::encode_text;
use regex::Regex;
use std::sync::LazyLock;
pub static PATTERN: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"^(?:https?://)?(?:www\.|mobile\.)?(?:x|twitter|fixvx|vxtwitter|fixupx|fxtwitter)\.com/[^.]+/status/(\d+)").unwrap()
});
pub fn enabled() -> bool {
true
}
pub async fn fetch_from_url(url: &str) -> Result<Fetched, FetchError> {
let id = PATTERN
.captures(url)
.and_then(|caps| caps.get(1))
.map(|m| m.as_str())
.ok_or(FetchError::NotFound)?;
match fetch(id).await {
Ok(tweet) => Ok(tweet.into()),
// Syndication withholds NSFW/age-restricted tweets (empty `{}`).
// Retry as the logged-in user when TWITTER_AUTH_TOKEN is set;
// otherwise degrade to an empty result (the bot replies
// "No media found").
Err(FetchError::Sensitive) => {
if super::auth::enabled() {
match super::auth::fetch(id).await {
Ok(tweet) => Ok(tweet.into()),
Err(e) => {
log::warn!("twitter auth fallback failed for {id}: {e}");
Ok(empty_fetched(url))
}
}
} else {
log::info!(
"tweet {id} is sensitive; set TWITTER_AUTH_TOKEN to fetch NSFW media"
);
Ok(empty_fetched(url))
}
}
Err(e) => Err(e),
}
}
/// A Fetched with no media for withheld tweets: the bot replies
/// "No media found" and moves on instead of erroring.
fn empty_fetched(url: &str) -> Fetched {
Fetched {
source_url: url.to_string(),
caption: url.to_string(),
title: String::new(),
media: vec![],
sensitive: true,
render_data: None,
_keep_alive: None,
}
}
/// Fetches a tweet from the syndication endpoint. Deleted/blocked tweets
/// surface as `FetchError::NotFound`.
pub async fn fetch(id: &str) -> Result<Tweet, FetchError> {
let id_num = id.parse::<u64>().map_err(|_| FetchError::NotFound)?;
let response = crate::site::CLIENT
.get(format!(
"https://cdn.syndication.twimg.com/tweet-result?id={id}&lang=en&token={}",
syndication_token(id_num)
))
.send()
.await?;
if !response.status().is_success() {
return Err(FetchError::NotFound);
}
let text = response.text().await?;
// Deleted tweets answer with {"errors": [...]} instead of a tweet.
if serde_json::from_str::<serde_json::Value>(&text)
.map(|v| v.get("errors").is_some())
.unwrap_or(false)
{
return Err(FetchError::NotFound);
}
// NSFW / age-restricted tweets exist but are served as an empty `{}` —
// they surface as FetchError::Sensitive so the caller can retry as a
// logged-in user.
if serde_json::from_str::<serde_json::Value>(&text)
.map(|v| v.get("id_str").is_none())
.unwrap_or(false)
{
return Err(FetchError::Sensitive);
}
Tweet::from_syndication_json(&text).map_err(FetchError::Json)
}
/// The syndication token: JS `((id / 1e15) * PI).toString(36)` (the
/// `replace('0.','')` is a no-op for realistic tweet ids). The endpoint
/// currently serves public tweets regardless of the token; the formula is
/// kept for parity with the known-good client behavior.
fn syndication_token(id: u64) -> String {
let value = (id as f64 / 1e15) * std::f64::consts::PI;
let integer = value.trunc() as u64;
let mut fraction = value.fract();
let mut digits = String::new();
if integer == 0 {
digits.push('0');
} else {
let mut n = integer;
let mut buf = Vec::new();
while n > 0 {
buf.push(char::from_digit((n % 36) as u32, 36).unwrap());
n /= 36;
}
digits.extend(buf.into_iter().rev());
}
digits.push('.');
for _ in 0..10 {
fraction *= 36.0;
let digit = fraction.trunc() as u32;
digits.push(char::from_digit(digit.min(35), 36).unwrap());
fraction -= digit as f64;
if fraction == 0.0 {
break;
}
}
digits
}
#[derive(Debug)]
pub struct Tweet {
id: String,
text: String,
author: String,
author_id: String,
media: Vec<Media>,
sensitive: bool,
}
impl Tweet {
fn url(&self) -> String {
format!("{}/status/{}", self.author_url(), self.id)
}
fn author_url(&self) -> String {
format!("https://x.com/{}", self.author_id)
}
pub fn caption(&self) -> String {
format!(
"{url}\n<a href=\"{author_url}\">{author}</a>: {text}",
url = self.url(),
author_url = self.author_url(),
author = encode_text(&self.author),
text = encode_text(&self.text),
)
}
pub fn from_syndication_json(raw_json: &str) -> Result<Self, serde_json::Error> {
let json: model::SyndicationTweet = serde_json::from_str(raw_json)?;
let id = json.id_str;
// Strip the appended media short link first, then expand the remaining
// t.co short links (the user's own URLs) to their real destinations.
let text = expand_links(
&strip_trailing_short_links(&json.text, json.display_text_range),
&json.entities.urls,
);
// `name` is the display name, `screen_name` the handle (Python's
// vxtwitter mapping: author = display name, author_id = handle).
let author = json.user.name;
let author_id = json.user.screen_name;
let mut media = vec![];
for item in json.media_details {
match item.media_type.as_str() {
"photo" => media.push(Media::Illustration {
title: None,
url: original_twimg_url(&item.media_url_https),
thumbnail_url: None,
// The param-less base URL is a reduced-size variant;
// used as the fallback when the original is too large.
fallback_url: Some(item.media_url_https.clone()),
}),
"video" => media.push(Media::Video {
title: None,
url: mp4_variant(&item),
thumbnail_url: item.media_url_https,
}),
"animated_gif" => media.push(Media::Animated {
title: None,
url: mp4_variant(&item),
thumbnail_url: item.media_url_https,
}),
_ => {}
}
}
let sensitive = json.possibly_sensitive.unwrap_or(false);
Ok(Self {
id,
text,
author,
author_id,
media,
sensitive,
})
}
}
/// The raw syndication `text` ends with the appended media short link
/// (" https://t.co/wmI8McgXul"). `display_text_range` marks the visible text;
/// a regex strips any remaining trailing t.co link when the range is absent
/// or a tweet ends in a URL short link.
///
/// X reports these indices in Unicode **code points**, not UTF-16 units
/// (verified against GraphQL responses containing emoji: cutting an emoji
/// tweet by UTF-16 units silently drops the character after the emoji).
fn strip_trailing_short_links(text: &str, display_text_range: Option<[usize; 2]>) -> String {
let mut out = match display_text_range {
Some([start, end]) if start < end => {
text.chars().skip(start).take(end - start).collect()
}
_ => text.to_string(),
};
while TRAILING_TCO.is_match(&out) {
out = TRAILING_TCO.replace(&out, "").into_owned();
}
out
}
/// Trailing Twitter short link, optionally preceded by whitespace.
static TRAILING_TCO: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"\s*https?://t\.co/[A-Za-z0-9]+$").unwrap()
});
/// Replaces every t.co short link that has an entity mapping with its
/// expanded URL. Short links without a mapping stay untouched.
fn expand_links(text: &str, urls: &[model::SyndicationEntityUrl]) -> String {
let mut out = text.to_string();
for entity in urls {
if let Some(expanded) = &entity.expanded_url {
out = out.replace(&entity.url, expanded);
}
}
out
}
/// pbs.twimg.com serves a reduced default size without size params; `name=orig`
/// returns the original file (fxtwitter used to hand out the original
/// directly, the syndication API does not). Non-twimg URLs pass through
/// unchanged.
fn original_twimg_url(url: &str) -> String {
if url.starts_with("https://pbs.twimg.com/")
&& (url.ends_with(".jpg") || url.ends_with(".png"))
{
format!("{url}?name=orig")
} else {
url.to_string()
}
}
fn mp4_variant(item: &model::SyndicationMedia) -> String {
item.video_info
.as_ref()
.and_then(|info| {
info.variants
.iter()
.find(|variant| variant.content_type == "video/mp4")
})
.map(|variant| variant.url.clone())
.unwrap_or_else(|| item.media_url_https.clone())
}
impl From<Tweet> for Fetched {
fn from(tweet: Tweet) -> Self {
let url = tweet.url();
let author_url = tweet.author_url();
let render_data = Some(crate::site::RenderData {
url: url.clone(),
author: encode_text(&tweet.author).into_owned(),
author_url: author_url.clone(),
title: encode_text(&tweet.text).into_owned(),
tags: String::new(),
});
Fetched {
source_url: url,
caption: tweet.caption(),
title: tweet.text.clone(),
media: tweet.media,
sensitive: tweet.sensitive,
render_data,
_keep_alive: None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture(media_details: serde_json::Value) -> serde_json::Value {
serde_json::json!({
"__typename": "Tweet",
"id_str": "861627479294746624",
"text": "a & b <c>",
"user": { "name": "Display Name", "screen_name": "author_handle" },
"possibly_sensitive": true,
"mediaDetails": media_details
})
}
#[test]
fn pattern_matches_all_domains() {
for url in [
"https://x.com/user/status/1234567890",
"https://twitter.com/user/status/1234567890",
"https://mobile.twitter.com/user/status/1234567890",
"https://www.x.com/user/status/1234567890",
"https://fxtwitter.com/user/status/1234567890",
"https://fixupx.com/user/status/1234567890",
"https://fixvx.com/user/status/1234567890",
"https://vxtwitter.com/user/status/1234567890",
] {
let caps = PATTERN.captures(url).unwrap_or_else(|| panic!("{url}"));
assert_eq!(caps.get(1).unwrap().as_str(), "1234567890");
}
}
#[test]
fn pattern_rejects_non_tweet_urls() {
for url in [
"https://x.com/user",
"https://x.com/user/status/abc",
"https://bsky.app/profile/u/post/3xxxx",
"https://pixiv.net/artworks/123",
"https://example.com/x.com/user/status/123",
] {
assert!(!PATTERN.is_match(url), "{url}");
}
}
#[test]
fn syndication_json_converts_to_fetched() {
let raw = fixture(serde_json::json!([
{ "type": "photo", "media_url_https": "https://pbs.twimg.com/media/photo.jpg" },
{
"type": "video",
"media_url_https": "https://pbs.twimg.com/thumb.jpg",
"video_info": {
"variants": [
{ "content_type": "application/x-mpegURL", "url": "https://x.com/pl.m3u8" },
{ "content_type": "video/mp4", "url": "https://video.twimg.com/v.mp4" }
]
}
}
]));
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
let fetched: Fetched = tweet.into();
assert_eq!(
fetched.source_url,
"https://x.com/author_handle/status/861627479294746624"
);
assert_eq!(fetched.title, "a & b <c>");
assert!(fetched.sensitive);
assert_eq!(fetched.media.len(), 2);
match &fetched.media[0] {
Media::Illustration { url, .. } => {
// Photo URL is rewritten to request the original file.
assert_eq!(
url,
"https://pbs.twimg.com/media/photo.jpg?name=orig"
);
}
other => panic!("expected illustration, got {other:?}"),
}
match &fetched.media[1] {
Media::Video { url, thumbnail_url, .. } => {
assert_eq!(url, "https://video.twimg.com/v.mp4");
assert_eq!(thumbnail_url, "https://pbs.twimg.com/thumb.jpg");
}
other => panic!("expected video, got {other:?}"),
}
assert!(
fetched
.caption
.contains("<a href=\"https://x.com/author_handle\">Display Name</a>: a &amp; b &lt;c&gt;"),
"caption: {}",
fetched.caption
);
}
#[test]
fn syndication_text_only_has_no_media() {
let raw = fixture(serde_json::json!([]));
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
let fetched: Fetched = tweet.into();
assert!(fetched.media.is_empty());
}
#[test]
fn syndication_gif_maps_to_animated() {
let raw = fixture(serde_json::json!([
{
"type": "animated_gif",
"media_url_https": "https://pbs.twimg.com/g.jpg",
"video_info": {
"variants": [{ "content_type": "video/mp4", "url": "https://video.twimg.com/g.mp4" }]
}
}
]));
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert!(matches!(&tweet.media[0], Media::Animated { .. }));
}
#[test]
fn syndication_text_strips_trailing_media_short_link() {
// Real syndication shape: the media short link sits after the visible
// text, and display_text_range marks where it begins.
let raw = serde_json::json!({
"__typename": "Tweet",
"id_str": "1",
"text": "hello world https://t.co/abc123",
"display_text_range": [0, 11],
"user": { "name": "N", "screen_name": "h" },
"mediaDetails": []
});
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert_eq!(tweet.text, "hello world");
assert!(!tweet.caption().contains("t.co"));
}
#[test]
fn syndication_text_strips_trailing_short_link_without_range() {
// No display_text_range: the regex fallback removes the trailing link.
let raw = serde_json::json!({
"__typename": "Tweet",
"id_str": "1",
"text": "hello https://t.co/abc123",
"user": { "name": "N", "screen_name": "h" },
"mediaDetails": []
});
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert_eq!(tweet.text, "hello");
}
#[test]
fn syndication_text_expands_url_entities() {
// Real FloodSocial shape: the user's own link is a t.co short link in
// the text; the entity mapping expands it, the trailing media short
// link is stripped.
let raw = serde_json::json!({
"__typename": "Tweet",
"id_str": "1",
"text": "Test Tweet with @mentionThis $twtr https://t.co/RzmrQ6wAzD #hashtag https://t.co/9r69akA484",
"display_text_range": [0, 67],
"user": { "name": "N", "screen_name": "h" },
"entities": {
"urls": [{
"url": "https://t.co/RzmrQ6wAzD",
"expanded_url": "http://bit.ly/2pUk4be",
"display_url": "bit.ly/2pUk4be"
}]
},
"mediaDetails": []
});
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert_eq!(
tweet.text,
"Test Tweet with @mentionThis $twtr http://bit.ly/2pUk4be #hashtag"
);
assert!(!tweet.caption().contains("t.co"));
}
#[test]
fn syndication_text_keeps_unmapped_short_links() {
// No entity mapping for the embedded link: it stays as-is. Only the
// trailing media link is stripped.
let raw = serde_json::json!({
"__typename": "Tweet",
"id_str": "1",
"text": "check https://t.co/abc123 #tag https://t.co/def456",
"display_text_range": [0, 30],
"user": { "name": "N", "screen_name": "h" },
"mediaDetails": []
});
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert_eq!(tweet.text, "check https://t.co/abc123 #tag");
}
#[test]
fn syndication_text_utf16_display_range_keeps_multibyte() {
// display_text_range is in UTF-16 units; a Japanese text must not be
// sliced by UTF-8 bytes.
let text = "コミティア落ちたので、明日は行きません。🙏ごめんなさい";
let units: Vec<u16> = text.encode_utf16().collect();
assert_eq!(units.len(), 28);
let raw = serde_json::json!({
"__typename": "Tweet",
"id_str": "1",
"text": text,
"display_text_range": [0, 28],
"user": { "name": "N", "screen_name": "h" },
"mediaDetails": []
});
let tweet = Tweet::from_syndication_json(&raw.to_string()).unwrap();
assert_eq!(tweet.text, text, "full text kept intact");
}
#[test]
fn original_twimg_url_rewrites_photo_urls() {
assert_eq!(
original_twimg_url("https://pbs.twimg.com/media/C_UdnvPUwAE3Dnn.jpg"),
"https://pbs.twimg.com/media/C_UdnvPUwAE3Dnn.jpg?name=orig"
);
assert_eq!(
original_twimg_url("https://pbs.twimg.com/media/abc.png"),
"https://pbs.twimg.com/media/abc.png?name=orig"
);
// Non-twimg URLs (videos, animated gifs) pass through unchanged.
assert_eq!(
original_twimg_url("https://video.twimg.com/v.mp4"),
"https://video.twimg.com/v.mp4"
);
assert_eq!(
original_twimg_url("https://pbs.twimg.com/media/abc.webp"),
"https://pbs.twimg.com/media/abc.webp"
);
}
#[test]
fn syndication_token_matches_js_formula() {
// JS: ((861627479294746624 / 1e15) * PI).toString(36) == "236.vrsocvda"
let token = syndication_token(861627479294746624);
assert!(token.starts_with("236.v"), "got {token}");
}
#[tokio::test]
async fn live_fetch_with_photos() {
let fetched = fetch("861627479294746624").await.unwrap();
assert_eq!(fetched.media.len(), 4);
}
#[tokio::test]
async fn live_fetch_text_only() {
let fetched = fetch("1992471125734142256").await.unwrap();
assert!(fetched.media.is_empty());
}
#[tokio::test]
async fn live_fetch_deleted_tweet_is_not_found() {
// Deleted tweet: the syndication endpoint answers with errors.
let result = fetch("0").await;
assert!(matches!(result, Err(FetchError::NotFound)), "got {result:?}");
}
}
+5
View File
@@ -0,0 +1,5 @@
mod auth;
mod interface;
mod model;
pub use interface::{PATTERN, Tweet, enabled, fetch_from_url};
+62
View File
@@ -0,0 +1,62 @@
use serde::Deserialize;
/// Response shape of the syndication endpoint
/// (`cdn.syndication.twimg.com/tweet-result`).
#[derive(Deserialize, Debug)]
pub struct SyndicationTweet {
pub id_str: String,
pub text: String,
pub user: SyndicationUser,
#[serde(default)]
pub possibly_sensitive: Option<bool>,
/// Visible-text span; the raw `text` field has the appended media short
/// link after it. Indices are Unicode code points (not UTF-16 units).
#[serde(default, rename = "display_text_range")]
pub display_text_range: Option<[usize; 2]>,
#[serde(default)]
pub entities: SyndicationEntities,
#[serde(default, rename = "mediaDetails")]
pub media_details: Vec<SyndicationMedia>,
}
#[derive(Deserialize, Debug, Default)]
pub struct SyndicationEntities {
#[serde(default)]
pub urls: Vec<SyndicationEntityUrl>,
}
/// A URL entity: `url` is the t.co short link as it appears in the text,
/// `expanded_url` the real destination.
#[derive(Deserialize, Debug)]
pub struct SyndicationEntityUrl {
pub url: String,
#[serde(default)]
pub expanded_url: Option<String>,
}
#[derive(Deserialize, Debug)]
pub struct SyndicationUser {
pub name: String,
pub screen_name: String,
}
#[derive(Deserialize, Debug)]
pub struct SyndicationMedia {
#[serde(rename = "type")]
pub media_type: String,
pub media_url_https: String,
#[serde(default)]
pub video_info: Option<SyndicationVideoInfo>,
}
#[derive(Deserialize, Debug)]
pub struct SyndicationVideoInfo {
#[serde(default)]
pub variants: Vec<SyndicationVariant>,
}
#[derive(Deserialize, Debug)]
pub struct SyndicationVariant {
pub content_type: String,
pub url: String,
}
+25
View File
@@ -0,0 +1,25 @@
[package]
name = "xmedia-bot"
version = "1.0.7"
edition = "2024"
[dependencies]
teloxide = { version = "0.17", features = ["webhooks-axum", "macros"] }
tokio = { version = "1.40", features = ["rt-multi-thread", "macros"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
log = "0.4"
pretty_env_logger = "0.5"
dotenv = "0.15"
url = "2.5.2"
regex = "1.12"
html-escape = "0.2"
rusqlite = { version = "0.32", features = ["bundled"] }
rand = "0.8"
tempfile = "3"
parking_lot = "0.12"
png = "0.18"
zune-jpeg = "0.5"
fast_image_resize = "6"
jpeg-encoder = "0.7"
x-media = { path = "../x-media" }
+73
View File
@@ -0,0 +1,73 @@
//! Central env handling. The only other places that read env are
//! `Bot::from_env` (TELOXIDE_TOKEN) and x-media (PIXIV_REFRESH_TOKEN).
use std::env;
use std::net::IpAddr;
use std::time::Duration;
pub struct Config {
/// BOT_ADMIN: comma-separated ints; empty when unset.
pub admin_ids: Vec<i64>,
/// EDIT_MESSAGE_TTL_SECONDS, default 86400 (24h).
pub edit_message_ttl: Duration,
/// LINK_CACHE_TTL_SECONDS, default 604800 (7 days).
pub link_cache_ttl: Duration,
// Webhook settings (moved out of main; names/defaults unchanged).
pub webhook_enabled: bool,
pub webhook_url: Option<url::Url>,
pub webhook_listen: Option<IpAddr>,
pub webhook_port: Option<u16>,
pub webhook_cert: Option<String>,
pub webhook_secret_token: Option<String>,
}
impl Config {
pub fn load() -> Config {
let admin_ids = env::var("BOT_ADMIN")
.ok()
.map(|s| {
s.split(',')
.filter_map(|part| part.trim().parse::<i64>().ok())
.collect()
})
.unwrap_or_default();
let edit_message_ttl = env::var("EDIT_MESSAGE_TTL_SECONDS")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.map(Duration::from_secs)
.unwrap_or(Duration::from_secs(86400));
let link_cache_ttl = env::var("LINK_CACHE_TTL_SECONDS")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.map(Duration::from_secs)
.unwrap_or(Duration::from_secs(7 * 24 * 3600));
let webhook_enabled = env::var("WEBHOOK")
.is_ok_and(|v| matches!(v.to_lowercase().as_str(), "true" | "yes" | "1"));
let webhook_url = env::var("WEBHOOK_URL").ok().and_then(|s| s.parse().ok());
let webhook_listen = env::var("WEBHOOK_LISTEN").ok().and_then(|s| s.parse().ok());
let webhook_port = env::var("WEBHOOK_PORT").ok().and_then(|s| s.parse().ok());
// Empty strings count as unset (e.g. `-e WEBHOOK_CERT=` to disable a
// value that would otherwise come from `.env`).
let webhook_cert = env::var("WEBHOOK_CERT")
.ok()
.filter(|s| !s.is_empty());
let webhook_secret_token = env::var("WEBHOOK_SECRET_TOKEN")
.ok()
.filter(|s| !s.is_empty());
Config {
admin_ids,
edit_message_ttl,
link_cache_ttl,
webhook_enabled,
webhook_url,
webhook_listen,
webhook_port,
webhook_cert,
webhook_secret_token,
}
}
}
+777
View File
@@ -0,0 +1,777 @@
use crate::config::Config;
use crate::link_cache::{CachedMediaKind, CachedPost, LinkCache};
use crate::queue::PersistentTaskQueue;
use crate::send::{self, MediaItemPayload, Task};
use crate::state::{ChatData, ChatStore, unix_now};
use std::collections::HashSet;
use std::sync::LazyLock;
use teloxide::prelude::*;
use tokio::sync::Semaphore;
use teloxide::types::{
CallbackQuery, ChatAction, ChatId, ChatKind, InlineQuery, InlineQueryResult,
InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, Message,
MessageEntityKind, MessageId, ParseMode, Recipient, ReplyParameters,
};
use teloxide::utils::command::BotCommands;
use teloxide::RequestError;
use x_media::media::Media;
pub static CHAT_STORE: LazyLock<ChatStore> = LazyLock::new(|| {
ChatStore::open("data/task_queue.db").expect("failed to open chat store")
});
pub static TASK_QUEUE: LazyLock<PersistentTaskQueue> =
LazyLock::new(|| PersistentTaskQueue::new("data/task_queue.db"));
pub static LINK_CACHE: LazyLock<LinkCache> =
LazyLock::new(|| LinkCache::open("data/task_queue.db"));
pub static CONFIG: LazyLock<Config> = LazyLock::new(Config::load);
/// Cap on concurrent per-URL processing. teloxide dispatches updates to a
/// per-chat worker that handles them sequentially, so a batch-forward of many
/// messages would otherwise be processed one at a time (fetch + send each,
/// roughly a second per message). Moving the work into spawned tasks trades
/// per-chat reply ordering for throughput; the semaphore bounds how many run
/// at once so a big burst cannot hammer Telegram's rate limits.
static URL_TASKS: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(8));
#[derive(BotCommands, Clone)]
#[command(rename_rule = "snake_case", description = "")]
enum Command {
#[command(description = "")]
Start,
#[command(description = "")]
Help,
#[command(description = "", parse_with = "split")]
SetForwardChannel(String),
#[command(description = "")]
RemoveForwardChannel,
#[command(description = "")]
EditBeforeForward,
#[command(description = "", parse_with = "split")]
SetTemplate(String),
#[command(description = "")]
BotDict,
#[command(description = "", parse_with = "split")]
SetFormat(String),
}
async fn reply<T>(bot: Bot, message: Message, text: T) -> Result<Message, RequestError>
where
T: Into<String>,
{
bot.send_message(message.chat.id, text)
.reply_parameters(ReplyParameters::new(message.id).allow_sending_without_reply())
.await
}
fn now_f64() -> f64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
/// Extracts URL and text-link entities (text + caption), deduped in order.
pub fn extract_urls(message: &Message) -> Vec<String> {
let mut urls = Vec::new();
for entity in message.parse_entities().into_iter().flatten() {
match entity.kind() {
MessageEntityKind::Url => urls.push(entity.text().to_string()),
MessageEntityKind::TextLink { url } => urls.push(url.to_string()),
_ => {}
}
}
for entity in message.parse_caption_entities().into_iter().flatten() {
match entity.kind() {
MessageEntityKind::Url => urls.push(entity.text().to_string()),
MessageEntityKind::TextLink { url } => urls.push(url.to_string()),
_ => {}
}
}
let mut seen = HashSet::new();
urls.retain(|url| seen.insert(url.clone()));
urls
}
/// Edit-before-forward: a reply to the prompt swaps the caption of the first
/// forwarded message. Returns true when the message was consumed as an edit.
async fn edit_message_handler(bot: &Bot, message: &Message) -> bool {
let Some(reply) = message.reply_to_message() else {
return false;
};
let chat_id = message.chat.id.0;
let Some(text) = message.text() else {
return false;
};
let chat_data = CHAT_STORE.get(chat_id).await;
let Some(edit) = chat_data.edit_message.get(&(reply.id.0 as i64)) else {
return false;
};
let Some(first_forward_id) = edit.forward_message_ids.first() else {
return false;
};
let link = format!(
"<a href=\"{0}\">{1}</a>",
edit.url,
html_escape::encode_text(text)
);
let new_text = if edit.template.is_empty() {
link
} else {
chat_data
.template
.get(&edit.template)
.map(|template| template.replace("[]", &link))
.unwrap_or(link)
};
let result = bot
.edit_message_caption(ChatId(chat_id), MessageId(*first_forward_id as i32))
.caption(new_text)
.parse_mode(ParseMode::Html)
.await;
match result {
Ok(_) => log::info!(
"edit-before-forward: caption swapped on message {first_forward_id} for prompt {}",
reply.id.0
),
Err(e) => log::error!("edit_message_caption failed: {e}"),
}
true
}
enum SetForwardChannelError {
EmptyParameter,
NotChannel,
NotAdmin,
NotBotAdmin(RequestError),
NotBotCanPost,
}
async fn set_forward_channel_handler(
bot: &Bot,
message: &Message,
channel: String,
) -> Result<i64, SetForwardChannelError> {
if channel.is_empty() {
return Err(SetForwardChannelError::EmptyParameter);
}
let channel = match channel.parse::<i64>() {
Ok(id) => Recipient::Id(ChatId(id)),
Err(_) => Recipient::ChannelUsername(channel),
};
if let Some(from) = &message.from {
log::info!(
"Set forward channel for {} ({}) to {}",
from.full_name(),
message.chat.id,
channel
);
}
let chat = match bot.get_chat(channel.clone()).await {
Err(e) => {
log::error!("Failed to get channel {}: {}", channel, e);
return Err(SetForwardChannelError::NotBotAdmin(e));
}
Ok(chat) => chat,
};
if !chat.is_channel() {
return Err(SetForwardChannelError::NotChannel);
}
let channel_id = chat.id.0;
match bot.get_chat_administrators(channel.clone()).await {
Err(e) => {
log::error!("Failed to get channel administrators {}: {}", channel, e);
return Err(SetForwardChannelError::NotBotAdmin(e));
}
Ok(admins) => {
if !admins.iter().any(|admin| admin.user.id == message.chat.id) {
return Err(SetForwardChannelError::NotAdmin);
}
let bot_id = bot.get_me().await.expect("Failed get bot id").user.id;
if let Some(bot_admin) = admins.iter().find(|admin| admin.user.id == bot_id)
&& !bot_admin.can_post_messages()
{
return Err(SetForwardChannelError::NotBotCanPost);
}
}
}
Ok(channel_id)
}
async fn execute_command(bot: &Bot, message: &Message, command: Command) -> Result<(), RequestError> {
match command {
Command::Start => {
bot.send_message(message.chat.id, "Hello!").await?;
}
Command::Help => {
bot.send_message(message.chat.id, Command::descriptions().to_string())
.await?;
}
Command::SetForwardChannel(channel) => {
let result = match set_forward_channel_handler(bot, message, channel).await {
Ok(channel_id) => {
let mut chat_data = CHAT_STORE.get(message.chat.id.0).await;
chat_data.forward_channel_id = Some(channel_id);
CHAT_STORE.set(message.chat.id.0, &chat_data).await;
"Add successfully.".to_string()
}
Err(SetForwardChannelError::EmptyParameter) => {
"Receive empty parameter.\nYou should enter a channel id or username".to_string()
}
Err(SetForwardChannelError::NotChannel) => {
"Given id / username is not a channel".to_string()
}
Err(SetForwardChannelError::NotAdmin) => {
"You are not an administrator of the channel".to_string()
}
Err(SetForwardChannelError::NotBotAdmin(e)) => {
e.to_string() + "\nPlease add the bot as an admin to the channel"
}
Err(SetForwardChannelError::NotBotCanPost) => {
"Bot can't post messages to the channel".to_string()
}
};
reply(bot.clone(), message.clone(), result).await?;
}
Command::RemoveForwardChannel => {
let chat_id = message.chat.id.0;
let mut chat_data = CHAT_STORE.get(chat_id).await;
let text = if chat_data.forward_channel_id.is_some() {
chat_data.forward_channel_id = None;
CHAT_STORE.set(chat_id, &chat_data).await;
"Remove successfully.".to_string()
} else {
"No channel to remove.".to_string()
};
reply(bot.clone(), message.clone(), text).await?;
}
Command::EditBeforeForward => {
let chat_id = message.chat.id.0;
let mut chat_data = CHAT_STORE.get(chat_id).await;
let text = if chat_data.forward_channel_id.is_none() {
"Please enable forward channel first.".to_string()
} else if chat_data.edit_before_forward {
chat_data.edit_before_forward = false;
chat_data.edit_message.clear();
CHAT_STORE.set(chat_id, &chat_data).await;
"Disable edit before forward.".to_string()
} else {
chat_data.edit_before_forward = true;
CHAT_STORE.set(chat_id, &chat_data).await;
"Enable edit before forward.".to_string()
};
reply(bot.clone(), message.clone(), text).await?;
}
Command::SetTemplate(name) => {
let chat_id = message.chat.id.0;
let text = match message.reply_to_message() {
None => "Please reply to a message to set as template.".to_string(),
Some(reply) => {
let reply_text = reply.text().unwrap_or_default();
if !reply_text.contains("[]") {
"Please reply to a message with [] to set as template.".to_string()
} else if name.is_empty() {
"Please provide a name for the template.".to_string()
} else {
let mut chat_data = CHAT_STORE.get(chat_id).await;
chat_data
.template
.insert(name, html_escape::encode_text(reply_text).into_owned());
CHAT_STORE.set(chat_id, &chat_data).await;
"Template set.".to_string()
}
}
};
reply(bot.clone(), message.clone(), text).await?;
}
Command::BotDict => {
let chat_data = CHAT_STORE.get(message.chat.id.0).await;
let debug = format!("{chat_data:?}");
let text = html_escape::encode_text(&debug).into_owned();
reply(bot.clone(), message.clone(), text).await?;
}
Command::SetFormat(arg) => {
let chat_id = message.chat.id.0;
let (site, format) = match arg.split_once(char::is_whitespace) {
Some((site, format)) if !format.trim().is_empty() => (site.trim(), format.trim().to_string()),
_ => {
reply(
bot.clone(),
message.clone(),
"Usage: /set_format <site> <format>",
)
.await?;
return Ok(());
}
};
if !["twitter", "bsky", "pixiv"].contains(&site) {
reply(
bot.clone(),
message.clone(),
"Unknown site. Use twitter, bsky or pixiv.",
)
.await?;
return Ok(());
}
let mut chat_data = CHAT_STORE.get(chat_id).await;
chat_data.message_format.insert(site.to_string(), format);
CHAT_STORE.set(chat_id, &chat_data).await;
reply(bot.clone(), message.clone(), "Format set.").await?;
}
}
Ok(())
}
/// For locally produced media (encoded ugoira MP4) the thumbnail URL is a
/// hotlink-protected remote URL Telegram may not fetch; let Telegram generate
/// its own thumbnail instead.
fn thumbnail_for(media: &Media) -> Option<String> {
let url = media.url();
if url.starts_with("http://") || url.starts_with("https://") {
media.thumbnail_url().map(str::to_string)
} else {
None
}
}
fn media_to_payload(media: &Media, sensitive: bool) -> MediaItemPayload {
let fallback_url = media.smaller_url().map(str::to_string);
match media {
// A gif inside a group becomes a video item; a lone gif takes the
// animation path (see url_media).
Media::Illustration { .. } => MediaItemPayload::Photo {
media: media.url().to_string(),
has_spoiler: sensitive,
fallback_url,
file_id: false,
},
Media::Video { .. } => MediaItemPayload::Video {
media: media.url().to_string(),
has_spoiler: sensitive,
thumbnail: thumbnail_for(media),
fallback_url,
file_id: false,
},
Media::Animated { .. } => MediaItemPayload::Video {
media: media.url().to_string(),
has_spoiler: sensitive,
thumbnail: thumbnail_for(media),
fallback_url,
file_id: false,
},
}
}
async fn enqueue_retry(task: Task, delay_seconds: f64) {
let payload = serde_json::to_value(task).expect("task serializes");
let run_after = now_f64() + delay_seconds;
if let Err(e) = TASK_QUEUE.enqueue(payload, run_after).await {
log::error!("failed to enqueue retry: {e}");
}
}
/// Sends a task and handles the outcome: post-send actions on success, retry
/// enqueue on retryable failure, reply + link-cache invalidation on
/// permanent failure (a stale cached file id must not repeat forever).
async fn dispatch_send(bot: Bot, message: &Message, task: &Task, url: &str) {
let result = match task {
Task::SendAnimation { .. } => send::send_animation(&bot, task).await,
Task::SendMediaSequence { .. } => send::send_media_sequence(&bot, task).await,
Task::ForwardMessages { .. } => unreachable!(),
};
match result {
Ok(message_ids) => {
log::info!("sent {} message(s) for {url}", message_ids.len());
send::post_send_actions(&bot, task, message_ids).await;
}
Err(send::SendError::Retryable { delay_seconds, task }) => {
log::info!("send for {url} failed, queued for retry in {delay_seconds:.1}s");
enqueue_retry(task, delay_seconds).await;
let _ = reply(bot, message.clone(), "Send failed. Task queued for retry.").await;
}
Err(send::SendError::Permanent {
message: err_message,
task,
}) => {
send::invalidate_cache(&task).await;
log::error!("send for {url} failed permanently: {err_message}");
let _ = reply(bot, message.clone(), format!("Send failed: {err_message}")).await;
}
}
}
/// Builds the send task from ready-made items, sharing the payload shape
/// between the fresh-fetch and link-cache paths.
#[allow(clippy::too_many_arguments)]
fn build_send_task(
chat_data: &ChatData,
message: &Message,
source_url: String,
caption: String,
items: Vec<MediaItemPayload>,
cache_data: Option<CachedPost>,
) -> Task {
let chat_id = message.chat.id.0;
if items.len() == 1 && matches!(items[0], MediaItemPayload::Animation { .. }) {
Task::SendAnimation {
chat_id,
reply_to_message_id: message.id.0 as i64,
caption,
animation: items.into_iter().next().unwrap(),
source_url,
edit_before_forward: chat_data.edit_before_forward,
forward_channel_id: chat_data.forward_channel_id,
notify_chat_id: Some(chat_id),
notify_message_id: Some(message.id.0 as i64),
cache_data,
}
} else {
Task::SendMediaSequence {
chat_id,
reply_to_message_id: message.id.0 as i64,
caption,
media_batches: send::chunk_media_items(items),
batch_index: 0,
sent_message_ids: vec![],
source_url,
edit_before_forward: chat_data.edit_before_forward,
forward_channel_id: chat_data.forward_channel_id,
notify_chat_id: Some(chat_id),
notify_message_id: Some(message.id.0 as i64),
cache_data,
}
}
}
async fn url_media(bot: Bot, message: &Message, url: &str) {
let chat_id = message.chat.id.0;
if let Err(e) = bot.send_chat_action(ChatId(chat_id), ChatAction::Typing).await {
log::error!("send_chat_action failed: {e}");
}
// Link cache: a post sent before is re-sent from Telegram file ids —
// no source-site request, no download, no upload. Keyed by the
// normalized post id so x.com / fxtwitter / /photo/N variants collide.
if let Some(key) = x_media::site::cache_key(url)
&& let Some(cached) = LINK_CACHE.get(&key, CONFIG.link_cache_ttl).await
{
log::info!("link cache hit for {url}");
let chat_data = CHAT_STORE.get(chat_id).await;
let site = key.split(':').next().unwrap_or("unknown");
let format = chat_data
.message_format
.get(site)
.cloned()
.unwrap_or_default();
let caption = if format.is_empty() {
cached.caption.clone()
} else {
x_media::site::caption_from_fields(
&format,
"",
&cached.url,
&cached.author,
&cached.author_url,
&cached.title,
&cached.tags,
)
};
let items: Vec<MediaItemPayload> = cached
.media
.iter()
.map(|m| match m.kind {
CachedMediaKind::Photo => MediaItemPayload::Photo {
media: m.file_id.clone(),
has_spoiler: cached.sensitive,
fallback_url: None,
file_id: true,
},
CachedMediaKind::Video => MediaItemPayload::Video {
media: m.file_id.clone(),
has_spoiler: cached.sensitive,
thumbnail: None,
fallback_url: None,
file_id: true,
},
CachedMediaKind::Animation => MediaItemPayload::Animation {
media: m.file_id.clone(),
has_spoiler: cached.sensitive,
file_id: true,
},
})
.collect();
let task = build_send_task(
&chat_data,
message,
cached.url.clone(),
caption,
items,
Some(cached),
);
dispatch_send(bot, message, &task, url).await;
return;
}
log::info!("fetching {url}");
match x_media::site::fetch(url).await {
// Unsupported links are ignored silently (Python parity).
Ok(None) => {
log::info!("no site pattern matches {url}; ignoring");
}
// Retries exhausted: notify the user (Rust-only requirement 3).
Err(e) => {
log::error!("fetch {url}: {e}");
let _ = reply(bot, message.clone(), "Failed to fetch media from this link.").await;
}
Ok(Some(fetched)) => {
if fetched.media.is_empty() {
let _ = reply(
bot,
message.clone(),
"No media found or media type is not supported.",
)
.await;
return;
}
let chat_data = CHAT_STORE.get(chat_id).await;
// Per-site caption format override (empty -> built-in caption).
let format = chat_data
.message_format
.get(fetched.site_name())
.cloned()
.unwrap_or_default();
let caption = fetched.caption_with(&format);
// Raw render data for the link cache; the send fills in the
// Telegram file ids and persists the entry.
let cache_data = fetched.render_fields().map(|(author, author_url, title, tags)| {
CachedPost {
url: fetched.source_url.clone(),
caption: fetched.caption.clone(),
title: title.to_string(),
author: author.to_string(),
author_url: author_url.to_string(),
tags: tags.to_string(),
sensitive: fetched.sensitive,
media: vec![],
}
});
let items: Vec<MediaItemPayload> = fetched
.media
.iter()
.map(|media| media_to_payload(media, fetched.sensitive))
.collect();
let task = build_send_task(
&chat_data,
message,
fetched.source_url.clone(),
caption,
items,
cache_data,
);
dispatch_send(bot, message, &task, url).await;
}
}
}
pub async fn message_handler(bot: Bot, message: Message) -> Result<(), RequestError> {
let is_private = matches!(message.chat.kind, ChatKind::Private(_));
let sender = message
.from
.as_ref()
.map(|from| from.full_name())
.unwrap_or_else(|| "unknown".to_string());
let text_preview = message
.text()
.map(|t| if t.len() > 120 { &t[..120] } else { t })
.unwrap_or("<no text>");
log::info!("message from {sender} in {} (private={is_private}): {text_preview}", message.chat.id);
// URL/edit flows only run in private chats; commands run in any chat.
if is_private && edit_message_handler(&bot, &message).await {
return respond(());
}
if let Some(text) = message.text()
&& let Ok(command) = Command::parse(text, "")
{
log::info!("command from {}: {text_preview}", message.chat.id);
execute_command(&bot, &message, command).await?;
return respond(());
}
if is_private {
let urls = extract_urls(&message);
if !urls.is_empty() {
log::info!("extracted {} URL(s): {urls:?}", urls.len());
}
for url in urls {
let bot = bot.clone();
let message = message.clone();
tokio::spawn(async move {
// Held for the whole task; the semaphore is never closed.
let _permit = URL_TASKS.acquire().await.expect("URL semaphore closed");
url_media(bot, &message, &url).await;
});
}
}
respond(())
}
pub async fn inline_query_handler(bot: Bot, query: InlineQuery) -> Result<(), RequestError> {
if query.query.is_empty() {
return respond(());
}
log::info!("inline query: {}", query.query);
match x_media::site::fetch(&query.query).await {
Ok(Some(fetched)) => {
let mut results: Vec<InlineQueryResult> = Vec::new();
for (i, media) in fetched.media.iter().enumerate() {
let id = format!("{i}");
let Some(url) = url::Url::parse(media.url()).ok() else {
continue;
};
let thumbnail = media
.thumbnail_url()
.and_then(|t| url::Url::parse(t).ok())
.unwrap_or_else(|| url.clone());
let caption = fetched.caption.clone();
let result = match media {
Media::Illustration { .. } => {
// Inline photo results have their own (smaller) size
// cap; use the reduced variant when one exists.
let photo_url = media
.smaller_url()
.and_then(|u| url::Url::parse(u).ok())
.unwrap_or_else(|| url.clone());
InlineQueryResult::Photo(
InlineQueryResultPhoto::new(id, photo_url, thumbnail)
.caption(caption)
.parse_mode(ParseMode::Html),
)
}
Media::Video { .. } => InlineQueryResult::Video(
InlineQueryResultVideo::new(
id,
url,
"video/mp4".parse().expect("valid mime"),
thumbnail,
fetched.title.clone(),
)
.caption(caption)
.parse_mode(ParseMode::Html),
),
Media::Animated { .. } => InlineQueryResult::Mpeg4Gif(
InlineQueryResultMpeg4Gif::new(id, url, thumbnail)
.caption(caption)
.parse_mode(ParseMode::Html),
),
};
results.push(result);
}
if !results.is_empty() {
bot.answer_inline_query(query.id, results).await?;
}
}
Ok(None) => {}
Err(e) => log::error!("inline fetch {}: {e}", query.query),
}
respond(())
}
pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<(), RequestError> {
let callback_query_id = query.id;
let data = query.data.clone();
let Some(message) = &query.message else {
return respond(());
};
let chat_id = message.chat().id.0;
let prompt_message_id = message.id().0 as i64;
let ttl_secs = CONFIG.edit_message_ttl.as_secs() as i64;
let mut chat_data = CHAT_STORE.get(chat_id).await;
let edit = chat_data.edit_message.get(&prompt_message_id).cloned();
let Some(edit) = edit else {
log::info!("callback from {}: no edit record for prompt {prompt_message_id}", chat_id);
bot.answer_callback_query(callback_query_id)
.text("Expired")
.await?;
return respond(());
};
// Lazy expiry: a stale record (past the TTL, not yet swept) is dropped.
if edit.created_at + ttl_secs <= unix_now() {
chat_data.edit_message.remove(&prompt_message_id);
CHAT_STORE.set(chat_id, &chat_data).await;
bot.answer_callback_query(callback_query_id)
.text("Expired")
.await?;
return respond(());
}
let Some(data) = data else {
return respond(());
};
log::info!("callback from {} on prompt {prompt_message_id}: {data}", chat_id);
if data == "forward" {
match chat_data.forward_channel_id {
Some(channel_id) => {
let forward_task = Task::ForwardMessages {
from_chat_id: edit.chat_id,
to_chat_id: channel_id,
message_ids: edit.forward_message_ids.clone(),
notify_chat_id: Some(chat_id),
notify_message_id: Some(prompt_message_id),
};
match send::forward_messages(&bot, &forward_task).await {
Ok(()) => {
log::info!(
"forwarded {} message(s) to channel {channel_id}",
edit.forward_message_ids.len()
);
bot.answer_callback_query(callback_query_id)
.text("✅ Forwarded")
.await?;
let _ = bot
.delete_message(ChatId(chat_id), MessageId(prompt_message_id as i32))
.await;
chat_data.edit_message.remove(&prompt_message_id);
CHAT_STORE.set(chat_id, &chat_data).await;
}
Err(send::SendError::Retryable { delay_seconds, task }) => {
log::info!("forward queued for retry in {delay_seconds:.1}s");
enqueue_retry(task, delay_seconds).await;
bot.answer_callback_query(callback_query_id)
.text("Forward queued for retry.")
.await?;
}
Err(send::SendError::Permanent { message, .. }) => {
log::error!("forward failed permanently: {message}");
bot.answer_callback_query(callback_query_id)
.text(format!("Forward failed: {message}"))
.await?;
}
}
}
None => {
log::info!("forward callback without a forward channel set");
bot.answer_callback_query(callback_query_id)
.text("No forward channel set.")
.await?;
}
}
return respond(());
}
if let Some(name) = data.strip_prefix("template|") {
if let Some(template_html) = chat_data.template.get(name).cloned()
&& let Some(first_forward_id) = edit.forward_message_ids.first().copied()
{
// Raw template including the [] placeholder (Python parity).
let _ = bot
.edit_message_caption(ChatId(chat_id), MessageId(first_forward_id as i32))
.caption(template_html)
.parse_mode(ParseMode::Html)
.await;
if let Some(entry) = chat_data.edit_message.get_mut(&prompt_message_id) {
entry.template = name.to_string();
}
CHAT_STORE.set(chat_id, &chat_data).await;
log::info!("template '{name}' applied to prompt {prompt_message_id}");
}
bot.answer_callback_query(callback_query_id).await?;
}
respond(())
}
+230
View File
@@ -0,0 +1,230 @@
//! Persistent cache of successfully sent posts.
//!
//! After a media send succeeds, the raw render data plus the Telegram
//! `file_id`s of the sent items are stored keyed by [`crate::site` cache
//! key]. A repeated link is then answered entirely from local state — no
//! re-fetch of the source site, no re-upload — and no media file is stored
//! on disk (the file ids point at Telegram's servers). Entries expire after
//! [`Config::link_cache_ttl`]; a stale entry is dropped lazily on read and
//! by the periodic prune in `main`.
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum CachedMediaKind {
Photo,
Video,
Animation,
}
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct CachedMedia {
pub kind: CachedMediaKind,
pub file_id: String,
}
/// Everything needed to re-send a post without touching the source site:
/// the canonical URL, pre-escaped caption fields, and the file ids produced
/// by the original successful send.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct CachedPost {
pub url: String,
/// The site's built-in caption (used when the chat has no format
/// override).
pub caption: String,
pub title: String,
pub author: String,
pub author_url: String,
pub tags: String,
pub sensitive: bool,
pub media: Vec<CachedMedia>,
}
/// SQLite-backed cache sharing `data/task_queue.db` with the queue and chat
/// state (same `open_db` pattern: busy timeout, `spawn_blocking` I/O).
pub struct LinkCache {
db_path: String,
}
fn open_db(path: &str) -> rusqlite::Result<Connection> {
let conn = Connection::open(path)?;
conn.busy_timeout(Duration::from_secs(5))?;
Ok(conn)
}
impl LinkCache {
pub fn open(db_path: &str) -> Self {
if let Ok(conn) = Connection::open(db_path)
&& let Err(e) = conn.execute_batch(
"CREATE TABLE IF NOT EXISTS link_cache (url TEXT PRIMARY KEY, \
payload TEXT NOT NULL, created_at REAL NOT NULL);",
)
{
log::error!("failed to initialize link cache schema: {e}");
}
Self {
db_path: db_path.to_string(),
}
}
/// Returns the cached post if present and not expired; a stale entry is
/// removed on the spot.
pub async fn get(&self, key: &str, ttl: Duration) -> Option<CachedPost> {
let db_path = self.db_path.clone();
let key = key.to_string();
let ttl = ttl.as_secs_f64();
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<CachedPost>> {
let conn = open_db(&db_path)?;
let mut stmt =
conn.prepare("SELECT payload, created_at FROM link_cache WHERE url = ?1")?;
let mut rows = stmt.query(params![key])?;
let Some(row) = rows.next()? else {
return Ok(None);
};
let payload: String = row.get(0)?;
let created_at: f64 = row.get(1)?;
if now_f64() - created_at > ttl {
conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?;
return Ok(None);
}
serde_json::from_str(&payload).map(Some).map_err(|e| {
rusqlite::Error::ToSqlConversionFailure(Box::new(e))
})
})
.await
.expect("link cache read worker panicked")
.unwrap_or_else(|e| {
log::error!("link cache read failed: {e}");
None
})
}
pub async fn put(&self, key: &str, post: &CachedPost) {
let db_path = self.db_path.clone();
let key = key.to_string();
let payload = serde_json::to_string(post).expect("cached post serializes");
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute(
"INSERT OR REPLACE INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)",
params![key, payload, now_f64()],
)?;
Ok(())
})
.await
.expect("link cache write worker panicked")
.unwrap_or_else(|e| log::error!("link cache write failed: {e}"));
}
/// Drops an entry (e.g. a cached file id that turned out invalid).
pub async fn remove(&self, key: &str) {
let db_path = self.db_path.clone();
let key = key.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?;
Ok(())
})
.await
.expect("link cache delete worker panicked")
.unwrap_or_else(|e| log::error!("link cache delete failed: {e}"));
}
/// Removes expired entries; returns how many were deleted.
pub async fn prune(&self, ttl: Duration) -> usize {
let db_path = self.db_path.clone();
let cutoff = now_f64() - ttl.as_secs_f64();
tokio::task::spawn_blocking(move || -> rusqlite::Result<usize> {
let conn = open_db(&db_path)?;
conn.execute(
"DELETE FROM link_cache WHERE created_at < ?1",
params![cutoff],
)
})
.await
.expect("link cache prune worker panicked")
.unwrap_or_else(|e| {
log::error!("link cache prune failed: {e}");
0
})
}
}
fn now_f64() -> f64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
#[cfg(test)]
mod tests {
use super::*;
fn entry() -> CachedPost {
CachedPost {
url: "https://x.com/u/status/1".into(),
caption: "cap".into(),
title: "t".into(),
author: "a".into(),
author_url: "au".into(),
tags: "".into(),
sensitive: true,
media: vec![CachedMedia {
kind: CachedMediaKind::Photo,
file_id: "AgAC...".into(),
}],
}
}
#[tokio::test]
async fn put_get_roundtrip() {
let dir = tempfile::tempdir().unwrap();
let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap());
cache.put("twitter:1", &entry()).await;
let got = cache.get("twitter:1", Duration::from_secs(3600)).await;
assert!(got.is_some());
let got = got.unwrap();
assert_eq!(got.url, "https://x.com/u/status/1");
assert_eq!(got.media[0].file_id, "AgAC...");
}
#[tokio::test]
async fn expired_entry_removed_on_read() {
let dir = tempfile::tempdir().unwrap();
let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap());
cache.put("twitter:1", &entry()).await;
// Force the row into the past so a 1s TTL expires it.
{
let conn = Connection::open(dir.path().join("c.db")).unwrap();
conn.execute(
"UPDATE link_cache SET created_at = created_at - 100",
[],
)
.unwrap();
}
assert!(cache.get("twitter:1", Duration::from_secs(1)).await.is_none());
assert!(cache.get("twitter:1", Duration::from_secs(3600)).await.is_none());
}
#[tokio::test]
async fn remove_and_prune() {
let dir = tempfile::tempdir().unwrap();
let cache = LinkCache::open(dir.path().join("c.db").to_str().unwrap());
cache.put("twitter:1", &entry()).await;
cache.put("pixiv:2", &entry()).await;
cache.remove("twitter:1").await;
assert!(cache.get("twitter:1", Duration::from_secs(3600)).await.is_none());
assert!(cache.get("pixiv:2", Duration::from_secs(3600)).await.is_some());
{
let conn = Connection::open(dir.path().join("c.db")).unwrap();
conn.execute("UPDATE link_cache SET created_at = created_at - 100", [])
.unwrap();
}
assert_eq!(cache.prune(Duration::from_secs(1)).await, 1);
assert!(cache.get("pixiv:2", Duration::from_secs(3600)).await.is_none());
}
}
+173
View File
@@ -0,0 +1,173 @@
use dotenv::dotenv;
use teloxide::dptree::endpoint;
use teloxide::stop::StopToken;
use teloxide::types::{ChatId, InputFile, MessageId};
use teloxide::update_listeners::{self, webhooks, UpdateListener};
use teloxide::prelude::*;
use tokio::sync::watch;
use x_media::site;
mod config;
mod handlers;
mod link_cache;
mod photo;
mod queue;
mod send;
mod state;
use handlers::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE};
/// Docker `stop` / `compose down` delivers SIGTERM, which teloxide's ctrlc
/// handler (SIGINT only) never sees — without this the process would die
/// before the graceful shutdown below (admin notice, queue drain). Stopping
/// the token unwinds the dispatcher exactly like Ctrl+C does.
#[cfg(unix)]
fn spawn_sigterm_handler(stop_token: StopToken) {
tokio::spawn(async move {
let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install SIGTERM handler");
sigterm.recv().await;
log::info!("SIGTERM received, stopping the dispatcher");
stop_token.stop();
});
}
#[cfg(not(unix))]
fn spawn_sigterm_handler(_stop_token: StopToken) {}
#[tokio::main]
async fn main() {
dotenv().ok();
pretty_env_logger::init();
log::info!("Starting bot");
let bot = Bot::from_env();
log::info!(
"config: {} admin(s), edit-message TTL {}s",
CONFIG.admin_ids.len(),
CONFIG.edit_message_ttl.as_secs()
);
// Queue worker: handles typed tasks, dead-letters failed sends to the
// task's chat.
TASK_QUEUE
.start(send::handle_task, send::dead_letter_notify)
.await;
log::info!("task queue worker started");
// Pixiv login validation (user request): a failed login notifies the
// admin and disables pixiv for this process.
if site::pixiv::enabled() {
match site::pixiv::validate().await {
Ok(()) => log::info!("pixiv login validated"),
Err(e) => {
log::error!("pixiv login failed: {e}");
if let Some(admin) = CONFIG.admin_ids.first() {
let _ = bot
.send_message(ChatId(*admin), format!("Pixiv login failed: {e}"))
.await;
}
site::pixiv::disable();
}
}
}
// Edit-expiry sweep: clears the prompt's buttons once the record expires.
log::info!("edit-expiry sweep: every 300s, ttl {}", CONFIG.edit_message_ttl.as_secs());
let (stop_tx, stop_rx) = watch::channel(false);
{
let bot = bot.clone();
let mut stop_rx = stop_rx;
tokio::spawn(async move {
loop {
tokio::select! {
_ = stop_rx.changed() => break,
_ = tokio::time::sleep(std::time::Duration::from_secs(300)) => {}
}
let ttl = CONFIG.edit_message_ttl;
let removed = CHAT_STORE.prune_expired(ttl).await;
let pruned = LINK_CACHE.prune(CONFIG.link_cache_ttl).await;
if pruned > 0 {
log::info!("link cache: pruned {pruned} expired entr(ies)");
}
for (chat_id, prompt_message_id) in removed {
// If the prompt was already deleted, this fails with a
// 400 "message to edit not found" — log and ignore.
if let Err(e) = bot
.edit_message_reply_markup(ChatId(chat_id), MessageId(prompt_message_id as i32))
.await
{
log::info!("edit-expiry sweep: prompt message gone: {e}");
}
}
}
});
}
let handler = dptree::entry()
.branch(Update::filter_message().branch(endpoint(handlers::message_handler)))
.branch(Update::filter_inline_query().branch(endpoint(handlers::inline_query_handler)))
.branch(Update::filter_callback_query().branch(endpoint(handlers::callback_query_handler)));
let mut dispatcher = Dispatcher::builder(bot.clone(), handler)
.dependencies(dptree::deps![""])
.enable_ctrlc_handler()
.build();
if CONFIG.webhook_enabled {
log::info!("running in webhook mode");
let url = CONFIG
.webhook_url
.clone()
.expect("WEBHOOK_URL is not set");
// `webhooks::axum` calls set_webhook itself (with the full options,
// secret token included) — no explicit registration here.
let listen = CONFIG.webhook_listen.expect("WEBHOOK_LISTEN is not set");
let port = CONFIG.webhook_port.expect("WEBHOOK_PORT is not set");
let mut options = webhooks::Options::new((listen, port).into(), url);
if let Some(cert) = &CONFIG.webhook_cert {
options = options.certificate(InputFile::file(cert));
}
if let Some(secret) = &CONFIG.webhook_secret_token {
options = options.secret_token(secret.clone());
}
let mut listener = webhooks::axum(bot.clone(), options)
.await
.expect("Failed to create webhook listener");
let stop_token = listener.stop_token();
spawn_sigterm_handler(stop_token);
dispatcher
.dispatch_with_listener(
listener,
LoggingErrorHandler::with_custom_text("Error from update listener"),
)
.await;
} else {
log::info!("running in polling mode");
// Same listener `dispatch()` builds internally — using
// `dispatch_with_listener` just exposes its stop token so SIGTERM can
// unwind the dispatcher before the graceful shutdown below.
let mut listener = update_listeners::polling_default(bot.clone()).await;
let stop_token = listener.stop_token();
spawn_sigterm_handler(stop_token);
dispatcher
.dispatch_with_listener(
listener,
LoggingErrorHandler::with_custom_text("Error from update listener"),
)
.await;
}
// Graceful stop (Ctrl+C / SIGTERM): stop the sweep, notify the admin, drain the queue.
log::info!("Stopping bot");
let _ = stop_tx.send(true);
if let Some(admin) = CONFIG.admin_ids.first() {
let _ = bot.send_message(ChatId(*admin), "Shutting down...").await;
}
TASK_QUEUE.stop().await;
log::info!("Bot stopped");
}
+515
View File
@@ -0,0 +1,515 @@
//! Pure-Rust photo processing: brings a downloaded photo within Telegram's
//! limits (width + height ≤ 10000 px, bytes ≤ 10 MiB) without ffmpeg.
//!
//! Stack: `png` (image-png) for PNG decode/encode, `zune-jpeg` for JPEG
//! decode, `fast_image_resize` (Lanczos3) for downsampling, `jpeg-encoder`
//! for JPEG output.
//!
//! Bit-depth rule: a PNG above 24 bits (32-bit RGBA or 16-bit per channel)
//! is reduced to 24-bit RGB; 24-bit and lower depths are left untouched —
//! gray stays gray, never upconverted. The only upconversion is palette
//! expansion, which resampling requires. Alpha is flattened onto white (JPEG
//! and 24-bit RGB have no alpha channel).
use std::io::Write;
use fast_image_resize as fir;
use tempfile::NamedTempFile;
/// Telegram rejects photos whose width + height exceed this limit
/// (PHOTO_INVALID_DIMENSIONS). Verified empirically: 6300x3730 (sum 10030)
/// fails, 6100x3900 (sum 10000) passes.
pub const PHOTO_MAX_DIMENSION_SUM: u32 = 10000;
/// Resize target with a safety margin so rounding cannot cross the cap.
pub const PHOTO_TARGET_DIMENSION_SUM: u32 = 9900;
/// Upload cap (bytes): files above this are not uploaded; the bot falls back
/// to a smaller media URL instead.
pub const MAX_UPLOAD_BYTES: u64 = 10 * 1024 * 1024;
/// Decode budget (bytes): a larger intermediate buffer is not worth the peak
/// memory; the photo degrades to the smaller URL instead.
const MAX_DECODE_BYTES: u64 = 512 * 1024 * 1024;
/// JPEG output quality (1-100).
const JPEG_QUALITY: u8 = 90;
/// What to upload for a downloaded photo.
pub enum PhotoPrep {
/// Upload this file (the original when within limits, else the processed
/// copy).
Upload(NamedTempFile),
/// The photo cannot be brought within Telegram's limits — the caller
/// falls back to the item's smaller URL.
UseFallback,
}
/// A decoded image buffer tagged with its channel layout.
#[derive(Debug)]
enum PixBuf {
Gray(Vec<u8>),
GrayAlpha(Vec<u8>),
Rgb(Vec<u8>),
}
impl PixBuf {
fn pixel_type(&self) -> fir::PixelType {
match self {
PixBuf::Gray(_) => fir::PixelType::U8,
PixBuf::GrayAlpha(_) => fir::PixelType::U8x2,
PixBuf::Rgb(_) => fir::PixelType::U8x3,
}
}
fn into_vec(self) -> Vec<u8> {
match self {
PixBuf::Gray(v) | PixBuf::GrayAlpha(v) | PixBuf::Rgb(v) => v,
}
}
}
/// Entry point: detects the format and processes the photo if needed.
pub fn prepare_photo(file: NamedTempFile) -> Result<PhotoPrep, String> {
let bytes = std::fs::read(file.path()).map_err(|e| format!("prepare read failed: {e}"))?;
if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
prepare_png(file, bytes)
} else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
prepare_jpeg(file, bytes)
} else {
log::warn!("photo in unsupported format; falling back to smaller media");
Ok(PhotoPrep::UseFallback)
}
}
/// Parses the PNG IHDR (bytes 8..26: signature + length + "IHDR" + width +
/// height + bit depth + color type).
fn parse_png_header(bytes: &[u8]) -> Option<(u32, u32, png::BitDepth, png::ColorType)> {
if !bytes.starts_with(b"\x89PNG\r\n\x1a\n") || bytes.len() < 26 {
return None;
}
let w = u32::from_be_bytes(bytes.get(16..20)?.try_into().ok()?);
let h = u32::from_be_bytes(bytes.get(20..24)?.try_into().ok()?);
let depth = match *bytes.get(24)? {
1 => png::BitDepth::One,
2 => png::BitDepth::Two,
4 => png::BitDepth::Four,
8 => png::BitDepth::Eight,
16 => png::BitDepth::Sixteen,
_ => return None,
};
let color = match *bytes.get(25)? {
0 => png::ColorType::Grayscale,
2 => png::ColorType::Rgb,
3 => png::ColorType::Indexed,
4 => png::ColorType::GrayscaleAlpha,
6 => png::ColorType::Rgba,
_ => return None,
};
Some((w, h, depth, color))
}
/// Output channels of a decoded frame for the given color type (post
/// STRIP_16; palette expands to RGB).
fn output_channels(color: png::ColorType) -> usize {
match color {
png::ColorType::Grayscale => 1,
png::ColorType::GrayscaleAlpha => 2,
png::ColorType::Rgb | png::ColorType::Indexed => 3,
png::ColorType::Rgba => 4,
}
}
/// The 32→24 rule: RGBA (32-bit) becomes RGB with alpha composited onto
/// white; 16-bit per channel was already stripped to 8-bit at decode.
fn flatten_rgba_to_rgb(rgba: &[u8]) -> Vec<u8> {
let mut rgb = Vec::with_capacity(rgba.len() / 4 * 3);
for px in rgba.chunks_exact(4) {
let a = px[3] as u32;
for v in &px[..3] {
// Over white: C = C*a/255 + 255*(1 - a/255).
let v = (*v as u32 * a + 255 * (255 - a)) / 255;
rgb.push(v.min(255) as u8);
}
}
rgb
}
/// Lanczos3 downsampling via fast_image_resize.
fn resize_pix(pix: PixBuf, w: u32, h: u32, nw: u32, nh: u32) -> Result<PixBuf, String> {
let pixel_type = pix.pixel_type();
let src = fir::images::Image::from_vec_u8(w, h, pix.into_vec(), pixel_type)
.map_err(|e| format!("resize input: {e}"))?;
let mut dst = fir::images::Image::new(nw, nh, pixel_type);
let mut resizer = fir::Resizer::new();
let options = fir::ResizeOptions::default()
.resize_alg(fir::ResizeAlg::Convolution(fir::FilterType::Lanczos3));
resizer
.resize(&src, &mut dst, &options)
.map_err(|e| format!("resize: {e}"))?;
let buf = dst.into_vec();
Ok(match pixel_type {
fir::PixelType::U8 => PixBuf::Gray(buf),
fir::PixelType::U8x2 => PixBuf::GrayAlpha(buf),
_ => PixBuf::Rgb(buf),
})
}
fn encode_png(out: &mut Vec<u8>, pix: &PixBuf, w: u32, h: u32) -> Result<(), png::EncodingError> {
let (color, buf) = match pix {
PixBuf::Gray(v) => (png::ColorType::Grayscale, v.as_slice()),
PixBuf::GrayAlpha(v) => (png::ColorType::GrayscaleAlpha, v.as_slice()),
PixBuf::Rgb(v) => (png::ColorType::Rgb, v.as_slice()),
};
let mut encoder = png::Encoder::new(out, w, h);
encoder.set_color(color);
encoder.set_depth(png::BitDepth::Eight);
let mut writer = encoder.write_header()?;
writer.write_image_data(buf)?;
Ok(())
}
fn encode_jpeg(pix: &PixBuf, w: u32, h: u32) -> Result<Vec<u8>, String> {
use jpeg_encoder::{ColorType, Encoder};
let mut out = Vec::new();
let encoder = Encoder::new(&mut out, JPEG_QUALITY);
match pix {
PixBuf::Gray(v) => encoder
.encode(v, w as u16, h as u16, ColorType::Luma)
.map_err(|e| format!("jpeg encode: {e}"))?,
PixBuf::GrayAlpha(v) => {
// JPEG has no alpha: composite onto white, output as gray.
let gray: Vec<u8> = v
.chunks_exact(2)
.map(|px| {
let (g, a) = (px[0] as u32, px[1] as u32);
((g * a + 255 * (255 - a)) / 255).min(255) as u8
})
.collect();
encoder
.encode(&gray, w as u16, h as u16, ColorType::Luma)
.map_err(|e| format!("jpeg encode: {e}"))?;
}
PixBuf::Rgb(v) => encoder
.encode(v, w as u16, h as u16, ColorType::Rgb)
.map_err(|e| format!("jpeg encode: {e}"))?,
}
Ok(out)
}
fn write_temp(bytes: &[u8], ext: &str) -> Result<NamedTempFile, String> {
let mut file = tempfile::Builder::new()
.suffix(&format!(".{ext}"))
.tempfile()
.map_err(|e| format!("temp file failed: {e}"))?;
file.as_file_mut()
.write_all(bytes)
.map_err(|e| format!("temp file write failed: {e}"))?;
Ok(file)
}
fn target_dims(w: u32, h: u32) -> (u32, u32) {
let scale = PHOTO_TARGET_DIMENSION_SUM as f64 / (w + h) as f64;
(
((w as f64 * scale).round() as u32).max(1),
((h as f64 * scale).round() as u32).max(1),
)
}
/// PNG branch: decode (16→8, palette→RGB; gray/GA stay), flatten RGBA to
/// RGB, Lanczos-downscale beyond the dimension cap, encode PNG — a PNG still
/// over the upload cap afterwards becomes JPEG.
fn prepare_png(file: NamedTempFile, bytes: Vec<u8>) -> Result<PhotoPrep, String> {
let (w, h, _bit_depth, color_type) =
parse_png_header(&bytes).ok_or("invalid PNG header")?;
let size_over = bytes.len() as u64 > MAX_UPLOAD_BYTES;
if w + h <= PHOTO_MAX_DIMENSION_SUM && !size_over {
return Ok(PhotoPrep::Upload(file));
}
log::info!("photo {w}x{h} ({_bit_depth:?} {color_type:?}, {} bytes) needs processing", bytes.len());
let channels = output_channels(color_type);
if (w as u64) * (h as u64) * channels as u64 > MAX_DECODE_BYTES {
log::warn!("photo decode buffer exceeds the memory budget; falling back to smaller media");
return Ok(PhotoPrep::UseFallback);
}
// STRIP_16 drops 16-bit to 8-bit (the depth-reduction step); palette
// expands to RGB (resampling requires it). Gray and gray-alpha are kept.
let transforms = match color_type {
png::ColorType::Indexed => png::Transformations::EXPAND,
_ => png::Transformations::STRIP_16,
};
let mut decoder = png::Decoder::new(std::io::Cursor::new(&bytes));
decoder.set_transformations(transforms);
let mut reader = decoder.read_info().map_err(|e| format!("png decode: {e}"))?;
let out_w = reader.info().width;
let out_h = reader.info().height;
let mut buf = vec![
0u8;
reader
.output_buffer_size()
.ok_or("png output buffer size")?
];
reader
.next_frame(&mut buf)
.map_err(|e| format!("png frame: {e}"))?;
let mut pix = match color_type {
png::ColorType::Rgba => PixBuf::Rgb(flatten_rgba_to_rgb(&buf)),
png::ColorType::Grayscale => PixBuf::Gray(buf),
png::ColorType::GrayscaleAlpha => PixBuf::GrayAlpha(buf),
png::ColorType::Rgb | png::ColorType::Indexed => PixBuf::Rgb(buf),
};
let (mut w, mut h) = (out_w, out_h);
if w + h > PHOTO_MAX_DIMENSION_SUM {
let (nw, nh) = target_dims(w, h);
pix = resize_pix(pix, w, h, nw, nh)?;
(w, h) = (nw, nh);
log::info!("downscaled photo to {w}x{h} (Lanczos3)");
}
let mut png_bytes = Vec::new();
encode_png(&mut png_bytes, &pix, w, h).map_err(|e| format!("png encode: {e}"))?;
if png_bytes.len() as u64 <= MAX_UPLOAD_BYTES {
return Ok(PhotoPrep::Upload(write_temp(&png_bytes, "png")?));
}
log::info!("PNG still over the upload cap after processing; transcoding to JPEG");
let jpeg_bytes = encode_jpeg(&pix, w, h)?;
if jpeg_bytes.len() as u64 <= MAX_UPLOAD_BYTES {
return Ok(PhotoPrep::Upload(write_temp(&jpeg_bytes, "jpg")?));
}
log::warn!("processed photo still exceeds the upload cap; falling back to smaller media");
Ok(PhotoPrep::UseFallback)
}
/// JPEG branch: zune-jpeg decode → Lanczos downscale → jpeg-encoder output.
fn prepare_jpeg(file: NamedTempFile, bytes: Vec<u8>) -> Result<PhotoPrep, String> {
let mut decoder = zune_jpeg::JpegDecoder::new(std::io::Cursor::new(&bytes));
// Decodes to RGB by default. Headers first so dimensions are known before
// the (potentially huge) pixel decode.
decoder
.decode_headers()
.map_err(|e| format!("jpeg headers: {e}"))?;
let info = decoder.info().ok_or("jpeg info unavailable")?;
let (w, h) = (info.width as u32, info.height as u32);
let size_over = bytes.len() as u64 > MAX_UPLOAD_BYTES;
if w + h <= PHOTO_MAX_DIMENSION_SUM && !size_over {
return Ok(PhotoPrep::Upload(file));
}
if (w as u64) * (h as u64) * 3 > MAX_DECODE_BYTES {
log::warn!("photo decode buffer exceeds the memory budget; falling back to smaller media");
return Ok(PhotoPrep::UseFallback);
}
let pixels = decoder.decode().map_err(|e| format!("jpeg decode: {e}"))?;
let mut pix = PixBuf::Rgb(pixels);
let (mut w, mut h) = (w, h);
if w + h > PHOTO_MAX_DIMENSION_SUM {
let (nw, nh) = target_dims(w, h);
pix = resize_pix(pix, w, h, nw, nh)?;
(w, h) = (nw, nh);
log::info!("downscaled jpeg to {w}x{h} (Lanczos3)");
}
let jpeg_bytes = encode_jpeg(&pix, w, h)?;
if jpeg_bytes.len() as u64 <= MAX_UPLOAD_BYTES {
return Ok(PhotoPrep::Upload(write_temp(&jpeg_bytes, "jpg")?));
}
log::warn!("processed photo still exceeds the upload cap; falling back to smaller media");
Ok(PhotoPrep::UseFallback)
}
#[cfg(test)]
mod tests {
use super::*;
fn png_header(w: u32, h: u32, depth: u8, color: u8) -> Vec<u8> {
let mut bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR".to_vec();
bytes.extend(w.to_be_bytes());
bytes.extend(h.to_be_bytes());
bytes.extend([depth, color, 0, 0, 0]);
bytes
}
#[test]
fn parses_png_header() {
let bytes = png_header(8979, 5316, 16, 6); // 16-bit RGBA
let (w, h, depth, color) = parse_png_header(&bytes).unwrap();
assert_eq!((w, h), (8979, 5316));
assert_eq!(depth, png::BitDepth::Sixteen);
assert_eq!(color, png::ColorType::Rgba);
let (_, _, depth, color) = parse_png_header(&png_header(10, 10, 8, 0)).unwrap();
assert_eq!(depth, png::BitDepth::Eight);
assert_eq!(color, png::ColorType::Grayscale);
assert!(parse_png_header(b"not a png").is_none());
}
#[test]
fn flatten_rgba_to_rgb_composites_over_white() {
// opaque red stays red
assert_eq!(flatten_rgba_to_rgb(&[255, 0, 0, 255]), vec![255, 0, 0]);
// fully transparent → white
assert_eq!(flatten_rgba_to_rgb(&[0, 0, 0, 0]), vec![255, 255, 255]);
// half alpha red → (255+255)/2 = 255, (0*128 + 255*127)/255 = 127
let out = flatten_rgba_to_rgb(&[255, 0, 0, 128]);
assert_eq!(out[0], 255);
assert_eq!(out[1], 127);
assert_eq!(out[2], 127);
}
#[test]
fn target_dims_stay_under_the_cap() {
for (w, h) in [(12000u32, 7000u32), (10000, 10000), (8979, 5316)] {
let (nw, nh) = target_dims(w, h);
assert!(nw + nh <= PHOTO_MAX_DIMENSION_SUM, "{w}x{h} -> {nw}x{nh}");
assert!(nw >= 1 && nh >= 1);
}
// already within limits: no change expected from the caller, but the
// helper must not produce zero dimensions.
let (nw, nh) = target_dims(500, 400);
assert!(nw >= 1 && nh >= 1);
}
#[test]
fn resize_pix_changes_dimensions() {
// 300x200 RGB → 100x66
let buf: Vec<u8> = (0..300 * 200 * 3).map(|i| (i % 251) as u8).collect();
let resized = resize_pix(PixBuf::Rgb(buf), 300, 200, 100, 66).unwrap();
match resized {
PixBuf::Rgb(v) => assert_eq!(v.len(), 100 * 66 * 3),
other => panic!("expected rgb, got {other:?}"),
}
}
#[test]
fn png_encode_roundtrip_keeps_gray() {
let gray = vec![128u8; 4 * 4];
let mut out = Vec::new();
encode_png(&mut out, &PixBuf::Gray(gray), 4, 4).unwrap();
assert!(!out.is_empty());
let (_, _, depth, color) = parse_png_header(&out).unwrap();
assert_eq!(depth, png::BitDepth::Eight);
assert_eq!(color, png::ColorType::Grayscale);
}
#[test]
fn jpeg_encode_produces_bytes() {
let rgb = vec![128u8; 8 * 8 * 3];
let out = encode_jpeg(&PixBuf::Rgb(rgb), 8, 8).unwrap();
assert!(out.len() > 100);
assert!(out.starts_with(&[0xFF, 0xD8]));
}
/// Writes a small dimension-oversized PNG (9999x2 → sum 10001) to a temp
/// file and runs the full pipeline.
fn run_pipeline(w: u32, h: u32, color: png::ColorType, fill: u8) -> Result<PhotoPrep, String> {
let (channels, data): (usize, Vec<u8>) = match color {
png::ColorType::Grayscale => (1, vec![fill; (w * h) as usize]),
png::ColorType::Rgb => (3, vec![fill; (w * h * 3) as usize]),
_ => unreachable!(),
};
let mut bytes = Vec::new();
{
let mut encoder = png::Encoder::new(&mut bytes, w, h);
encoder.set_color(color);
encoder.set_depth(png::BitDepth::Eight);
let mut writer = encoder.write_header().unwrap();
writer.write_image_data(&data).unwrap();
}
assert_eq!(data.len(), channels * (w * h) as usize);
let mut file = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
std::io::Write::write_all(file.as_file_mut(), &bytes).unwrap();
prepare_photo(file)
}
#[test]
fn pipeline_downscales_oversized_png_keeping_format() {
let prep = run_pipeline(9999, 2, png::ColorType::Rgb, 128).unwrap();
match prep {
PhotoPrep::Upload(file) => {
let out = std::fs::read(file.path()).unwrap();
let (w, h, depth, color) = parse_png_header(&out).unwrap();
assert!(w + h <= PHOTO_MAX_DIMENSION_SUM, "{w}x{h}");
assert_eq!(depth, png::BitDepth::Eight);
assert_eq!(color, png::ColorType::Rgb);
}
PhotoPrep::UseFallback => panic!("over-dimension PNG should have been resized"),
}
}
#[test]
fn pipeline_keeps_gray_png_gray() {
let prep = run_pipeline(9999, 2, png::ColorType::Grayscale, 200).unwrap();
match prep {
PhotoPrep::Upload(file) => {
let out = std::fs::read(file.path()).unwrap();
let (_, _, _, color) = parse_png_header(&out).unwrap();
assert_eq!(color, png::ColorType::Grayscale, "gray must not upconvert");
}
PhotoPrep::UseFallback => panic!("over-dimension gray PNG should have been resized"),
}
}
#[test]
fn pipeline_resizes_oversized_jpeg() {
// Build a small over-dimension JPEG with jpeg-encoder.
let (w, h) = (9999u16, 2u16);
let rgb = vec![90u8; (w as usize) * (h as usize) * 3];
let mut bytes = Vec::new();
{
let encoder = jpeg_encoder::Encoder::new(&mut bytes, 90);
encoder.encode(&rgb, w, h, jpeg_encoder::ColorType::Rgb).unwrap();
}
let mut file = tempfile::Builder::new().suffix(".jpg").tempfile().unwrap();
std::io::Write::write_all(file.as_file_mut(), &bytes).unwrap();
match prepare_photo(file).unwrap() {
PhotoPrep::Upload(file) => {
let out = std::fs::read(file.path()).unwrap();
assert!(out.starts_with(&[0xFF, 0xD8]), "output must stay jpeg");
// 9999x2 downscaled: the buffer length tells the new dims.
assert!(out.len() > 100);
}
PhotoPrep::UseFallback => panic!("over-dimension JPEG should have been resized"),
}
}
#[test]
#[ignore = "heavy: generates a >10 MiB PNG (run explicitly)"]
fn pipeline_transcodes_oversized_png_to_jpeg() {
// 6000x4000 (sum 10000 — under the dimension cap) smooth gradient with
// small per-pixel noise: PNG-incompressible (delta filters defeated)
// but JPEG-friendly (DCT smooths the small noise). Verified with
// ffmpeg: 8000x6000 amp-5 variant is a 59 MB PNG / 3.3 MB JPEG.
let (w, h) = (6000u32, 4000u32);
let mut rng = 0x1234_5678_9abc_def0u64;
let mut data = Vec::with_capacity((w * h * 3) as usize);
for y in 0..h {
for x in 0..w {
let base = (x + y) * 255 / (w + h);
rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let n = ((rng >> 33) % 11) as i32 - 5; // noise in [-5, 5]
let v = (base as i32 + n).clamp(0, 255) as u8;
data.extend_from_slice(&[v, v, v]);
}
}
let mut bytes = Vec::new();
{
let mut encoder = png::Encoder::new(&mut bytes, w, h);
encoder.set_color(png::ColorType::Rgb);
encoder.set_depth(png::BitDepth::Eight);
let mut writer = encoder.write_header().unwrap();
writer.write_image_data(&data).unwrap();
}
assert!(bytes.len() as u64 > MAX_UPLOAD_BYTES, "test needs a >10MiB PNG, got {}", bytes.len());
let mut file = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
std::io::Write::write_all(file.as_file_mut(), &bytes).unwrap();
match prepare_photo(file).unwrap() {
PhotoPrep::Upload(file) => {
let out = std::fs::read(file.path()).unwrap();
assert!(out.starts_with(&[0xFF, 0xD8]), "must transcode to JPEG");
assert!(out.len() as u64 <= MAX_UPLOAD_BYTES);
}
PhotoPrep::UseFallback => panic!("PNG over the byte cap must transcode to JPEG"),
}
}
}
+513
View File
@@ -0,0 +1,513 @@
//! Generic persistent task queue backed by SQLite (table `tasks`).
//!
//! Concepts kept from the Python `utils/task_queue.py` (untrusted, redesigned):
//! the table schema, the lease/lock/recovery model, and the retry→dead-letter
//! flow. The Python dict-mutation hack (attempts inside the payload) is
//! replaced by dedicated columns.
use parking_lot::Mutex;
use rusqlite::{params, Connection, TransactionBehavior};
use serde_json::Value;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::Notify;
use tokio::task::JoinHandle;
pub const MAX_RETRIES: u32 = 2;
pub const LOCK_TTL_SECONDS: f64 = 120.0;
/// Number of concurrent worker loops. Tasks are independent (retries and
/// forward resumes); leases serialize row claims via SQLite transactions, so
/// extra workers drain backlogs faster. Each worker can be mid-send to
/// Telegram at the same time as handler tasks, so keep this modest.
const QUEUE_WORKERS: usize = 4;
/// What a handler returns instead of throwing. The payload it carries is the
/// (possibly updated) task state to persist for the next attempt.
pub enum QueueError {
/// Reschedule with the given delay; after `MAX_RETRIES` attempts the task
/// is dead-lettered instead.
Retryable {
delay_seconds: f64,
payload: Value,
},
/// Give up now.
Permanent {
message: String,
payload: Value,
},
}
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
type Handler = dyn Fn(Value) -> BoxFuture<'static, Result<(), QueueError>> + Send + Sync;
type DeadLetter = dyn Fn(Value, String) -> BoxFuture<'static, ()> + Send + Sync;
pub struct PersistentTaskQueue {
db_path: String,
notify: Arc<Notify>,
stop: Arc<AtomicBool>,
worker: Mutex<Vec<JoinHandle<()>>>,
counter: AtomicU64,
}
struct LeasedRow {
id: String,
payload: String,
attempts: i32,
}
/// Owned worker state so the spawned loop does not borrow the queue handle.
struct QueueWorker {
db_path: String,
notify: Arc<Notify>,
stop: Arc<AtomicBool>,
handler: Arc<Handler>,
dead_letter: Arc<DeadLetter>,
}
fn now_f64() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
/// Opens the queue DB with a busy timeout. Handler tasks enqueue while
/// workers lease/update rows concurrently; without the timeout a concurrent
/// write fails immediately with SQLITE_BUSY and the operation is lost.
fn open_db(path: &str) -> rusqlite::Result<Connection> {
let conn = Connection::open(path)?;
conn.busy_timeout(Duration::from_secs(5))?;
Ok(conn)
}
fn ensure_schema(conn: &Connection) -> rusqlite::Result<()> {
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \
run_after REAL NOT NULL, attempts INTEGER NOT NULL, status TEXT NOT NULL, \
locked_until REAL NOT NULL, created_at REAL NOT NULL);",
)
}
impl PersistentTaskQueue {
pub fn new(db_path: &str) -> Self {
// Ensure the parent dir and table exist even if only the queue (not
// ChatStore) is used — a fresh container without a mounted data dir
// must still be able to open the DB.
if let Some(parent) = std::path::Path::new(db_path).parent()
&& !parent.as_os_str().is_empty()
&& let Err(e) = std::fs::create_dir_all(parent)
{
log::error!("failed to create queue dir: {e}");
}
if let Ok(conn) = Connection::open(db_path)
&& let Err(e) = ensure_schema(&conn)
{
log::error!("failed to initialize queue schema: {e}");
}
Self {
db_path: db_path.to_string(),
notify: Arc::new(Notify::new()),
stop: Arc::new(AtomicBool::new(false)),
worker: Mutex::new(Vec::new()),
counter: AtomicU64::new(0),
}
}
/// Starts the worker loops. Also recovers rows left `in_progress` by a
/// previous process (lease expired).
pub async fn start<H, F, D, G>(&self, handler: H, dead_letter: D)
where
H: Fn(Value) -> F + Send + Sync + 'static,
F: Future<Output = Result<(), QueueError>> + Send + 'static,
D: Fn(Value, String) -> G + Send + Sync + 'static,
G: Future<Output = ()> + Send + 'static,
{
let handler: Arc<Handler> = Arc::new(move |payload| Box::pin(handler(payload)));
let dead_letter: Arc<DeadLetter> =
Arc::new(move |payload, message| Box::pin(dead_letter(payload, message)));
self.recover_stale().await;
let mut handles = Vec::with_capacity(QUEUE_WORKERS);
for _ in 0..QUEUE_WORKERS {
let worker = QueueWorker {
db_path: self.db_path.clone(),
notify: Arc::clone(&self.notify),
stop: Arc::clone(&self.stop),
handler: Arc::clone(&handler),
dead_letter: Arc::clone(&dead_letter),
};
handles.push(tokio::spawn(worker.run_loop()));
}
*self.worker.lock() = handles;
}
pub async fn stop(&self) {
self.stop.store(true, Ordering::Relaxed);
self.notify.notify_waiters();
let handles = std::mem::take(&mut *self.worker.lock());
for handle in handles {
let _ = handle.await;
}
}
/// Persists a task. `run_after` is an absolute unix timestamp (seconds).
/// Notifies the worker only after the insert has committed, so the worker
/// never wakes to an invisible row.
pub async fn enqueue(&self, payload: Value, run_after: f64) -> rusqlite::Result<()> {
let id = format!(
"task_{}_{}",
(now_f64() * 1000.0) as u64,
self.counter.fetch_add(1, Ordering::Relaxed)
);
let payload = payload.to_string();
let db_path = self.db_path.clone();
log::info!("enqueued {id} (run_after {run_after:.1})");
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute(
"INSERT OR REPLACE INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
VALUES (?1, ?2, ?3, 0, 'pending', 0, ?4)",
params![id, payload, run_after, now_f64()],
)?;
Ok(())
})
.await
.expect("queue insert worker panicked")?;
// Wake every sleeping worker: with several workers the one that finds
// nothing due must not starve the newly inserted row.
self.notify.notify_waiters();
Ok(())
}
async fn recover_stale(&self) {
let db_path = self.db_path.clone();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute(
"UPDATE tasks SET status='pending', locked_until=0 WHERE status='in_progress' AND locked_until < ?1",
params![now_f64()],
)?;
Ok(())
})
.await
.expect("queue recovery worker panicked")
.unwrap_or_else(|e| log::error!("queue recovery failed: {e}"));
}
}
impl QueueWorker {
async fn run_loop(self) {
while !self.stop.load(Ordering::Relaxed) {
match self.lease_next().await {
Some(row) => self.process(row).await,
None => {
let wait_until = self.earliest_run_after().await;
let notified = self.notify.notified();
tokio::pin!(notified);
match wait_until {
Some(until) => {
let delay = (until - now_f64()).max(0.0);
tokio::select! {
_ = &mut notified => {}
_ = tokio::time::sleep(Duration::from_secs_f64(delay)) => {}
}
}
None => {
notified.await;
}
}
}
}
}
}
/// Leases the oldest due row (sets it `in_progress` with a lock TTL).
async fn lease_next(&self) -> Option<LeasedRow> {
let db_path = self.db_path.clone();
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<LeasedRow>> {
let mut conn = open_db(&db_path)?;
// BEGIN IMMEDIATE: with several workers, a deferred transaction
// that read before another worker's lease commit would fail with
// SQLITE_BUSY_SNAPSHOT. Taking the write lock up front serializes
// leases and re-reads the freshest committed state.
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let now = now_f64();
let row = tx.query_row(
"SELECT id, payload, attempts FROM tasks WHERE status='pending' AND run_after <= ?1 AND locked_until <= ?1 \
ORDER BY run_after LIMIT 1",
params![now],
|r| {
Ok((
r.get::<_, String>(0)?,
r.get::<_, String>(1)?,
r.get::<_, i32>(2)?,
))
},
);
let (id, payload, attempts) = match row {
Ok(row) => row,
Err(rusqlite::Error::QueryReturnedNoRows) => {
tx.commit()?;
return Ok(None);
}
Err(e) => return Err(e),
};
tx.execute(
"UPDATE tasks SET status='in_progress', locked_until=?1 WHERE id=?2",
params![now + LOCK_TTL_SECONDS, id],
)?;
tx.commit()?;
Ok(Some(LeasedRow {
id,
payload,
attempts,
}))
})
.await
.expect("queue lease worker panicked")
.unwrap_or_else(|e| {
log::error!("queue lease failed: {e}");
None
})
}
async fn earliest_run_after(&self) -> Option<f64> {
let db_path = self.db_path.clone();
tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<f64>> {
let conn = open_db(&db_path)?;
let mut stmt = conn.prepare("SELECT MIN(run_after) FROM tasks WHERE status='pending'")?;
let mut rows = stmt.query([])?;
match rows.next()? {
Some(row) => Ok(row.get::<_, Option<f64>>(0)?),
None => Ok(None),
}
})
.await
.expect("queue timing worker panicked")
.unwrap_or_else(|e| {
log::error!("queue timing query failed: {e}");
None
})
}
async fn process(&self, row: LeasedRow) {
let payload: Value = match serde_json::from_str(&row.payload) {
Ok(value) => value,
Err(e) => {
log::error!("queue: unparseable payload for {}: {e}", row.id);
self.delete_row(&row.id).await;
(self.dead_letter)(Value::Null, format!("invalid stored payload: {e}")).await;
return;
}
};
log::info!("processing {} (attempt {})", row.id, row.attempts + 1);
match (self.handler)(payload).await {
Ok(()) => {
log::info!("task {} completed", row.id);
self.delete_row(&row.id).await;
}
Err(QueueError::Retryable {
delay_seconds,
payload,
}) => {
if row.attempts as u32 >= MAX_RETRIES {
let message = format!("task failed after {MAX_RETRIES} retries");
log::error!("dead-lettering {}: {message}", row.id);
self.delete_row(&row.id).await;
(self.dead_letter)(payload, message).await;
} else {
log::info!(
"task {} rescheduled in {delay_seconds:.1}s (attempt {})",
row.id,
row.attempts + 1
);
self.reschedule(&row.id, payload, delay_seconds, row.attempts + 1)
.await;
}
}
Err(QueueError::Permanent { message, payload }) => {
log::error!("dead-lettering {}: {message}", row.id);
self.delete_row(&row.id).await;
(self.dead_letter)(payload, message).await;
}
}
}
async fn delete_row(&self, id: &str) {
let db_path = self.db_path.clone();
let id = id.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute("DELETE FROM tasks WHERE id = ?1", params![id])?;
Ok(())
})
.await
.expect("queue delete worker panicked")
.unwrap_or_else(|e| log::error!("queue delete failed: {e}"));
}
async fn reschedule(&self, id: &str, payload: Value, delay_seconds: f64, attempts: i32) {
let db_path = self.db_path.clone();
let id = id.to_string();
let payload = payload.to_string();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = open_db(&db_path)?;
conn.execute(
"UPDATE tasks SET payload=?1, run_after=?2, attempts=?3, status='pending', locked_until=0 WHERE id=?4",
params![payload, now_f64() + delay_seconds, attempts, id],
)?;
Ok(())
})
.await
.expect("queue reschedule worker panicked")
.unwrap_or_else(|e| log::error!("queue reschedule failed: {e}"));
self.notify.notify_waiters();
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
async fn new_queue() -> (PersistentTaskQueue, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("queue.db");
let queue = PersistentTaskQueue::new(path.to_str().unwrap());
(queue, dir)
}
#[tokio::test]
async fn enqueue_runs_handler_once() {
let (queue, _dir) = new_queue().await;
let calls = Arc::new(AtomicUsize::new(0));
let calls_worker = calls.clone();
queue
.start(
move |payload| {
assert_eq!(payload["n"], 42);
calls_worker.fetch_add(1, AtomicOrdering::SeqCst);
async { Ok(()) }
},
|_payload, _message| async {},
)
.await;
queue
.enqueue(serde_json::json!({"n": 42}), now_f64())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
queue.stop().await;
}
#[tokio::test]
async fn retryable_reschedules_then_dead_letters() {
let (queue, _dir) = new_queue().await;
let calls = Arc::new(AtomicUsize::new(0));
let dead_calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
let d = dead_calls.clone();
queue
.start(
move |payload| {
c.fetch_add(1, AtomicOrdering::SeqCst);
let payload = payload.clone();
async move {
Err(QueueError::Retryable {
delay_seconds: 0.001,
payload,
})
}
},
move |_payload, _message| {
d.fetch_add(1, AtomicOrdering::SeqCst);
async {}
},
)
.await;
queue
.enqueue(serde_json::json!({"a": 1}), now_f64())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(600)).await;
assert_eq!(
calls.load(AtomicOrdering::SeqCst),
MAX_RETRIES as usize + 1,
"handler should run once per attempt"
);
assert_eq!(dead_calls.load(AtomicOrdering::SeqCst), 1);
queue.stop().await;
}
#[tokio::test]
async fn permanent_error_dead_letters_immediately() {
let (queue, _dir) = new_queue().await;
let calls = Arc::new(AtomicUsize::new(0));
let dead_calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
let d = dead_calls.clone();
queue
.start(
move |payload| {
c.fetch_add(1, AtomicOrdering::SeqCst);
let payload = payload.clone();
async move {
Err(QueueError::Permanent {
message: "nope".into(),
payload,
})
}
},
move |_payload, message| {
assert_eq!(message, "nope");
d.fetch_add(1, AtomicOrdering::SeqCst);
async {}
},
)
.await;
queue
.enqueue(serde_json::json!({"a": 1}), now_f64())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
assert_eq!(dead_calls.load(AtomicOrdering::SeqCst), 1);
queue.stop().await;
}
#[tokio::test]
async fn stale_in_progress_row_is_recovered_on_start() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("queue.db");
// Insert a stale leased row directly (lease expired).
{
let conn = Connection::open(&path).unwrap();
ensure_schema(&conn).unwrap();
conn.execute(
"INSERT INTO tasks (id, payload, run_after, attempts, status, locked_until, created_at) \
VALUES ('task_stale', '{\"s\":1}', 0, 0, 'in_progress', ?1, 0)",
params![now_f64() - 10.0],
)
.unwrap();
}
let queue = PersistentTaskQueue::new(path.to_str().unwrap());
let calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
queue
.start(
move |payload| {
assert_eq!(payload["s"], 1);
c.fetch_add(1, AtomicOrdering::SeqCst);
async { Ok(()) }
},
|_payload, _message| async {},
)
.await;
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(calls.load(AtomicOrdering::SeqCst), 1);
queue.stop().await;
}
}
File diff suppressed because it is too large Load Diff
+156
View File
@@ -0,0 +1,156 @@
//! Per-chat state with SQLite persistence (table `chat_state` in
//! `data/task_queue.db`, shared with the task queue).
use parking_lot::Mutex;
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[derive(Serialize, Deserialize, Default, Clone, Debug)]
pub struct ChatData {
pub forward_channel_id: Option<i64>,
pub edit_before_forward: bool,
/// Key: prompt message id.
pub edit_message: HashMap<i64, EditMessage>,
/// name -> HTML template containing "[]"
pub template: HashMap<String, String>,
/// site name (twitter/bsky/pixiv) -> user-supplied caption format with
/// {url} {author} {author_url} {title} {tags} placeholders.
pub message_format: HashMap<String, String>,
}
#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct EditMessage {
pub url: String,
pub chat_id: i64,
pub forward_message_ids: Vec<i64>,
pub template: String,
/// Unix seconds at registration; expiry = created_at + ttl.
pub created_at: i64,
}
pub struct ChatStore {
/// In-memory cache; the DB is the source of truth on first access.
cache: Mutex<HashMap<i64, ChatData>>,
db_path: String,
}
pub fn unix_now() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
impl ChatStore {
/// Creates the parent directory and both tables (idempotent).
pub fn open(path: &str) -> rusqlite::Result<Self> {
if let Some(parent) = Path::new(path).parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
}
let conn = Connection::open(path)?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, payload TEXT NOT NULL, \
run_after REAL NOT NULL, attempts INTEGER NOT NULL, status TEXT NOT NULL, \
locked_until REAL NOT NULL, created_at REAL NOT NULL); \
CREATE TABLE IF NOT EXISTS chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL);",
)?;
drop(conn);
Ok(ChatStore {
cache: Mutex::new(HashMap::new()),
db_path: path.to_string(),
})
}
pub async fn get(&self, chat_id: i64) -> ChatData {
if let Some(data) = self.cache.lock().get(&chat_id) {
return data.clone();
}
let db_path = self.db_path.clone();
let payload = tokio::task::spawn_blocking(move || -> rusqlite::Result<Option<String>> {
let conn = Connection::open(&db_path)?;
// Concurrent handler tasks (batch-forwards) may write chat_state
// while this read runs; without a busy timeout a write lock
// collision fails the query immediately.
conn.busy_timeout(std::time::Duration::from_secs(5))?;
let mut stmt = conn.prepare("SELECT payload FROM chat_state WHERE chat_id = ?1")?;
let mut rows = stmt.query(params![chat_id.to_string()])?;
match rows.next()? {
Some(row) => Ok(Some(row.get(0)?)),
None => Ok(None),
}
})
.await
.expect("chat_state worker panicked")
.unwrap_or_else(|e| {
log::error!("chat_state read failed: {e}");
None
})
.unwrap_or_default();
let data: ChatData = serde_json::from_str(&payload).unwrap_or_default();
self.cache.lock().insert(chat_id, data.clone());
data
}
/// Write-through: update the cache and the DB.
pub async fn set(&self, chat_id: i64, data: &ChatData) {
self.cache.lock().insert(chat_id, data.clone());
let payload = serde_json::to_string(data).expect("chat state serializes");
let db_path = self.db_path.clone();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = Connection::open(&db_path)?;
conn.busy_timeout(std::time::Duration::from_secs(5))?;
conn.execute(
"INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)",
params![chat_id.to_string(), payload],
)?;
Ok(())
})
.await
.expect("chat_state worker panicked")
.unwrap_or_else(|e| log::error!("chat_state write failed: {e}"));
}
/// Removes edit-before-forward records whose `created_at + ttl` is in the
/// past. Returns the removed `(chat_id, prompt_message_id)` pairs so the
/// caller can clear the prompt's buttons.
pub async fn prune_expired(&self, ttl: Duration) -> Vec<(i64, i64)> {
let now = unix_now();
let ttl_secs = ttl.as_secs() as i64;
let mut removed = Vec::new();
let changed: Vec<(i64, ChatData)> = {
let mut cache = self.cache.lock();
let mut out = Vec::new();
for (chat_id, data) in cache.iter_mut() {
let keys: Vec<i64> = data.edit_message.keys().copied().collect();
let mut kept = HashMap::new();
for key in keys {
if let Some(entry) = data.edit_message.get(&key) {
if entry.created_at + ttl_secs > now {
kept.insert(key, entry.clone());
} else {
removed.push((*chat_id, key));
}
}
}
if kept.len() != data.edit_message.len() {
data.edit_message = kept;
out.push((*chat_id, data.clone()));
}
}
out
};
for (chat_id, data) in changed {
self.set(chat_id, &data).await;
}
if !removed.is_empty() {
log::info!("pruned {} expired edit-before-forward record(s)", removed.len());
}
removed
}
}
+59 -14
View File
@@ -1,21 +1,66 @@
services:
tgxmb: image: yoursfunny/telegram-twitter-media-bot: latest
nginx-proxy:
image: nginxproxy/nginx-proxy:1.11.6-alpine
restart: always
ports:
- "8443:8443"
- '80:80'
- '443:443'
volumes:
- /var/run/docker.sock:/tmp/docker.sock:ro
- certs:/etc/nginx/certs:ro
- html:/usr/share/nginx/html:ro
networks: [proxy]
labels:
- 'com.github.nginx-proxy.nginx=true'
container_name: nginx-proxy
acme-companion:
image: nginxproxy/acme-companion
restart: always
environment:
DEFAULT_EMAIL: 'admin@yoursfunny.top'
volumes:
- /var/run/docker.sock:/var/run/docker.sock:ro
- certs:/etc/nginx/certs:rw
- html:/usr/share/nginx/html:rw
- acme:/etc/acme.sh
networks: [proxy]
container_name: acme-companion
depends_on:
- nginx-proxy
tgxmb:
image: yoursfunny/telegram-twitter-media-bot:latest
restart: always
environment:
LOCAL_USER_ID: '1000'
BOT_TOKEN: ''
TELOXIDE_TOKEN: ''
BOT_ADMIN: ''
WEBHOOK: false
WEBHOOK_LISTEN: '127.0.0.1'
WEBHOOK_PORT: 8443
WEBHOOK_URL: 'https://example.com'
WEBHOOK_KEY: './cert/private.key'
WEBHOOK_CERT: './cert/cert.pem'
WEBHOOK_SECRET_TOKEN: 'secret-token'
# LOG_LEVEL: 'WARNING'
volumes:
- ./data: /app/data
# - ./cert:/app/cert
PIXIV_REFRESH_TOKEN: ''
TWITTER_AUTH_TOKEN: ''
EDIT_MESSAGE_TTL_SECONDS: '86400'
LINK_CACHE_TTL_SECONDS: '604800'
RUST_LOG: 'info'
VIRTUAL_HOST: '<YOUR_DOMAIN>'
VIRTUAL_PORT: '8443'
# ACME_HOST: 'your.domain.com'
WEBHOOK: 'true'
WEBHOOK_LISTEN: '0.0.0.0'
WEBHOOK_PORT: '8443'
WEBHOOK_URL: 'https://<YOUR_DOMAIN>/'
WEBHOOK_SECRET_TOKEN: ''
volumes:
- ./data:/app/data
networks: [proxy]
depends_on:
- nginx-proxy
container_name: tgxmb
volumes:
certs:
html:
acme:
networks:
proxy:
name: proxy
+18 -4
View File
@@ -5,12 +5,26 @@ if [ "$(id -u)" -eq '0' ]
then
USER_ID=${LOCAL_USER_ID:-9001}
useradd --shell /bin/bash -u ${USER_ID} -o -c "" -m user > /dev/null 2>&1
usermod -a -G root user > /dev/null 2>&1
chown -R `id -u user`:`id -u user` /app > /dev/null 2>&1
# `docker compose restart` / `docker restart` reuse the same container, so
# the overlay fs keeps the user created on first boot. A second `useradd`
# then fails with exit code 9, which would trip `set -e` and kill the
# container on every restart. Create only if missing; align the UID
# otherwise so LOCAL_USER_ID changes still apply.
if ! id user > /dev/null 2>&1
then
useradd --shell /bin/bash -u ${USER_ID} -o -c "" -m user > /dev/null 2>&1 || true
else
usermod -u ${USER_ID} -o user > /dev/null 2>&1 || true
fi
usermod -a -G root user > /dev/null 2>&1 || true
# Bind-mounted volumes may not support chown; a failure here must not kill
# the container either.
chown -R `id -u user`:`id -u user` /app > /dev/null 2>&1 || true
export HOME=/home/user
exec gosu user "$0" "$@"
# setpriv (util-linux, present in bookworm-slim) replaces gosu: drop to the
# target user and exec, keeping the process as PID 1.
exec setpriv --reuid=`id -u user` --regid=`id -g user` --init-groups "$@"
fi
exec "$@"
-261
View File
@@ -1,261 +0,0 @@
import html
from functools import wraps
from typing import TYPE_CHECKING
from telegram import InlineKeyboardButton, InlineKeyboardMarkup
from telegram.constants import ChatAction, ChatType, ParseMode
from telegram.ext import (ApplicationBuilder, CallbackQueryHandler, CommandHandler, Defaults,
InlineQueryHandler, MessageHandler, PicklePersistence, filters)
import common
from tweet import TGTweet
if TYPE_CHECKING:
from telegram import Chat, Message, Update
from telegram.ext import Application, ContextTypes
logger = common.get_logger(__name__)
def send_action(action):
def decorator(func):
@wraps(func)
async def command_func(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs):
await update.effective_chat.send_action(action)
return await func(update, context, *args, **kwargs)
return command_func
return decorator
async def inline_query(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
query = update.inline_query.query
if query == "":
return
logger.info(f"Query: {query}")
async with TGTweet(query) as tweet:
result = list(tweet.inline_query_generator)
await update.inline_query.answer(result)
@send_action(ChatAction.UPLOAD_PHOTO)
async def url_media(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
url = update.message.text
logger.info(f"Receiving url: {url}")
async with TGTweet(url) as tweet:
media = list(tweet.pm_media_generator)
message_to_send = await update.effective_message.reply_media_group(
media,
caption=tweet.message_text,
reply_to_message_id=update.message.message_id,
) if not tweet.is_single_gif else await update.effective_message.reply_animation(
media[0][0],
caption=tweet.message_text,
reply_to_message_id=update.message.message_id,
has_spoiler=media[0][1]
)
if not isinstance(message_to_send, tuple):
message_to_send = (message_to_send,)
url = tweet.url
if context.user_data.get('edit_before_forward', False):
message_reply = await update.effective_message.reply_text(
"Reply to edit message. [URL]",
reply_markup=InlineKeyboardMarkup.from_button(
InlineKeyboardButton("↩️ Confirm", callback_data="forward")
),
reply_to_message_id=update.message.message_id,
)
context.user_data['message_reply'] = message_reply
context.user_data['message_to_send'] = message_to_send
context.user_data['message_url'] = url
return
if 'forward_channel_id' in context.user_data:
await forward_message(update, context, message_to_send)
async def forward_message(
update: Update,
context: ContextTypes.DEFAULT_TYPE,
message_to_send: tuple[Message, ...],
) -> None:
try:
await update.effective_chat.copy_messages(
context.user_data['forward_channel_id'],
[m.id for m in message_to_send]
)
except Exception as e:
await update.effective_message.reply_text(str(e))
async def edit_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
if 'message_reply' not in context.user_data:
return
if update.message.reply_to_message != context.user_data['message_reply']:
return
template = context.user_data.get('template', None)
message_url = '<a href="{0}">{1}</a>'
url = context.user_data['message_url']
if template:
update_text = template.replace("[]", message_url.format(
url,
html.escape(update.message.text)
))
else:
update_text = html.escape(update.message.text)
match = common.message_url_regex.search(update_text)
if match:
match = match.span()
update_text = update_text[:match[0]] + message_url.format(
url,
update_text[match[0] + 1:match[1] - 1]
) + update_text[match[1]:]
message_to_send = context.user_data['message_to_send']
await message_to_send[0].edit_caption(update_text)
async def query_forward_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
message_to_send = context.user_data['message_to_send']
await forward_message(update, context, message_to_send)
await update.callback_query.answer('✅ Forwarded')
await update.callback_query.delete_message()
del context.user_data['message_reply']
del context.user_data['message_to_send']
del context.user_data['message_url']
@send_action(ChatAction.TYPING)
async def cmd_set_forward_channel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
if not context.args:
await update.effective_message.reply_text("Please provide a channel username or id.")
return
channel = context.args[0]
try:
channel: Chat = await context.bot.get_chat(channel)
except Exception as e:
await update.effective_message.reply_text(str(e))
return
if channel.type != ChatType.CHANNEL:
await update.effective_message.reply_text("That is not a channel.")
return
try:
channel_admin = await channel.get_administrators()
except Exception as e:
await update.effective_message.reply_text(str(e) + "\nPlease add the bot to the channel and set as admin")
return
user_bot = filter(lambda x: x.user.id == context.bot.id, channel_admin)
user_bot = next(user_bot, None)
if user_bot.can_post_messages:
context.user_data['forward_channel_id'] = channel.id
await update.effective_message.reply_text("Add successfully.")
@send_action(ChatAction.TYPING)
async def cmd_remove_forward_channel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
if 'forward_channel_id' in context.user_data:
del context.user_data['forward_channel_id']
await update.effective_message.reply_text("Remove successfully.")
return
await update.effective_message.reply_text("No channel to remove.")
@send_action(ChatAction.TYPING)
async def cmd_edit_before_forward(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
if context.user_data.get('forward_channel_id', None) is None:
await update.effective_message.reply_text("Please enable forward channel first.")
return
ebf_status = context.user_data.get('edit_before_forward', False)
if ebf_status:
context.user_data['edit_before_forward'] = False
context.user_data.pop('message_reply', None)
context.user_data.pop('message_to_send', None)
context.user_data.pop('message_url', None)
await update.effective_message.reply_text("Disable edit before forward.")
return
context.user_data['edit_before_forward'] = True
await update.effective_message.reply_text("Enable edit before forward.")
@send_action(ChatAction.TYPING)
async def cmd_set_template(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
reply = update.effective_message.reply_to_message
if not reply:
await update.effective_message.reply_text("Please reply to a message to set as template.")
return
if '[]' not in reply.text_html:
await update.effective_message.reply_text("Please reply to a message with [] to set as template.")
return
context.user_data['template'] = reply.text_html
await update.effective_message.reply_text("Template set.")
@send_action(ChatAction.TYPING)
async def cmd_user_dict(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await update.effective_message.reply_text(str(context.user_data))
async def post_init(application: Application) -> None:
# commands = [
# BotCommand('start', CMD_START),
# ]
# await application.bot.set_my_commands(commands)
DESCRIPTION = "A bot to fetch tweets from Twitter."
await application.bot.set_my_description(DESCRIPTION)
await application.bot.set_my_short_description(DESCRIPTION)
TGTweet.init_client()
async def post_stop(application: Application) -> None:
await application.bot.send_message(common.ADMIN[0], "Shutting down...")
async def post_shutdown(application: Application) -> None:
await TGTweet.close_client()
def main():
defaults = Defaults(parse_mode=ParseMode.HTML, allow_sending_without_reply=True)
persistence = PicklePersistence(filepath='data/pers.pkl')
application = (ApplicationBuilder()
.token(common.BOT_TOKEN)
.defaults(defaults)
.persistence(persistence)
.post_init(post_init)
.post_stop(post_stop)
.post_shutdown(post_shutdown)
.concurrent_updates(True)
.build()
)
user_filter = filters.User()
user_filter.add_user_ids(common.ADMIN)
handlers = [
InlineQueryHandler(inline_query, common.x_url_regex),
MessageHandler(filters.Regex(common.x_url_regex) & filters.ChatType.PRIVATE, url_media),
CommandHandler("set_forward_channel", cmd_set_forward_channel),
CommandHandler("remove_forward_channel", cmd_remove_forward_channel),
CommandHandler("edit_before_forward", cmd_edit_before_forward),
CommandHandler("set_template", cmd_set_template),
MessageHandler(~filters.COMMAND & filters.ChatType.PRIVATE, edit_message),
CallbackQueryHandler(query_forward_message, pattern="forward"),
CommandHandler("bot_dict", cmd_user_dict, filters=user_filter),
]
application.add_handlers(handlers)
if common.WEBHOOK:
application.run_webhook(
listen=common.WEBHOOK_LISTEN,
port=common.WEBHOOK_PORT,
secret_token=common.WEBHOOK_SECRET_TOKEN,
key=common.WEBHOOK_KEY,
cert=common.WEBHOOK_CERT,
webhook_url=common.WEBHOOK_URL
)
else:
application.run_polling()
if __name__ == '__main__':
main()
-3
View File
@@ -1,3 +0,0 @@
python-telegram-bot[webhooks]~=21.3
httpx[http2]~=0.27.0
uvloop~=0.19.0; sys_platform != 'win32'
-254
View File
@@ -1,254 +0,0 @@
import html
from typing import Generator
from httpx import AsyncClient
from telegram import (InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, InputMediaPhoto,
InputMediaVideo)
from common import get_logger, x_media_regex, x_tco_regex, x_url_regex
logger = get_logger(__name__)
twimg_url = 'https://pbs.twimg.com/'
vx_api_url = 'https://api.vxtwitter.com/{0}/status/{1}'
message_raw_text = """{url}
<a href="{author_url}">{author}</a>: {text}
"""
def create_client() -> 'AsyncClient':
return AsyncClient(http2=True)
async def close_client(_client: 'AsyncClient') -> None:
await _client.aclose()
async def fetch_json(_client: 'AsyncClient', url: str) -> dict:
logger.info(f"Fetching {url}")
response = await _client.get(url)
assert response.status_code == response.is_success, f"Failed to fetch {url}, status code {response.status_code}"
return response.json()
class TweetMedia:
def __init__(self, url: str, thumb: str, media_type: str):
self._url: str = url
self._thumb: str = thumb
self._type: str = media_type
def __str__(self):
return f"Media[url: {self.url} thumb: {self.thumb} type: {self.type}]"
@property
def _uri(self) -> str | None:
match = x_media_regex.match(self._url)
if match:
return match.group(2).removesuffix('.jpg').removesuffix('.png')
return None
@property
def url(self) -> str:
match self._type:
case "image":
return f"{twimg_url}{self._uri}?format=jpg&name=4096x4096"
case "video":
return self._url
case "gif":
return self._url
case _:
return self._url
@property
def thumb(self) -> str:
match self._type:
case "image":
return f"{twimg_url}{self._uri}?format=jpg&name=thumb"
case "video":
return self._thumb
case "gif":
return self._thumb
case _:
return self._thumb
@property
def type(self) -> str:
return self._type
class Tweet:
def __init__(
self,
tweet_id: str,
author: str,
author_id: str,
text: str,
media: list[TweetMedia],
sensitive: bool = False
):
self._id: str = tweet_id
self._author: str = author
self._author_id: str = author_id
self._text: str = text
self._media: list[TweetMedia] = media
self._sensitive: bool = sensitive
@property
def id(self) -> str:
return self._id
@property
def url(self) -> str:
return f"https://twitter.com/{self._author_id}/status/{self._id}"
@property
def author(self) -> str:
return self._author
@property
def author_url(self) -> str:
return f"https://twitter.com/{self._author_id}"
@property
def text(self) -> str:
return self._text
@property
def media(self) -> list[TweetMedia]:
return self._media
@property
def sensitive(self) -> bool:
return self._sensitive
class TGTweet(Tweet):
_httpx_client: AsyncClient
def __init__(self, url: str):
self._url: str = url
self._api_param: tuple[str] = self._tweet_id
self._is_single_gif: bool = False
assert self._api_param
async def __aenter__(self):
self._tweet: dict = await self._fetch_tweet(self._api_param)
super().__init__(*self._init_properties)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
@classmethod
def init_client(cls) -> None:
cls._httpx_client = create_client()
@classmethod
async def close_client(cls) -> None:
await close_client(cls._httpx_client)
async def _fetch_tweet(self, api_param: tuple[str]) -> dict:
return await fetch_json(self._httpx_client, vx_api_url.format(*api_param))
@property
def _tweet_id(self) -> tuple[str] | None:
match = x_url_regex.match(self._url)
if match:
return match.groups()
return None
@property
def _tweet_text(self) -> str:
match = x_tco_regex.search(self._tweet['text'])
return self._tweet['text'][:match.start()].strip(" ") if match else self._tweet['text']
@property
def _tweet_media(self) -> list[TweetMedia]:
return [
TweetMedia(
url=x['url'],
thumb=x['thumbnail_url'],
media_type=x['type']
)
for x in self._tweet['media_extended']
]
@property
def _init_properties(self) -> tuple:
id = self._tweet['tweetID']
author = self._tweet['user_name']
author_id = self._tweet['user_screen_name']
text = self._tweet_text
media = self._tweet_media
sensitive = self._tweet['possibly_sensitive']
return id, author, author_id, text, media, sensitive
@property
def is_single_gif(self) -> bool:
return self._is_single_gif
@property
def message_text(self) -> str:
return message_raw_text.format(
url=self.url,
author_url=self.author_url,
author=html.escape(self.author),
text=html.escape(self.text)
)
@property
def inline_query_generator(self) -> Generator[
InlineQueryResultPhoto | InlineQueryResultVideo | InlineQueryResultMpeg4Gif, None, None
]:
for i, tweet_media in enumerate(self.media):
logger.info(str(tweet_media))
if tweet_media.type == "image":
yield InlineQueryResultPhoto(
id=str(i),
photo_url=tweet_media.url,
thumbnail_url=tweet_media.thumb,
caption=self.message_text
)
elif tweet_media.type == "video":
yield InlineQueryResultVideo(
id=str(i),
video_url=tweet_media.url,
mime_type="video/mp4",
thumbnail_url=tweet_media.thumb,
title=self.text,
caption=self.message_text
)
elif tweet_media.type == "gif":
yield InlineQueryResultMpeg4Gif(
id=str(i),
mpeg4_url=tweet_media.url,
thumbnail_url=tweet_media.thumb,
caption=self.message_text
)
@property
def pm_media_generator(self) -> Generator[InputMediaPhoto | InputMediaVideo | tuple[str, bool], None, None]:
for tweet_media in self.media:
logger.info(str(tweet_media))
if tweet_media.type == "image":
yield InputMediaPhoto(
media=tweet_media.url,
has_spoiler=self.sensitive
)
elif tweet_media.type == "video":
yield InputMediaVideo(
media=tweet_media.url,
has_spoiler=self.sensitive,
thumbnail=tweet_media.thumb
)
elif tweet_media.type == "gif":
if len(self.media) == 1:
self._is_single_gif = True
yield tweet_media.url, self.sensitive
yield InputMediaVideo(
media=tweet_media.url,
has_spoiler=self.sensitive,
thumbnail=tweet_media.thumb
)