merge refactor-rs into master

This commit is contained in:
2026-08-03 23:34:45 +08:00
43 changed files with 8101 additions and 1314 deletions
+2 -1
View File
@@ -24,8 +24,9 @@
**/secrets.dev.yaml
**/values.dev.yaml
*.db
.python-version
LICENSE
README.md
data/
cert/
**/target/
.idea/
-1
View File
@@ -3,7 +3,6 @@ __pycache__/
cert/
data/
docker-compose.yml
utils/x.py
.env
Generated
+3536
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"
+54 -19
View File
@@ -1,26 +1,61 @@
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/xz for the ffmpeg download.
FROM rust:1-bookworm AS builder
ARG APP_NAME=telegram-twitter-media-bot
# Statically compiled ffmpeg (ugoira MP4 encoding). amd64 by default; override
# for other platforms or pin a different johnvansickle build.
ARG FFMPEG_URL=https://johnvansickle.com/ffmpeg/releases/ffmpeg-7.0.2-amd64-static.tar.xz
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 johnvansickle tarball has a
# `{build}/ffmpeg` layout, so strip one path component.
RUN wget -q -O /tmp/ffmpeg.tar.xz "$FFMPEG_URL" \
&& tar -xJf /tmp/ffmpeg.tar.xz -C /usr/local/bin --strip-components=1 --wildcards '*/ffmpeg' \
&& rm /tmp/ffmpeg.tar.xz \
&& /usr/local/bin/ffmpeg -version >/dev/null
# 3. Real sources last: only our crates recompile on source changes.
COPY crates/ ./crates/
RUN 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 git 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"]
+50
View File
@@ -0,0 +1,50 @@
# TelegramXMediaBot
Telegram 机器人,将 X / Twitter、Pixiv、Bluesky 的帖子链接转换为媒体消息发送,附带帖子标题、作者与标签。
## 功能
- 私聊发送链接后自动抓取并发送图片、视频与 GIF,超量图片自动分批
- 纯文字帖提示无媒体;不支持的链接静默忽略
- 支持内联查询(`@机器人 <链接>`
- 可绑定转发频道自动转发;支持转发前编辑 caption 与自定义模板
- 发送失败自动重试并持久化,重试耗尽后通知用户
- Pixiv ugoira 动图自动转码为 MP4
## 快速开始
```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``RUST_LOG``WEBHOOK*`
## 命令
| 命令 | 说明 |
|---|---|
| `/set_forward_channel <频道>` | 设置转发频道 |
| `/remove_forward_channel` | 取消转发频道 |
| `/edit_before_forward` | 开关转发前编辑 |
| `/set_template <名称>` | 将回复的消息(含 `[]`)保存为模板 |
| `/set_format <站点> <格式>` | 自定义 caption 格式(占位符 `{url}` `{title}` `{tags}` 等) |
| `/bot_dict` | 查看聊天状态 |
链接处理仅限私聊;命令在任意聊天可用。
## 备注
- 数据持久化于 `data/task_queue.db`,容器部署需挂载该目录
- 运行环境需安装 ffmpeg(Docker 镜像已内置)
- 测试:`cargo test --workspace`
-23
View File
@@ -1,23 +0,0 @@
import os
try:
import uvloop
import 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(",") if i]
PIXIV_REFRESH_TOKEN = os.getenv("PIXIV_REFRESH_TOKEN")
WEBHOOK = os.getenv("WEBHOOK").strip().lower() in ("true", "yes", "1")
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")
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "x-media"
version = "1.0.0"
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"
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;
+37
View File
@@ -0,0 +1,37 @@
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),
}
}
}
#[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,
}
+242
View File
@@ -0,0 +1,242 @@
//! 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) => {
let escaped = html_escape::encode_text(format).into_owned();
escaped
.replace("{url}", &data.url)
.replace("{author}", &data.author)
.replace("{author_url}", &data.author_url)
.replace("{title}", &data.title)
.replace("{tags}", &data.tags)
}
_ => self.caption.clone(),
}
}
}
#[derive(Debug)]
pub enum FetchError {
Http(reqwest::Error),
Json(serde_json::Error),
Pixiv(PixivError),
NotFound,
Blocked,
}
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"),
}
}
}
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 => 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.
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::*;
#[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,
}
@@ -0,0 +1,463 @@
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)?;
Ok(fetch(id).await?.into())
}
/// 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);
}
Ok(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: item.media_url_https,
thumbnail_url: None,
fallback_url: None,
}),
"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` (UTF-16 indices) 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.
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 => {
let units: Vec<u16> = text
.encode_utf16()
.skip(start)
.take(end - start)
.collect();
// Drop the replacement char that a surrogate cut at the boundary
// would produce (the range end is a valid UTF-16 boundary in
// practice, so this is just a safety net).
String::from_utf16_lossy(&units).replace('\u{FFFD}', "")
}
_ => 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
}
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[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 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:?}");
}
}
+4
View File
@@ -0,0 +1,4 @@
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 UTF-16 code 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,
}
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "xmedia-bot"
version = "1.0.0"
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"
x-media = { path = "../x-media" }
+58
View File
@@ -0,0 +1,58 @@
//! 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,
// 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 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());
let webhook_cert = env::var("WEBHOOK_CERT").ok();
let webhook_secret_token = env::var("WEBHOOK_SECRET_TOKEN").ok();
Config {
admin_ids,
edit_message_ttl,
webhook_enabled,
webhook_url,
webhook_listen,
webhook_port,
webhook_cert,
webhook_secret_token,
}
}
}
+639
View File
@@ -0,0 +1,639 @@
use crate::config::Config;
use crate::queue::PersistentTaskQueue;
use crate::send::{self, MediaItemPayload, Task};
use crate::state::{ChatStore, unix_now};
use std::collections::HashSet;
use std::sync::LazyLock;
use teloxide::prelude::*;
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 CONFIG: LazyLock<Config> = LazyLock::new(Config::load);
#[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 {
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,
},
Media::Video { .. } => MediaItemPayload::Video {
media: media.url().to_string(),
has_spoiler: sensitive,
thumbnail: thumbnail_for(media),
},
Media::Animated { .. } => MediaItemPayload::Video {
media: media.url().to_string(),
has_spoiler: sensitive,
thumbnail: thumbnail_for(media),
},
}
}
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}");
}
}
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}");
}
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);
let task = if fetched.media.len() == 1
&& matches!(fetched.media[0], Media::Animated { .. })
{
Task::SendAnimation {
chat_id,
reply_to_message_id: message.id.0 as i64,
caption: caption.clone(),
animation: MediaItemPayload::Animation {
media: fetched.media[0].url().to_string(),
has_spoiler: fetched.sensitive,
},
source_url: fetched.source_url.clone(),
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),
}
} else {
let items: Vec<MediaItemPayload> = fetched
.media
.iter()
.map(|media| media_to_payload(media, fetched.sensitive))
.collect();
Task::SendMediaSequence {
chat_id,
reply_to_message_id: message.id.0 as i64,
caption: caption.clone(),
media_batches: send::chunk_media_items(items),
batch_index: 0,
sent_message_ids: vec![],
source_url: fetched.source_url.clone(),
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),
}
};
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,
..
}) => {
log::error!("send for {url} failed permanently: {err_message}");
let _ = reply(bot, message.clone(), format!("Send failed: {err_message}")).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 {
url_media(bot.clone(), &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 { .. } => InlineQueryResult::Photo(
InlineQueryResultPhoto::new(id, 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(())
}
+131
View File
@@ -0,0 +1,131 @@
use dotenv::dotenv;
use teloxide::dptree::endpoint;
use teloxide::types::{ChatId, InputFile, MessageId};
use teloxide::update_listeners::webhooks;
use teloxide::prelude::*;
use tokio::sync::watch;
use x_media::site;
mod config;
mod handlers;
mod queue;
mod send;
mod state;
use handlers::{CHAT_STORE, CONFIG, TASK_QUEUE};
#[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;
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");
bot.set_webhook(url.clone()).await.unwrap();
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());
}
dispatcher
.dispatch_with_listener(
webhooks::axum(bot.clone(), options)
.await
.expect("Failed to create webhook listener"),
LoggingErrorHandler::with_custom_text("Error from update listener"),
)
.await;
} else {
log::info!("running in polling mode");
dispatcher.dispatch().await;
}
// Graceful stop (Ctrl+C): 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");
}
+488
View File
@@ -0,0 +1,488 @@
//! 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};
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;
/// 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<Option<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)
}
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(None),
counter: AtomicU64::new(0),
}
}
/// Starts the worker loop. 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 worker = QueueWorker {
db_path: self.db_path.clone(),
notify: Arc::clone(&self.notify),
stop: Arc::clone(&self.stop),
handler,
dead_letter,
};
let worker = tokio::spawn(worker.run_loop());
*self.worker.lock() = Some(worker);
}
pub async fn stop(&self) {
self.stop.store(true, Ordering::Relaxed);
self.notify.notify_one();
if let Some(handle) = self.worker.lock().take() {
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})");
let result = tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = Connection::open(&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")?;
self.notify.notify_one();
Ok(result)
}
async fn recover_stale(&self) {
let db_path = self.db_path.clone();
tokio::task::spawn_blocking(move || -> rusqlite::Result<()> {
let conn = Connection::open(&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 = Connection::open(&db_path)?;
let tx = conn.transaction()?;
let now = now_f64();
let row = tx.query_row(
"SELECT id, payload, attempts FROM tasks WHERE status='pending' AND run_after <= ?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 = Connection::open(&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 = Connection::open(&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 = Connection::open(&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_one();
}
}
#[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;
}
}
+917
View File
@@ -0,0 +1,917 @@
//! Typed task payloads and send/forward executors with retry classification
//! and the download-and-reupload fallback (Telegram's own fetch of a media
//! URL is blocked by hotlink protection; the bot downloads the file itself
//! and uploads it via multipart).
use crate::handlers::{CHAT_STORE, TASK_QUEUE};
use crate::queue::QueueError;
use crate::state::{EditMessage, unix_now};
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tempfile::NamedTempFile;
use teloxide::prelude::*;
use teloxide::types::{
ChatId, InlineKeyboardButton, InlineKeyboardMarkup, InputFile, InputMedia,
InputMediaAnimation, InputMediaPhoto, InputMediaVideo, Message, MessageId, ParseMode,
ReplyParameters,
};
use teloxide::{ApiError, RequestError};
use x_media::site::FetchError;
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum MediaItemPayload {
Photo {
media: String,
has_spoiler: bool,
},
Video {
media: String,
has_spoiler: bool,
thumbnail: Option<String>,
},
Animation {
media: String,
has_spoiler: bool,
},
}
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Task {
SendMediaSequence {
chat_id: i64,
reply_to_message_id: i64,
caption: String,
media_batches: Vec<Vec<MediaItemPayload>>,
batch_index: usize,
sent_message_ids: Vec<i64>,
source_url: String,
edit_before_forward: bool,
forward_channel_id: Option<i64>,
notify_chat_id: Option<i64>,
notify_message_id: Option<i64>,
},
SendAnimation {
chat_id: i64,
reply_to_message_id: i64,
caption: String,
animation: MediaItemPayload,
source_url: String,
edit_before_forward: bool,
forward_channel_id: Option<i64>,
notify_chat_id: Option<i64>,
notify_message_id: Option<i64>,
},
ForwardMessages {
from_chat_id: i64,
to_chat_id: i64,
message_ids: Vec<i64>,
notify_chat_id: Option<i64>,
notify_message_id: Option<i64>,
},
}
pub const MAX_MEDIA_GROUP: usize = 9;
pub const MAX_UPLOAD_BYTES: u64 = 50 * 1024 * 1024; // Telegram Bot API upload cap
/// Splits media into batches of at most [`MAX_MEDIA_GROUP`] items.
pub fn chunk_media_items<T: Clone>(items: Vec<T>) -> Vec<Vec<T>> {
items.chunks(MAX_MEDIA_GROUP).map(|chunk| chunk.to_vec()).collect()
}
/// Exponential backoff with jitter, capped at 30s.
pub fn retry_delay_seconds(attempts: u32) -> f64 {
let jitter: f64 = rand::thread_rng().gen_range(0.2..0.8);
(2f64.powi(attempts as i32) + jitter).min(30.0)
}
/// Telegram's servers failed to fetch a media URL (hotlink protection etc.):
/// these errors are handled by the download-and-reupload fallback, NOT by a
/// queue retry (resending the URL cannot succeed).
pub fn is_media_fetch_failure(e: &ApiError) -> bool {
const MARKERS: [&str; 5] = [
"webpage_media_empty",
"media_empty",
"empty_web_media",
"webpage_curl_failed",
"timeout",
];
let description = e.to_string().to_lowercase();
MARKERS.iter().any(|marker| description.contains(marker))
}
/// Task-free classification of a Telegram request error. The callers attach
/// the (updated) task when building a [`SendError`].
pub enum Classification {
Retryable { delay_seconds: f64 },
Permanent { message: String },
/// Handled by the download fallback, not a queue retry.
MediaFetchFailure,
}
pub fn classify_request_error(e: &RequestError) -> Classification {
match e {
RequestError::RetryAfter(seconds) => {
Classification::Retryable { delay_seconds: seconds.seconds() as f64 }
}
RequestError::Network(_) => Classification::Retryable {
delay_seconds: retry_delay_seconds(0),
},
RequestError::Api(api) if is_media_fetch_failure(api) => Classification::MediaFetchFailure,
RequestError::Api(api) => Classification::Permanent { message: api.to_string() },
RequestError::MigrateToChatId(_)
| RequestError::InvalidJson { .. }
| RequestError::Io(_) => Classification::Permanent { message: e.to_string() },
}
}
pub enum SendError {
Retryable { delay_seconds: f64, task: Task },
Permanent { message: String, task: Task },
}
fn classify_to_send_error(e: &RequestError, task: Task) -> SendError {
match classify_request_error(e) {
Classification::Retryable { delay_seconds } => SendError::Retryable {
delay_seconds,
task,
},
Classification::Permanent { message } => SendError::Permanent { message, task },
Classification::MediaFetchFailure => SendError::Permanent {
message: "media fetch failed".into(),
task,
},
}
}
fn parse_media_url(s: &str) -> Result<url::Url, String> {
url::Url::parse(s).map_err(|e| format!("invalid media URL: {e}"))
}
fn item_url(item: &MediaItemPayload) -> &str {
match item {
MediaItemPayload::Photo { media, .. }
| MediaItemPayload::Video { media, .. }
| MediaItemPayload::Animation { media, .. } => media,
}
}
/// Remote http(s) URLs are handed to Telegram to fetch; everything else
/// (e.g. a locally encoded ugoira MP4) is uploaded directly.
fn input_file_for(media: &str) -> Result<InputFile, String> {
if media.starts_with("http://") || media.starts_with("https://") {
Ok(InputFile::url(parse_media_url(media)?))
} else {
Ok(InputFile::file(media))
}
}
fn photo_media(file: InputFile, caption: Option<&str>, spoiler: bool) -> InputMedia {
let mut photo = InputMediaPhoto::new(file).parse_mode(ParseMode::Html);
if let Some(caption) = caption {
photo = photo.caption(caption);
}
if spoiler {
photo = photo.spoiler();
}
InputMedia::Photo(photo)
}
fn video_media(file: InputFile, caption: Option<&str>, spoiler: bool) -> InputMedia {
let mut video = InputMediaVideo::new(file).parse_mode(ParseMode::Html);
if let Some(caption) = caption {
video = video.caption(caption);
}
if spoiler {
video = video.spoiler();
}
InputMedia::Video(video)
}
fn animation_media(file: InputFile, caption: Option<&str>, spoiler: bool) -> InputMedia {
let mut animation = InputMediaAnimation::new(file).parse_mode(ParseMode::Html);
if let Some(caption) = caption {
animation = animation.caption(caption);
}
if spoiler {
animation = animation.spoiler();
}
InputMedia::Animation(animation)
}
/// Builds a media group from payloads; only the first item of the batch gets
/// the caption (Telegram rejects captions on later items).
fn build_media_group(
batch: &[MediaItemPayload],
caption: Option<&str>,
) -> Result<Vec<InputMedia>, String> {
batch
.iter()
.enumerate()
.map(|(i, item)| {
let item_caption = if i == 0 { caption } else { None };
Ok(match item {
MediaItemPayload::Photo {
media,
has_spoiler,
} => photo_media(input_file_for(media)?, item_caption, *has_spoiler),
MediaItemPayload::Video {
media,
has_spoiler,
thumbnail,
} => {
let mut video = video_media(input_file_for(media)?, item_caption, *has_spoiler);
if let (Some(thumb), InputMedia::Video(v)) = (thumbnail, &mut video) {
*v = v.clone().thumbnail(input_file_for(thumb)?);
}
video
}
MediaItemPayload::Animation {
media,
has_spoiler,
} => animation_media(input_file_for(media)?, item_caption, *has_spoiler),
})
})
.collect()
}
/// Infers a file extension from magic bytes so Telegram detects the mime type
/// on multipart uploads.
fn sniff_ext(bytes: &[u8]) -> &'static str {
if bytes.starts_with(&[0xFF, 0xD8]) {
"jpg"
} else if bytes.starts_with(b"\x89PNG") {
"png"
} else if bytes.starts_with(b"RIFF") && bytes.len() >= 12 && &bytes[8..12] == b"WEBP" {
"webp"
} else if bytes.starts_with(b"GIF8") {
"gif"
} else if bytes.len() >= 12 && &bytes[4..8] == b"ftyp" {
"mp4"
} else {
"bin"
}
}
enum FallbackError {
Retryable { delay_seconds: f64 },
Permanent { message: String },
}
/// Downloads one media item to a temp file (deleted on drop). Network errors
/// are retryable; size over the upload cap and other download errors are not.
async fn download_to_temp(item: &MediaItemPayload) -> Result<NamedTempFile, FallbackError> {
let media_url = match item {
MediaItemPayload::Photo { media, .. }
| MediaItemPayload::Video { media, .. }
| MediaItemPayload::Animation { media, .. } => media,
};
let bytes = match x_media::site::download_media(media_url).await {
Ok(bytes) => bytes,
Err(FetchError::Http(_)) => {
return Err(FallbackError::Retryable {
delay_seconds: retry_delay_seconds(0),
});
}
Err(e) => {
return Err(FallbackError::Permanent {
message: format!("download failed: {e}"),
});
}
};
if bytes.len() as u64 > MAX_UPLOAD_BYTES {
return Err(FallbackError::Permanent {
message: "media too large".into(),
});
}
let ext = sniff_ext(&bytes);
let mut file = tempfile::Builder::new()
.suffix(&format!(".{ext}"))
.tempfile()
.map_err(|e| FallbackError::Permanent {
message: format!("temp file failed: {e}"),
})?;
use std::io::Write;
file.as_file_mut()
.write_all(&bytes)
.map_err(|e| FallbackError::Permanent {
message: format!("temp file write failed: {e}"),
})?;
Ok(file)
}
/// Download-and-reupload fallback for one media batch.
async fn send_batch_via_upload(
bot: &Bot,
chat_id: i64,
reply_to: i64,
batch: &[MediaItemPayload],
caption: Option<&str>,
) -> Result<Vec<Message>, FallbackError> {
let mut files = Vec::new();
let mut items = Vec::new();
for (i, item) in batch.iter().enumerate() {
let file = download_to_temp(item).await?;
let path = file.path().to_path_buf();
let item_caption = if i == 0 { caption } else { None };
let media = match item {
MediaItemPayload::Photo { has_spoiler, .. } => {
photo_media(InputFile::file(path), item_caption, *has_spoiler)
}
MediaItemPayload::Video { has_spoiler, .. } => {
video_media(InputFile::file(path), item_caption, *has_spoiler)
}
MediaItemPayload::Animation { has_spoiler, .. } => {
animation_media(InputFile::file(path), item_caption, *has_spoiler)
}
};
items.push(media);
files.push(file);
}
let result = bot
.send_media_group(ChatId(chat_id), items)
.reply_parameters(ReplyParameters::new(MessageId(reply_to as i32)).allow_sending_without_reply())
.await;
match result {
Ok(messages) => Ok(messages),
Err(e) => Err(match classify_request_error(&e) {
Classification::Retryable { delay_seconds } => FallbackError::Retryable {
delay_seconds,
},
Classification::Permanent { message } => FallbackError::Permanent { message },
Classification::MediaFetchFailure => FallbackError::Permanent {
message: "upload failed".into(),
},
}),
}
}
fn updated_sequence_task(task: &Task, batch_index: usize, sent_message_ids: Vec<i64>) -> Task {
match task {
Task::SendMediaSequence {
chat_id,
reply_to_message_id,
caption,
media_batches,
source_url,
edit_before_forward,
forward_channel_id,
notify_chat_id,
notify_message_id,
..
} => Task::SendMediaSequence {
chat_id: *chat_id,
reply_to_message_id: *reply_to_message_id,
caption: caption.clone(),
media_batches: media_batches.clone(),
batch_index,
sent_message_ids,
source_url: source_url.clone(),
edit_before_forward: *edit_before_forward,
forward_channel_id: *forward_channel_id,
notify_chat_id: *notify_chat_id,
notify_message_id: *notify_message_id,
},
_ => unreachable!("updated_sequence_task requires a SendMediaSequence task"),
}
}
/// Sends the media batches starting at `task.batch_index`, extending
/// `sent_message_ids`. Returns all sent message ids on full success; on
/// failure returns a [`SendError`] whose task carries the resumed state.
pub async fn send_media_sequence(bot: &Bot, task: &Task) -> Result<Vec<i64>, SendError> {
let Task::SendMediaSequence {
chat_id,
reply_to_message_id,
caption,
media_batches,
batch_index,
sent_message_ids,
..
} = task
else {
unreachable!("send_media_sequence requires a SendMediaSequence task")
};
let chat_id = *chat_id;
let reply_to = *reply_to_message_id;
let mut sent = sent_message_ids.clone();
for idx in *batch_index..media_batches.len() {
let batch = &media_batches[idx];
let caption = if idx == 0 { Some(caption.as_str()) } else { None };
let items = match build_media_group(batch, caption) {
Ok(items) => items,
Err(message) => {
return Err(SendError::Permanent {
message,
task: updated_sequence_task(task, idx, sent),
});
}
};
match bot
.send_media_group(ChatId(chat_id), items)
.reply_parameters(ReplyParameters::new(MessageId(reply_to as i32)).allow_sending_without_reply())
.await
{
Ok(messages) => {
log::info!(
"media group batch {idx}/{} sent ({} item(s))",
media_batches.len(),
batch.len()
);
sent.extend(messages.into_iter().map(|m| m.id.0 as i64));
}
Err(RequestError::Api(api)) if is_media_fetch_failure(&api) => {
log::info!(
"Telegram could not fetch media for batch {idx} ({}), downloading and reuploading",
batch
.first()
.map(|item| item_url(item))
.unwrap_or("?")
);
match send_batch_via_upload(bot, chat_id, reply_to, batch, caption).await {
Ok(messages) => sent.extend(messages.into_iter().map(|m| m.id.0 as i64)),
Err(FallbackError::Retryable { delay_seconds }) => {
return Err(SendError::Retryable {
delay_seconds,
task: updated_sequence_task(task, idx, sent),
});
}
Err(FallbackError::Permanent { message }) => {
return Err(SendError::Permanent {
message,
task: updated_sequence_task(task, idx, sent),
});
}
}
}
Err(e) => {
return Err(classify_to_send_error(
&e,
updated_sequence_task(task, idx, sent),
));
}
}
}
Ok(sent)
}
async fn send_animation_inner(
bot: &Bot,
chat_id: i64,
reply_to: i64,
caption: &str,
spoiler: bool,
file: InputFile,
) -> Result<Message, RequestError> {
let mut request = bot
.send_animation(ChatId(chat_id), file)
.caption(caption)
.parse_mode(ParseMode::Html)
.reply_parameters(ReplyParameters::new(MessageId(reply_to as i32)).allow_sending_without_reply());
if spoiler {
request = request.has_spoiler(true);
}
request.await
}
/// Sends a lone animation (gif), URL first with the download fallback.
pub async fn send_animation(bot: &Bot, task: &Task) -> Result<Vec<i64>, SendError> {
let Task::SendAnimation {
chat_id,
reply_to_message_id,
caption,
animation,
..
} = task
else {
unreachable!("send_animation requires a SendAnimation task")
};
let chat_id = *chat_id;
let reply_to = *reply_to_message_id;
let (media_url, has_spoiler) = match animation {
MediaItemPayload::Animation {
media,
has_spoiler,
} => (media, *has_spoiler),
MediaItemPayload::Photo { .. } | MediaItemPayload::Video { .. } => {
unreachable!("SendAnimation carries an Animation payload")
}
};
let url_file = match input_file_for(media_url) {
Ok(file) => file,
Err(message) => return Err(SendError::Permanent { message, task: task.clone() }),
};
match send_animation_inner(bot, chat_id, reply_to, caption, has_spoiler, url_file)
.await
{
Ok(message) => Ok(vec![message.id.0 as i64]),
Err(RequestError::Api(api)) if is_media_fetch_failure(&api) => {
log::info!(
"Telegram could not fetch animation URL, downloading and reuploading: {}",
media_url
);
let file = match download_to_temp(animation).await {
Ok(file) => file,
Err(FallbackError::Retryable { delay_seconds }) => {
return Err(SendError::Retryable { delay_seconds, task: task.clone() });
}
Err(FallbackError::Permanent { message }) => {
return Err(SendError::Permanent { message, task: task.clone() });
}
};
let path = file.path().to_path_buf();
match send_animation_inner(
bot,
chat_id,
reply_to,
caption,
has_spoiler,
InputFile::file(path),
)
.await
{
Ok(message) => Ok(vec![message.id.0 as i64]),
Err(e) => Err(classify_to_send_error(&e, task.clone())),
}
}
Err(e) => Err(classify_to_send_error(&e, task.clone())),
}
}
/// Copies already-sent messages to the forward channel. No download fallback:
/// the files are already on Telegram's servers.
pub async fn forward_messages(bot: &Bot, task: &Task) -> Result<(), SendError> {
let Task::ForwardMessages {
from_chat_id,
to_chat_id,
message_ids,
..
} = task
else {
unreachable!("forward_messages requires a ForwardMessages task")
};
let message_ids = message_ids
.iter()
.map(|id| MessageId(*id as i32))
.collect::<Vec<_>>();
match bot
.copy_messages(ChatId(*to_chat_id), ChatId(*from_chat_id), message_ids.clone())
.await
{
Ok(_) => {
log::info!(
"copied {} message(s) from {} to {}",
message_ids.len(),
from_chat_id,
to_chat_id
);
Ok(())
}
Err(e) => Err(classify_to_send_error(&e, task.clone())),
}
}
/// One button per template name (column layout), then the confirm button.
pub fn build_edit_markup(templates: &HashMap<String, String>) -> InlineKeyboardMarkup {
let mut rows = Vec::new();
for name in templates.keys() {
rows.push(vec![InlineKeyboardButton::callback(
name.clone(),
format!("template|{name}"),
)]);
}
rows.push(vec![InlineKeyboardButton::callback(
"↩️ Confirm",
"forward",
)]);
InlineKeyboardMarkup::new(rows)
}
/// Notifies a chat about a dead-lettered task (skips when `notify_chat_id` is
/// absent).
pub async fn notify_failure(bot: &Bot, chat_id: Option<i64>, message_id: Option<i64>, message: &str) {
let Some(chat_id) = chat_id else { return };
let mut request = bot.send_message(ChatId(chat_id), message);
if let Some(message_id) = message_id {
request = request
.reply_parameters(ReplyParameters::new(MessageId(message_id as i32)).allow_sending_without_reply());
}
if let Err(e) = request.await {
log::error!("failed to notify about failed task: {e}");
}
}
/// After a successful send: either open the edit-before-forward prompt or
/// forward to the configured channel (with retry/queue handling).
pub async fn post_send_actions(bot: &Bot, task: &Task, message_ids: Vec<i64>) {
let (chat_id, reply_to, source_url, edit_before_forward, forward_channel_id, notify_chat_id, notify_message_id) =
match task {
Task::SendMediaSequence {
chat_id,
reply_to_message_id,
source_url,
edit_before_forward,
forward_channel_id,
notify_chat_id,
notify_message_id,
..
}
| Task::SendAnimation {
chat_id,
reply_to_message_id,
source_url,
edit_before_forward,
forward_channel_id,
notify_chat_id,
notify_message_id,
..
} => (
*chat_id,
*reply_to_message_id,
source_url.clone(),
*edit_before_forward,
*forward_channel_id,
*notify_chat_id,
*notify_message_id,
),
Task::ForwardMessages { .. } => return,
};
if edit_before_forward {
let mut chat_data = CHAT_STORE.get(chat_id).await;
let keyboard = build_edit_markup(&chat_data.template);
match bot
.send_message(ChatId(chat_id), "Reply to edit message.")
.reply_markup(keyboard)
.reply_parameters(
ReplyParameters::new(MessageId(reply_to as i32)).allow_sending_without_reply(),
)
.await
{
Ok(prompt) => {
log::info!(
"edit-before-forward prompt {} opened for {} message(s)",
prompt.id.0,
message_ids.len()
);
chat_data.edit_message.insert(
prompt.id.0 as i64,
EditMessage {
url: source_url,
chat_id,
forward_message_ids: message_ids,
template: String::new(),
created_at: unix_now(),
},
);
CHAT_STORE.set(chat_id, &chat_data).await;
}
Err(e) => log::error!("failed to send edit prompt: {e}"),
}
return;
}
if let Some(channel_id) = forward_channel_id {
log::info!("forwarding {} message(s) to channel {channel_id}", message_ids.len());
let forward_task = Task::ForwardMessages {
from_chat_id: chat_id,
to_chat_id: channel_id,
message_ids,
notify_chat_id,
notify_message_id,
};
match forward_messages(bot, &forward_task).await {
Ok(()) => {}
Err(SendError::Retryable { delay_seconds, task }) => {
let payload = serde_json::to_value(task).expect("task serializes");
let run_after = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
+ delay_seconds;
if let Err(e) = TASK_QUEUE.enqueue(payload, run_after).await {
log::error!("failed to enqueue forward retry: {e}");
}
}
Err(SendError::Permanent { message, .. }) => {
notify_failure(
bot,
notify_chat_id,
notify_message_id,
&format!("Task failed after retries: {message}"),
)
.await;
}
}
}
}
/// Queue entry point: parses the stored task and dispatches.
pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> {
let task: Task = match serde_json::from_value(payload.clone()) {
Ok(task) => task,
Err(e) => {
return Err(QueueError::Permanent {
message: format!("invalid task payload: {e}"),
payload,
});
}
};
let bot = Bot::from_env();
match task {
Task::SendMediaSequence { .. } | Task::SendAnimation { .. } => {
let message_ids = match send_media_or_animation(&bot, &task).await {
Ok(ids) => ids,
Err(SendError::Retryable { delay_seconds, task }) => {
return Err(QueueError::Retryable {
delay_seconds,
payload: serde_json::to_value(task).expect("task serializes"),
});
}
Err(SendError::Permanent { message, task }) => {
return Err(QueueError::Permanent {
message,
payload: serde_json::to_value(task).expect("task serializes"),
});
}
};
post_send_actions(&bot, &task, message_ids).await;
Ok(())
}
Task::ForwardMessages { .. } => match forward_messages(&bot, &task).await {
Ok(()) => Ok(()),
Err(SendError::Retryable { delay_seconds, task }) => Err(QueueError::Retryable {
delay_seconds,
payload: serde_json::to_value(task).expect("task serializes"),
}),
Err(SendError::Permanent { message, task }) => Err(QueueError::Permanent {
message,
payload: serde_json::to_value(task).expect("task serializes"),
}),
},
}
}
async fn send_media_or_animation(bot: &Bot, task: &Task) -> Result<Vec<i64>, SendError> {
match task {
Task::SendMediaSequence { .. } => send_media_sequence(bot, task).await,
Task::SendAnimation { .. } => send_animation(bot, task).await,
Task::ForwardMessages { .. } => unreachable!(),
}
}
/// Dead-letter callback wired to the queue in main: notifies the task's chat.
pub async fn dead_letter_notify(payload: serde_json::Value, message: String) {
let notify_chat_id = payload.get("notify_chat_id").and_then(|v| v.as_i64());
let notify_message_id = payload.get("notify_message_id").and_then(|v| v.as_i64());
if notify_chat_id.is_some() {
let bot = Bot::from_env();
notify_failure(
&bot,
notify_chat_id,
notify_message_id,
&format!("Task failed after retries: {message}"),
)
.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chunk_media_items_sizes() {
assert_eq!(chunk_media_items::<i32>(vec![]), Vec::<Vec<i32>>::new());
assert_eq!(chunk_media_items((0..9).collect()).len(), 1);
assert_eq!(chunk_media_items((0..10).collect()).len(), 2);
assert_eq!(chunk_media_items((0..25).collect()).len(), 3);
assert_eq!(chunk_media_items((0..25).collect())[2].len(), 7);
assert!(chunk_media_items((0..25).collect()).iter().all(|c| c.len() <= 9));
}
#[test]
fn retry_delay_seconds_bounds() {
for attempts in 0..10 {
let delay = retry_delay_seconds(attempts);
assert!(delay >= 1.0, "attempts={attempts}: {delay}");
assert!(delay <= 30.0, "attempts={attempts}: {delay}");
}
}
#[test]
fn is_media_fetch_failure_matches_markers() {
for description in [
"Bad Request: WEBPAGE_MEDIA_EMPTY",
"Bad Request: media_empty",
"Bad Request: EMPTY_WEB_MEDIA",
"Bad Request: webpage_curl_failed",
"Bad Request: request timeout",
] {
let api = ApiError::Unknown(description.to_string());
assert!(is_media_fetch_failure(&api), "{description}");
}
for description in ["Bad Request: message is not modified", "Forbidden: bot was blocked by the user"] {
let api = ApiError::Unknown(description.to_string());
assert!(!is_media_fetch_failure(&api), "{description}");
}
}
#[test]
fn classification_mapping() {
use teloxide::types::Seconds;
// RetryAfter -> Retryable with its delay
let e = RequestError::RetryAfter(Seconds::from_seconds(7));
assert!(matches!(
classify_request_error(&e),
Classification::Retryable { delay_seconds } if delay_seconds == 7.0
));
// Api error -> Permanent
let e = RequestError::Api(ApiError::Unknown("Bad Request: something".into()));
assert!(matches!(
classify_request_error(&e),
Classification::Permanent { .. }
));
// Api media-fetch marker -> MediaFetchFailure
let e = RequestError::Api(ApiError::Unknown("Bad Request: WEBPAGE_MEDIA_EMPTY".into()));
assert!(matches!(
classify_request_error(&e),
Classification::MediaFetchFailure
));
// MigrateToChatId -> Permanent
let e = RequestError::MigrateToChatId(ChatId(123));
assert!(matches!(
classify_request_error(&e),
Classification::Permanent { .. }
));
}
#[test]
fn task_serde_round_trip_preserves_resume_state() {
let task = Task::SendMediaSequence {
chat_id: 111,
reply_to_message_id: 222,
caption: "cap".into(),
media_batches: vec![
vec![MediaItemPayload::Photo {
media: "https://a/b.jpg".into(),
has_spoiler: true,
}],
vec![MediaItemPayload::Video {
media: "https://a/v.mp4".into(),
has_spoiler: false,
thumbnail: Some("https://a/t.jpg".into()),
}],
],
batch_index: 1,
sent_message_ids: vec![11, 12],
source_url: "https://x.com/u/status/1".into(),
edit_before_forward: true,
forward_channel_id: Some(333),
notify_chat_id: Some(111),
notify_message_id: Some(222),
};
let json = serde_json::to_value(&task).unwrap();
assert_eq!(json["type"], "send_media_sequence");
assert_eq!(json["batch_index"], 1);
let decoded: Task = serde_json::from_value(json).unwrap();
match decoded {
Task::SendMediaSequence {
batch_index,
sent_message_ids,
forward_channel_id,
media_batches,
..
} => {
assert_eq!(batch_index, 1);
assert_eq!(sent_message_ids, vec![11, 12]);
assert_eq!(forward_channel_id, Some(333));
assert_eq!(media_batches.len(), 2);
assert!(matches!(media_batches[0][0], MediaItemPayload::Photo { has_spoiler: true, .. }));
}
other => panic!("expected SendMediaSequence, got {other:?}"),
}
}
#[test]
fn media_item_payload_serde_tags() {
let photo = MediaItemPayload::Photo {
media: "https://a/b.jpg".into(),
has_spoiler: false,
};
let json = serde_json::to_value(&photo).unwrap();
assert_eq!(json["kind"], "photo");
}
#[test]
fn sniff_ext_detects_formats() {
assert_eq!(sniff_ext(&[0xFF, 0xD8, 0xFF, 0xE0]), "jpg");
assert_eq!(sniff_ext(b"\x89PNG\r\n\x1a\n"), "png");
assert_eq!(sniff_ext(b"RIFF\x00\x00\x00\x00WEBPVP8 "), "webp");
assert_eq!(sniff_ext(b"GIF89a"), "gif");
assert_eq!(sniff_ext(b"\x00\x00\x00\x18ftypisom"), "mp4");
assert_eq!(sniff_ext(b"something else"), "bin");
}
}
+151
View File
@@ -0,0 +1,151 @@
//! 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)?;
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.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
}
}
+17 -9
View File
@@ -1,22 +1,30 @@
services:
tgxmb: image: yoursfunny/telegram-twitter-media-bot: latest
tgxmb:
image: yoursfunny/telegram-twitter-media-bot:latest
restart: always
# ports:
# - "8443:8443"
environment:
# docker-entrypoint.sh drops privileges to this uid.
LOCAL_USER_ID: '1000'
BOT_TOKEN: ''
# Bot token (BotFather). Required.
TELOXIDE_TOKEN: ''
# Comma-separated admin chat ids; receives startup/shutdown notices.
BOT_ADMIN: ''
# Required for pixiv support; pixiv is disabled when unset.
PIXIV_REFRESH_TOKEN: ''
WEBHOOK: false
WEBHOOK_LISTEN: '127.0.0.1'
WEBHOOK_PORT: 8443
# Edit-before-forward records expire after this many seconds (default 86400 = 24h).
EDIT_MESSAGE_TTL_SECONDS: '86400'
RUST_LOG: 'info'
# Webhook mode is off by default (polling). The listener binds inside the
# container, so use 0.0.0.0 and publish the port if you enable it.
WEBHOOK: 'false'
WEBHOOK_LISTEN: '0.0.0.0'
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
volumes:
- ./data:/app/data
# - ./cert:/app/cert
container_name: tgxmb
+3 -1
View File
@@ -10,7 +10,9 @@ then
chown -R `id -u user`:`id -u user` /app > /dev/null 2>&1
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 "$@"
-336
View File
@@ -1,336 +0,0 @@
from __future__ import annotations
import html
from functools import wraps
from typing import TYPE_CHECKING
from telegram import InlineKeyboardButton, InlineKeyboardMarkup, MessageEntity
from telegram.constants import ChatAction, ChatType, ParseMode
from telegram.ext import (ApplicationBuilder, CallbackQueryHandler, CommandHandler, ContextTypes, Defaults,
InlineQueryHandler, MessageHandler, PicklePersistence, filters)
import common
import utils.regex as regex
from utils.context import ChatData, CustomContext, EditMessage
from utils.logger import get_logger
from utils.net import NetClient
from utils.pixiv import ProcessPixiv
from utils.telegram import Telegram
if TYPE_CHECKING:
from telegram import Message, Update
from telegram.ext import Application
logger = get_logger(__name__)
def send_action(action):
def decorator(func):
@wraps(func)
async def command_func(update: Update, context: CustomContext, *args, **kwargs):
try:
await update.effective_chat.send_action(action)
finally:
return await func(update, context, *args, **kwargs)
return command_func
return decorator
def extract_urls(message: Message) -> set[str]:
types = [MessageEntity.URL, MessageEntity.TEXT_LINK]
res = message.parse_entities(types)
res.update(message.parse_caption_entities(types))
res.update({key: key.url for key in res if key.type == MessageEntity.TEXT_LINK})
return set(res.values())
async def inline_query(update: Update, context: CustomContext) -> None:
query = update.inline_query.query
if query == "":
return
logger.info(f"Query: {query}")
async with Telegram(query) as tweet:
if not tweet:
return
await update.inline_query.answer(tweet.inline_query_result())
@send_action(ChatAction.UPLOAD_PHOTO)
async def url_media(update: Update, context: CustomContext, url: str) -> None:
async with Telegram(url) as tweet:
if not tweet:
return
media = tweet.message_media_result()
if not media:
await update.effective_message.reply_text(
"No media found or media type is not supported.",
reply_to_message_id=update.message.message_id,
)
return
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 isinstance(media[0], tuple) 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.chat_data.edit_before_forward:
message_reply = await update.effective_message.reply_text(
"Reply to edit message.",
reply_markup=InlineKeyboardMarkup.from_column(
[InlineKeyboardButton(name, callback_data=f"template|{name}") for name in
context.chat_data.template.keys()] + [InlineKeyboardButton("↩️ Confirm", callback_data="forward")]
),
reply_to_message_id=update.message.message_id,
)
context.chat_data.edit_message[message_reply.id] = EditMessage(
url=url,
forward=message_to_send
)
return
if context.chat_data.forward_channel_id:
await forward_message(update, context, message_to_send)
async def handel_url_media(update: Update, context: CustomContext) -> None:
url = update.message.text
logger.info(f"Receiving url: {url}")
await url_media(update, context, url)
async def forward_message(
update: Update,
context: CustomContext,
message_to_send: tuple[Message, ...],
) -> None:
try:
await update.effective_chat.copy_messages(
context.chat_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: CustomContext) -> bool:
if not (reply := update.message.reply_to_message):
return False
_edit_message = context.chat_data.edit_message.get(reply.id, None)
if not _edit_message:
return False
new_text = '<a href="{0}">{1}</a>'.format(
_edit_message.url,
html.escape(update.message.text)
)
update_text = context.chat_data.template[template].replace("[]", new_text) if (
template := _edit_message.template) else new_text
await _edit_message.forward[0].edit_caption(update_text)
return True
async def handle_message(update: Update, context: CustomContext) -> None:
if await edit_message(update, context):
return
if not (urls := extract_urls(update.message)):
return
for url in urls:
await url_media(update, context, url)
async def query_forward_message(update: Update, context: CustomContext) -> None:
_edit_message = context.chat_data.edit_message[update.effective_message.id]
await forward_message(update, context, _edit_message.forward)
await update.callback_query.answer('✅ Forwarded')
await update.callback_query.delete_message()
del _edit_message
async def query_template(update: Update, context: CustomContext) -> None:
query = update.callback_query
await query.answer()
name = query.data.split("|")[1]
_edit_message = context.chat_data.edit_message[query.message.message_id]
_edit_message.template = name
await _edit_message.forward[0].edit_caption(context.chat_data.template[name])
@send_action(ChatAction.TYPING)
async def cmd_start(update: Update, context: CustomContext) -> None:
await update.effective_message.reply_text(
"Welcome to the Twitter Fetcher Bot!\n"
"You can use this bot to fetch tweets from Twitter and forward them to a channel.\n"
"Use /set_forward_channel to set a channel to forward tweets.\n"
"Use /remove_forward_channel to remove the channel.\n"
"Use /edit_before_forward to enable or disable edit before forward.\n"
"Use /set_template to set a template for the forwarded message.\n"
"Use /bot_dict to see the bot's data.\n"
"Use /clear_edit_message to clear the edit message cache.\n"
"You can also reply to a message with a tweet URL to fetch the tweet and forward it to the channel.\n"
"You can also use inline query to search for tweets."
)
@send_action(ChatAction.TYPING)
async def cmd_set_forward_channel(update: Update, context: CustomContext) -> 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 = 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 = filter(lambda x: x.user.id == update.effective_user.id, channel_admin)
user = next(user, None)
if not user:
await update.effective_message.reply_text("You are not an admin of the channel.")
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.chat_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: CustomContext) -> None:
if context.chat_data.forward_channel_id:
context.chat_data.forward_channel_id = None
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: CustomContext) -> None:
if context.chat_data.forward_channel_id is None:
await update.effective_message.reply_text("Please enable forward channel first.")
return
if context.chat_data.edit_before_forward:
context.chat_data.edit_before_forward = False
context.chat_data.edit_message.clear()
await update.effective_message.reply_text("Disable edit before forward.")
return
context.chat_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: CustomContext) -> 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 (template := reply.text_html):
await update.effective_message.reply_text("Please reply to a message with [] to set as template.")
return
if not context.args:
await update.effective_message.reply_text("Please provide a name for the template.")
return
context.chat_data.template[''.join(context.args)] = template
await update.effective_message.reply_text("Template set.")
@send_action(ChatAction.TYPING)
async def cmd_user_dict(update: Update, context: CustomContext) -> None:
await update.effective_message.reply_text(html.escape(str(context.chat_data)), disable_web_page_preview=True)
@send_action(ChatAction.TYPING)
async def cmd_clear_edit_message(update: Update, context: CustomContext) -> None:
context.chat_data.edit_message.clear()
await update.effective_message.reply_text("Edit message cleared.")
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)
NetClient.init_client()
if common.PIXIV_REFRESH_TOKEN:
await ProcessPixiv.init_client(common.PIXIV_REFRESH_TOKEN)
async def post_stop(application: Application) -> None:
if common.ADMIN:
await application.bot.send_message(common.ADMIN[0], "Shutting down...")
async def post_shutdown(application: Application) -> None:
await NetClient.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)
.context_types(ContextTypes(context=CustomContext, chat_data=ChatData))
.post_init(post_init)
.post_stop(post_stop)
.post_shutdown(post_shutdown)
.concurrent_updates(True)
.http_version('2')
.build()
)
user_filter = filters.User()
user_filter.add_user_ids(common.ADMIN)
handlers = [
InlineQueryHandler(inline_query),
MessageHandler((filters.Regex(regex.x_url) | filters.Regex(regex.pixiv_url)) | filters.Regex(
regex.bsky_url) & filters.ChatType.PRIVATE,
handel_url_media),
CommandHandler("start", cmd_start),
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, handle_message),
CallbackQueryHandler(query_forward_message, pattern="forward"),
CallbackQueryHandler(query_template, pattern=r"^template\|"),
CommandHandler("bot_dict", cmd_user_dict),
CommandHandler("clear_edit_message", cmd_clear_edit_message),
]
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()
-4
View File
@@ -1,4 +0,0 @@
python-telegram-bot[webhooks]~=22.8
httpx[http2]~=0.27
uvloop~=0.22; sys_platform != 'win32'
async-pixiv~=1.1.2
View File
-163
View File
@@ -1,163 +0,0 @@
from __future__ import annotations
from functools import cached_property
from typing import TYPE_CHECKING
from utils.net import NetClient
from utils.regex import bsky_url
if TYPE_CHECKING:
from utils.types import BskyEmbedImages, BskyInfo, BskyEmbedVideo, BskyEmbedExternal
bsky_api_url = 'https://public.api.bsky.app/xrpc/app.bsky.feed.getPostThread'
SENSITIVE_TAG = {'sexual', 'nudity', 'porn', 'graphic-media'}
class BskyMedia:
__slots__ = ('_url', '_thumb', '_type', '__dict__')
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"BskyMedia(url={self.url}, thumb={self.thumb}, type={self.type})"
@property
def url(self) -> str:
return self._url
@property
def thumb(self) -> str:
return self._thumb
@property
def type(self) -> str:
return self._type
class Bsky:
__slots__ = ('_id', '_author', '_author_id', '_text', '_media', '_sensitive', '__dict__')
def __init__(
self,
id: str,
author: str,
author_id: str,
text: str,
media: list[BskyMedia],
sensitive: bool = False
):
self._id: str = id
self._author: str = author
self._author_id: str = author_id
self._text: str = text
self._media: list[BskyMedia] = media
self._sensitive: bool = sensitive
@property
def id(self) -> str:
return self._id
@cached_property
def url(self) -> str:
return f"https://bsky.app/profile/{self._author_id}/post/{self._id}"
@property
def author(self) -> str:
return self._author
@cached_property
def author_url(self) -> str:
return f"https://bsky.app/profile/{self._author_id}"
@property
def text(self) -> str:
return self._text
@property
def media(self) -> list[BskyMedia]:
return self._media
@property
def sensitive(self) -> bool:
return self._sensitive
class ProcessBsky:
__slots__ = ('_url', '_id', '_bsky')
def __init__(self, url: str):
self._url: str = url
async def __aenter__(self):
bsky = await self._fetch_bsky()
if not bsky['thread'].get('post'):
raise ValueError(f"BSky post not found: {bsky}")
self._bsky = bsky['thread']['post']
return Bsky(
id=self._id,
author=self._bsky['author']['displayName'],
author_id=self._bsky['author']['handle'],
text=self._bsky['record']['text'],
media=self._bsky_media,
sensitive=self._sensitive
)
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
async def _fetch_bsky(self) -> BskyInfo:
match = bsky_url.match(self._url)
if not match:
raise ValueError(f"Invalid Bsky URL: {self._url}")
auther_id, self._id = match.groups()
return await NetClient.fetch_json(
bsky_api_url,
params={'uri': f'at://{auther_id}/app.bsky.feed.post/{self._id}', 'depth': 0}
)
@property
def _bsky_media(self) -> list[BskyMedia]:
if not self._bsky.get('embed'):
return []
match (embed := self._bsky['embed'])['$type']:
case 'app.bsky.embed.images#view':
embed: BskyEmbedImages
return [
BskyMedia(
url=image['fullsize'],
thumb=image['thumb'],
media_type='image'
)
for image in embed['images']
]
case 'app.bsky.embed.video#view':
embed: BskyEmbedVideo
return [
BskyMedia(
url=embed['playlist'],
thumb=embed['thumbnail'],
media_type='video'
)
]
case 'app.bsky.embed.external#view':
embed: BskyEmbedExternal
return [
BskyMedia(
url=embed['external']['uri'],
thumb=embed['external']['thumb'],
media_type='external'
)
]
case _:
raise NotImplementedError(f"Unknown Bsky embed type: {embed['$type']}")
@property
def _sensitive(self) -> bool:
return any(
tag in label['val']
for label in self._bsky['labels']
for tag in SENSITIVE_TAG
)
-42
View File
@@ -1,42 +0,0 @@
import dataclasses
from typing import Optional
from telegram import Message
from telegram.ext import Application, CallbackContext, ExtBot
@dataclasses.dataclass(repr=False)
class EditMessage:
url: str
forward: tuple[Message, ...]
template: str = ""
def __str__(self):
forward = ", ".join(f"Message({f.id})" for f in self.forward)
return f"EditMessage(url={self.url}, forward={forward}, template={self.template})"
__repr__ = __str__
class ChatData:
def __init__(self):
self.forward_channel_id: Optional[int] = None
self.edit_before_forward: bool = False
self.edit_message: dict[int, EditMessage] = {}
self.template: dict[str, str] = {}
def __str__(self):
return f"ChatData(forward_channel_id={self.forward_channel_id}, edit_before_forward={self.edit_before_forward}, " \
f"edit_message={self.edit_message}, template={self.template})"
__repr__ = __str__
class CustomContext(CallbackContext[ExtBot, dict, ChatData, dict]):
def __init__(
self,
application: Application,
chat_id: Optional[int] = None,
user_id: Optional[int] = None
):
super().__init__(application=application, chat_id=chat_id, user_id=user_id)
-11
View File
@@ -1,11 +0,0 @@
import logging
import os
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)
-37
View File
@@ -1,37 +0,0 @@
from __future__ import annotations
from httpx import AsyncClient
def create_client() -> AsyncClient:
return AsyncClient(http2=True)
async def close_client(_client: AsyncClient) -> None:
return await _client.aclose()
async def fetch_json(_client: AsyncClient, url: str, params: dict = None) -> dict:
response = await _client.get(url, params=params)
assert response.is_success, f"Failed to fetch {url}, status code {response.status_code}"
return response.json()
class NetClient:
_httpx_client: AsyncClient
@classmethod
def init_client(cls) -> None:
cls._httpx_client = create_client()
@classmethod
async def close_client(cls) -> None:
await close_client(cls._httpx_client)
@classmethod
def get_client(cls) -> AsyncClient:
return cls._httpx_client
@classmethod
async def fetch_json(cls, url: str, params: dict = None) -> dict:
return await fetch_json(cls._httpx_client, url, params)
-141
View File
@@ -1,141 +0,0 @@
from __future__ import annotations
from typing import Literal, TYPE_CHECKING
from async_pixiv import PixivClient
from async_pixiv.error import APIError
if TYPE_CHECKING:
from async_pixiv.model.illust import Illust
class PixivMedia:
__slots__ = ('_url', '_thumb')
def __init__(self, url: str, thumb: str):
self._url: str = url
self._thumb: str = thumb
def __str__(self):
return f"PixivMedia(url={self.url}, thumb={self.thumb}, large={self.large})"
@property
def url(self) -> str:
return self._url
@property
def thumb(self) -> str:
return self._thumb
@property
def large(self) -> str:
url = self._url.replace("img-original", "img-master").removesuffix(".jpg").removesuffix(".png")
return url + "_master1200.jpg"
class Pixiv:
__slots__ = ('_illust',)
def __init__(self, illust: Illust):
self._illust: Illust = illust
@property
def url(self) -> str:
return str(self._illust.link).rstrip('/')
@property
def type(self) -> Literal["illust", "manga", "ugoira"]:
return self._illust.type.value
@property
def title(self) -> str:
return self._illust.title
@property
def author(self) -> str:
return self._illust.user.name
@property
def author_url(self) -> str:
return str(self._illust.user.link).rstrip('/')
@property
def description(self) -> str:
return self._illust.caption
@property
def tags(self) -> list[str]:
return [tag.name for tag in self._illust.tags]
@property
def is_multiple_pages(self) -> bool:
return self._illust.page_count > 1
@property
def is_nsfw(self) -> bool:
return self._illust.sanity_level > 5
@property
def is_ai(self) -> bool:
return self._illust.ai_type == 2
@property
def images(self) -> list[PixivMedia]:
if self.is_multiple_pages:
return [
PixivMedia(
url=str(page.image_urls.original),
thumb=str(page.image_urls.medium)
)
for page in self._illust.meta_pages
]
else:
return [
PixivMedia( # should observe if image_url.original is always None for single page
url=str(self._illust.image_urls.original or self._illust.meta_single_page.original),
thumb=str(self._illust.image_urls.medium)
)
]
class _ProcessPixiv:
_client: PixivClient
@classmethod
async def init_client(cls, token: str) -> None:
cls._client = PixivClient()
await cls._client.login_with_token(token)
cls._token = token
@classmethod
async def close_client(cls) -> None:
await cls._client.close()
@classmethod
async def refresh_token(cls):
await cls._client.login_with_token(cls._token)
class ProcessPixiv(_ProcessPixiv):
__slots__ = ('_url', '_illust')
def __init__(self, url: str):
self._url: str = url
async def __aenter__(self):
self._illust = await self._fetch_illust()
return Pixiv(self._illust)
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
async def _fetch_illust(self) -> Illust:
illust_id = self._parse_illust_id()
try:
return (await self._client.ILLUST.detail(illust_id)).illust
except APIError:
await self.refresh_token()
return (await self._client.ILLUST.detail(illust_id)).illust # TODO use retry here
def _parse_illust_id(self) -> int:
return int(self._url.split("/")[-1])
-11
View File
@@ -1,11 +0,0 @@
import re
x_url = re.compile(
r"^(?:https?://)?(?:www\.|mobile\.)?(?:x|twitter|fixvx|vxtwitter|fixupx|fxtwitter)\.com/(.+)/status/(\d+)")
x_media_url = re.compile(r"^(?:https?://)?(pbs|video)\.twimg\.com/(.*)")
x_tco_url = re.compile(r"(?:https?://)?t\.co/.+$", re.M)
message_url = re.compile(r"\[.+]", re.S)
pixiv_url = re.compile(r"^(?:https?://)?(?:www\.)?pixiv\.net/(?:en/)?(?:artworks/|i/)(\d+)")
bsky_url = re.compile(r"^(?:https?://)?bsky\.app/profile/(.+)/post/(.+)")
-288
View File
@@ -1,288 +0,0 @@
from __future__ import annotations
import html
from functools import cached_property
from typing import Generator, TYPE_CHECKING
from uuid import uuid4
from telegram import InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, InputMediaPhoto, \
InputMediaVideo
from common import PIXIV_REFRESH_TOKEN
from .bsky import ProcessBsky
from .logger import get_logger
from .pixiv import ProcessPixiv
from .regex import bsky_url, pixiv_url, x_url
from .tweet import ProcessTweet
if TYPE_CHECKING:
from .types import TypeInlineQueryResult, TypeMessageMediaResult
logger = get_logger(__name__)
message_raw_text_tweet = """{url}
<a href="{author_url}">{author}</a>: {text}
"""
message_raw_text_pixiv = """<a href="{url}">{text}</a> / <a href="{author_url}">{author}</a>
{tags}
"""
class Telegram:
def __init__(self, url: str):
self._url = url
async def __aenter__(self):
if x_url.match(self._url):
async with TelegramTweet(self._url) as tweet:
return tweet
elif PIXIV_REFRESH_TOKEN and pixiv_url.match(self._url):
async with TelegramPixiv(self._url) as pixiv:
return pixiv
elif bsky_url.match(self._url):
async with TelegramBsky(self._url) as bsky:
return bsky
else:
return None # TODO add raise and catch
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
class TelegramTweet:
message_raw_text = message_raw_text_tweet
__slots__ = ('_url', '_tweet', '__dict__')
def __init__(self, url: str):
self._url: str = url
async def __aenter__(self):
async with ProcessTweet(self._url) as tweet:
self._tweet = tweet
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
@property
def url(self) -> str:
return self._tweet.url
@cached_property
def message_text(self) -> str:
tweet = self._tweet
return self.message_raw_text.format(
url=tweet.url,
author_url=tweet.author_url,
author=html.escape(tweet.author),
text=html.escape(tweet.text)
)
def inline_query_result(self) -> tuple[TypeInlineQueryResult, ...]:
return tuple(self.inline_query_generator())
def message_media_result(self) -> tuple[TypeMessageMediaResult, ...]:
return tuple(self.message_media_generator())
def inline_query_generator(self) -> Generator[TypeInlineQueryResult, None, None]:
tweet = self._tweet
for tweet_media in tweet.media:
logger.info(str(tweet_media))
if tweet_media.type == "image":
yield InlineQueryResultPhoto(
id=str(uuid4()),
photo_url=tweet_media.url,
thumbnail_url=tweet_media.thumb,
caption=self.message_text
)
elif tweet_media.type == "video":
yield InlineQueryResultVideo(
id=str(uuid4()),
video_url=tweet_media.url,
mime_type="video/mp4",
thumbnail_url=tweet_media.thumb,
title=tweet.text,
caption=self.message_text
)
elif tweet_media.type == "gif":
yield InlineQueryResultMpeg4Gif(
id=str(uuid4()),
mpeg4_url=tweet_media.url,
thumbnail_url=tweet_media.thumb,
caption=self.message_text
)
def message_media_generator(self) -> Generator[TypeMessageMediaResult, None, None]:
tweet = self._tweet
for tweet_media in tweet.media:
logger.info(str(tweet_media))
if tweet_media.type == "image":
yield InputMediaPhoto(
media=tweet_media.url,
has_spoiler=tweet.sensitive
)
elif tweet_media.type == "video":
yield InputMediaVideo(
media=tweet_media.url,
has_spoiler=tweet.sensitive,
thumbnail=tweet_media.thumb
)
elif tweet_media.type == "gif":
if len(tweet.media) == 1:
yield tweet_media.url, tweet.sensitive
yield InputMediaVideo(
media=tweet_media.url,
has_spoiler=tweet.sensitive,
thumbnail=tweet_media.thumb
)
class TelegramPixiv:
message_raw_text = message_raw_text_pixiv
__slots__ = ('_url', '_pixiv')
def __init__(self, url: str):
self._url = url
async def __aenter__(self):
async with ProcessPixiv(self._url) as pixiv:
self._pixiv = pixiv
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
@property
def url(self) -> str:
return self._pixiv.url
@property
def message_text(self) -> str:
pixiv = self._pixiv
return self.message_raw_text.format(
url=pixiv.url,
author_url=pixiv.author_url,
author=html.escape(pixiv.author),
text=html.escape(pixiv.title),
tags=html.escape(" ".join(f"#{name}" for name in pixiv.tags))
)
def inline_query_result(self) -> tuple[TypeInlineQueryResult, ...]:
return tuple(i for i in self.inline_query_generator() if i)
def message_media_result(self) -> tuple[TypeMessageMediaResult, ...]:
return tuple(i for i in self.message_media_generator() if i)
def inline_query_generator(self) -> Generator[TypeInlineQueryResult, None, None]:
pixiv = self._pixiv
for media in pixiv.images:
logger.info(str(media))
if pixiv.type in ("illust", "manga"):
yield InlineQueryResultPhoto(
id=str(uuid4()),
photo_url=media.large,
thumbnail_url=media.thumb,
caption=self.message_text
)
else:
yield
def message_media_generator(self) -> Generator[TypeMessageMediaResult, None, None]:
pixiv = self._pixiv
for media in pixiv.images:
logger.info(str(media))
if pixiv.type in ("illust", "manga"):
yield InputMediaPhoto(
media=media.large,
has_spoiler=pixiv.is_nsfw
)
else:
yield
class TelegramBsky:
message_raw_text = message_raw_text_tweet
__slots__ = ('_url', '_bsky', '__dict__')
def __init__(self, url: str):
self._url: str = url
async def __aenter__(self):
async with ProcessBsky(self._url) as bsky:
self._bsky = bsky
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
@property
def url(self) -> str:
return self._bsky.url
@cached_property
def message_text(self) -> str:
bsky = self._bsky
return self.message_raw_text.format(
url=bsky.url,
author_url=bsky.author_url,
author=html.escape(bsky.author),
text=html.escape(bsky.text)
)
def inline_query_result(self) -> tuple[TypeInlineQueryResult, ...]:
return tuple(i for i in self.inline_query_generator() if i)
def message_media_result(self) -> tuple[TypeMessageMediaResult, ...]:
return tuple(i for i in self.message_media_generator() if i)
def inline_query_generator(self) -> Generator[TypeInlineQueryResult, None, None]:
bsky = self._bsky
for bsky_media in bsky.media:
logger.info(str(bsky_media))
if bsky_media.type == "image":
yield InlineQueryResultPhoto(
id=str(uuid4()),
photo_url=bsky_media.url,
thumbnail_url=bsky_media.thumb,
caption=self.message_text
)
elif bsky_media.type == "video":
# yield InlineQueryResultVideo(
# id=str(uuid4()),
# video_url=bsky_media.url,
# mime_type="video/mp4",
# thumbnail_url=bsky_media.thumb,
# title=bsky.text,
# caption=self.message_text
# )
yield
elif bsky_media.type == "external":
# yield InlineQueryResultVideo(
# id=str(uuid4()),
# video_url=bsky_media.url,
# mime_type="image/gif",
# thumbnail_url=bsky_media.thumb,
# title=bsky.text,
# caption=self.message_text
# )
yield
def message_media_generator(self) -> Generator[TypeMessageMediaResult, None, None]:
bsky = self._bsky
for bsky_media in bsky.media:
logger.info(str(bsky_media))
if bsky_media.type == "image":
yield InputMediaPhoto(
media=bsky_media.url,
has_spoiler=bsky.sensitive
)
elif bsky_media.type == "video":
# yield InputMediaVideo(
# media=bsky_media.url,
# has_spoiler=bsky.sensitive,
# thumbnail=bsky_media.thumb
# )
yield
elif bsky_media.type == "external":
yield bsky_media.url, bsky.sensitive
-152
View File
@@ -1,152 +0,0 @@
from __future__ import annotations
from functools import cached_property
from typing import TYPE_CHECKING
from .net import NetClient
from .regex import x_media_url, x_tco_url, x_url
if TYPE_CHECKING:
from .types import TweetInfo
twimg_url = 'https://pbs.twimg.com/'
vx_api_url = 'https://api.vxtwitter.com/{0}/status/{1}'
class TweetMedia:
__slots__ = ('_url', '_thumb', '_type', '__dict__')
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"TweetMedia(url={self.url}, thumb={self.thumb}, type={self.type})"
@cached_property
def _uri(self) -> str | None:
if match := x_media_url.match(self._url):
return match.group(2).removesuffix('.jpg').removesuffix('.png')
return None
@cached_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
@cached_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:
__slots__ = ('_id', '_author', '_author_id', '_text', '_media', '_sensitive', '__dict__')
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
@cached_property
def url(self) -> str:
return f"https://x.com/{self._author_id}/status/{self._id}"
@property
def author(self) -> str:
return self._author
@cached_property
def author_url(self) -> str:
return f"https://x.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 ProcessTweet:
__slots__ = ('_url', '_tweet')
def __init__(self, url: str):
self._url: str = url
async def __aenter__(self):
self._tweet = await self._fetch_tweet()
return Tweet(
tweet_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"]
)
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
async def _fetch_tweet(self) -> TweetInfo:
match = x_url.match(self._url)
assert match, f"Invalid URL: {self._url}"
auther_id, tweet_id = match.groups()
return await NetClient.fetch_json(vx_api_url.format(auther_id, tweet_id))
@property
def _tweet_text(self) -> str:
match = x_tco_url.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=tweet_media['url'],
thumb=tweet_media['thumbnail_url'],
media_type=tweet_media['type']
)
for tweet_media in self._tweet['media_extended']
]
-75
View File
@@ -1,75 +0,0 @@
from __future__ import annotations
from typing import Literal, TypedDict
from telegram import InlineQueryResultMpeg4Gif, InlineQueryResultPhoto, InlineQueryResultVideo, InputMediaPhoto, \
InputMediaVideo
TypeInlineQueryResult = InlineQueryResultMpeg4Gif | InlineQueryResultPhoto | InlineQueryResultVideo
InputMediaAnimation = tuple[str, bool]
TypeMessageMediaResult = InputMediaPhoto | InputMediaVideo | InputMediaAnimation
class TweetInfo(TypedDict):
tweetID: str
user_name: str
user_screen_name: str
text: str
media_extended: list[dict]
possibly_sensitive: bool
class BskyInfo(TypedDict):
thread: BskyThread
class BskyThread(TypedDict):
post: BskyPost
class BskyPost(TypedDict):
author: BskyAuthor
record: BskyPostRecord
embed: BskyEmbedImages | BskyEmbedVideo | BskyEmbedExternal
labels: list[BskyLabel]
class BskyAuthor(TypedDict):
handle: str
displayName: str
class BskyPostRecord(TypedDict):
text: str
class BskyLabel(TypedDict):
val: str
class BskyEmbedImage(TypedDict):
thumb: str
fullsize: str
BskyEmbedImages = TypedDict('BskyEmbedImages', {
'$type': Literal['app.bsky.embed.images#view'],
'images': list[BskyEmbedImage]
})
BskyEmbedVideo = TypedDict('BskyEmbedVideo', {
'$type': Literal['app.bsky.embed.video#view'],
'playlist': str,
'thumbnail': str
})
class BskyEmbedExternalItem(TypedDict):
uri: str
thumb: str
BskyEmbedExternal = TypedDict('BskyEmbedExternal', {
'$type': Literal['app.bsky.embed.external#view'],
'external': BskyEmbedExternalItem
})