main
  1#!/usr/bin/env python
  2# -*- coding: utf-8 -*-
  3import asyncio
  4
  5from loguru import logger
  6from pyrogram.client import Client
  7from pyrogram.types import Message
  8
  9from ai.chat_summary import ai_chat_summary
 10from ai.main import ai_image_generation, ai_text_generation, ai_video_generation
 11from asr.voice_recognition import voice_to_text
 12from bridge.ocr import send_to_ocr_bridge
 13from config import FAVORITE, PREFIX, PROXY
 14from danmu.entrypoint import query_danmu
 15from database.database import del_db
 16from history.query import query_chat_history
 17from messages.details import show_msg_info
 18from messages.help import social_media_help
 19from messages.modify import parse_kwargs
 20from messages.parser import parse_msg
 21from messages.sender import send2tg
 22from messages.utils import delete_message, equal_prefix, startswith_prefix
 23from networking import match_social_media_link
 24from others.convert_chinese import chinese_conversion
 25from others.convert_img_file import convert_image_file
 26from others.download_external import download_url_in_message
 27from others.extract_audio import extract_audio_file
 28from others.favorite import save_favorite, send_favorite
 29from others.ffmpeg import ffmpeg_cut, ffmpeg_h264, ffprobe
 30from others.search_google import search_google
 31from others.search_ytb import search_youtube
 32from others.tmdb import search_tmdb
 33from others.version import get_bot_version
 34from others.watermark import add_watermark
 35from preview.arxiv import preview_arxiv
 36from preview.bilibili import preview_bilibili
 37from preview.douyin import preview_douyin
 38from preview.github import preview_github
 39from preview.instagram import preview_instagram
 40from preview.netease import preview_music163
 41from preview.reddit import preview_reddit
 42from preview.spotify import preview_spotify
 43from preview.twitter import preview_twitter
 44from preview.v2ex import preview_v2ex
 45from preview.wechat import preview_wechat
 46from preview.weibo import preview_weibo
 47from preview.xiaohongshu import preview_xhs
 48from price.entrypoint import get_asset_price
 49from quotly.quotly import quote_message
 50from subtitles.subtitle import get_subtitle
 51from tts.tts import text_to_speech
 52from utils import to_int, true
 53from ytdlp.main import preview_ytdlp
 54from ytdlp.utils import ProxyError
 55
 56
 57async def process_message(
 58    client: Client,
 59    message: Message,
 60    target_chat: int | str | None = None,
 61    reply_msg_id: int = 0,
 62    *,
 63    ai: bool = True,  # GPT Chat
 64    asr: bool = True,  # Voice to Text
 65    audio_extract: bool = True,  # Extract Audio from message. /audio -> m4a file. /voice -> wav file
 66    danmu: bool = True,  # Query Danmu database.
 67    favorite: bool = True,  # Send & Save message to favorite.
 68    google_search: bool = True,  # Search on Google.
 69    ocr: bool = True,  # OCR image file.
 70    history: bool = True,  # Query Chat History.
 71    price: bool = True,  # Query Stock / Fiat / Crypto price.
 72    subtitle: bool = True,  # Download YouTube subtitle.
 73    summary: bool = True,  # AI Chat Summary.
 74    tts: bool = True,  # Text to Speech.
 75    quotly: bool = True,  # Quote message.
 76    convert_chinese: bool = True,  # Traditional Chinese <-> Simplified Chinese.
 77    wget: bool = True,  # Download external URL.
 78    ytb: bool = True,  # Search on YouTube.
 79    tmdb: bool = True,  # Search on TMDB.
 80    ffmpeg: bool = True,  # FFMpeg processing
 81    watermark: bool = True,  # Add watermark to video.
 82    version: bool = True,  # Show bot version with dependencies.
 83    convert_img: bool = True,  # Convert image file.
 84    show_progress: bool = True,  # Show progress bar.
 85    detail_progress: bool = False,  # Show detailed progress bar.
 86    **kwargs,
 87):
 88    """Call utility functions to handle the message.
 89
 90    Args:
 91        client (Client): The Pyrogram client.
 92        message (Message): The trigger message object.
 93        target_chat (int | str, optional): Send result to this telegram target chat. If not set, send to the trigger message's chat.
 94        reply_msg_id (int, optional): If set to integer > 0, the result is sent as a reply message to this message_id.
 95                                             If set to 0, reply to the trigger message itself.
 96                                             If set to -1, do not send as a reply message.
 97    """
 98    kwargs |= {"target_chat": target_chat, "reply_msg_id": reply_msg_id, "show_progress": show_progress, "detail_progress": detail_progress}
 99    if true(kwargs.get("target_chat")):
100        kwargs["target_chat"] = to_int(kwargs["target_chat"])
101    message, msg_kwargs = parse_kwargs(message)
102    kwargs |= msg_kwargs  # merge the kwargs from the message text
103    if ai:
104        await ai_text_generation(client, message, **kwargs)  # /ai
105        await ai_image_generation(client, message, **kwargs)  # /gen
106        await ai_video_generation(client, message, **kwargs)  # /gvid
107    if asr:
108        await voice_to_text(client, message, **kwargs)  # /asr
109    if audio_extract:
110        await extract_audio_file(client, message, **kwargs)  # /audio
111    if subtitle:
112        await get_subtitle(client, message, **kwargs)  # /subtitle
113    if wget:
114        await download_url_in_message(client, message, **kwargs)  # /wget
115    if google_search:
116        await search_google(client, message, **kwargs)  # /google
117    if ytb:
118        await search_youtube(client, message, **kwargs)  # /ytb
119    if ocr:
120        await send_to_ocr_bridge(client, message, **kwargs)  # /ocr
121    if price:
122        await get_asset_price(client, message, **kwargs)  # /price
123    if history:
124        await query_chat_history(client, message, **kwargs)  # /history
125    if summary:
126        await ai_chat_summary(client, message, **kwargs)  # /summary
127    if danmu:
128        await query_danmu(client, message, **kwargs)  # /danmu
129    if favorite:
130        await save_favorite(client, message, **kwargs)  # /save
131        await send_favorite(client, message, **kwargs)  # /fav
132    if tts:
133        await text_to_speech(client, message, **kwargs)  # /tts
134    if convert_chinese:
135        await chinese_conversion(client, message, **kwargs)  # /sc
136    if quotly:
137        await quote_message(client, message, **kwargs)  # /quote
138    if tmdb:
139        await search_tmdb(client, message, **kwargs)  # /tmdb
140    if convert_img:
141        await convert_image_file(client, message, **kwargs)
142    if ffmpeg:
143        await ffmpeg_cut(client, message, **kwargs)
144        await ffmpeg_h264(client, message, **kwargs)
145        await ffprobe(client, message, **kwargs)
146    if watermark:
147        await add_watermark(client, message, **kwargs)
148    if version:
149        await get_bot_version(client, message, **kwargs)
150
151    await show_msg_info(client, message, **kwargs)
152    await preview_social_media(client, message, **kwargs)
153
154
155async def preview_social_media(
156    client: Client,
157    message: Message,
158    *,
159    need_prefix: bool = True,  # Need prefix to parse social media. (/dl)
160    prepend_sender_user: bool = False,
161    douyin: bool = True,  # Parse Douyin
162    tiktok: bool = True,  # Parse TikTok
163    instagram: bool = True,  # Parse Instagram
164    twitter: bool = True,  # Parse Twitter
165    weibo: bool = True,  # Parse Weibo
166    reddit: bool = True,  # Parse Reddit
167    github: bool = True,  # Parse Github
168    xhs: bool = True,  # Parse XiaoHongShu
169    v2ex: bool = True,  # Parse V2EX
170    music163: bool = True,  # Parse Music163
171    spotify: bool = True,  # Parse Spotify
172    ytdlp: bool = True,  # Parse YT-DLP link
173    ytdlp_bilibili: bool = True,  # Parse YT-DLP Bilibili link
174    ytdlp_youtube: bool = True,  # Parse YT-DLP YouTube link
175    arxiv: bool = True,  # Parse arXiv
176    **kwargs,
177):
178    """Preview social media link in the message.
179
180    Args:
181        client (Client): The Pyrogram client.
182        message (Message): The trigger message object.
183        target_chat (int | str, optional): Send result to this telegram target chat. If not set, send to the trigger message's chat.
184        reply_msg_id (int, optional): If set to integer > 0, the result is sent as a reply message to this message_id.
185                                             If set to 0, reply to the trigger message itself.
186                                             If set to -1, do not send as a reply message.
187        need_prefix (bool, optional): Need to start with PREFIX to call this funciton. Defaults to True.
188        cmd_prefix (str, optional): prefix to call this function.
189        prepend_sender_user (bool, optional): Prepend the sender's username to the message. Defaults to False.
190        show_progress (bool, optional): Show a progress message on Telegram. Defaults to True.
191        detail_progress (bool, optional): Show detailed progress (Only if show_proress is set to True). Defaults to False.
192    """
193    # these commands are handled in `process_message`
194    ignore_prefix = [
195        PREFIX.ASR,
196        PREFIX.AI_SUMMARY,
197        PREFIX.AI_TEXT_GENERATION,
198        PREFIX.AI_IMG_GENERATION,
199        PREFIX.AUDIO,
200        PREFIX.COMBINATION,
201        PREFIX.CONVERT,
202        PREFIX.CRYPTO,
203        PREFIX.DANMU,
204        PREFIX.OCR,
205        PREFIX.PRICE,
206        PREFIX.SEARCH_GOOGLE,
207        PREFIX.SEARCH_YOUTUBE,
208        PREFIX.STOCK,
209        PREFIX.SUBTITLE,
210        PREFIX.VOICE,
211        PREFIX.WGET,
212        PREFIX.FAYAN,
213        PREFIX.TTS,
214        PREFIX.TMDB,
215        PREFIX.CONVERT_TO_SC,
216        PREFIX.CONVERT_TO_TC,
217        FAVORITE.SAVE_PREFIX,
218        FAVORITE.SEND_PREFIX,
219    ]
220
221    info = parse_msg(message)
222    this_msg = message
223    this_texts = info["text"]  # texts of the trigger message
224    if startswith_prefix(this_texts, prefix=ignore_prefix):
225        return None
226    if need_prefix and not startswith_prefix(this_texts, prefix=[PREFIX.SOCIAL_MEDIA, "/help", "/retry"]):
227        return None
228
229    # message only contains prefix command
230    if equal_prefix(this_texts, prefix=[PREFIX.SOCIAL_MEDIA, "/help", "/retry"]):
231        # without reply, send docs if message only contains prefix command
232        if not message.reply_to_message:
233            help_msg = social_media_help(info["cid"], info["ctype"], PREFIX.SOCIAL_MEDIA)
234            return await send2tg(client, message, texts=help_msg, **kwargs)
235        # with reply, treat the reply_msg as the trigger to preview social media link
236        message = message.reply_to_message
237        info = parse_msg(message, silent=True, use_cache=False)  # parse again
238
239    # add send_from_user.
240    if prepend_sender_user:
241        # Caution: this format should be consistent with `save_messages` function in `message.database.py`
242        kwargs["send_from_user"] = f"👤[@{info['full_name']}](tg://user?id={info['uid']})//"
243
244    warn_msg = None
245    try:
246        matched = await match_social_media_link(info["text"])  # match "platform" and "url" (and other info)
247        if matched["platform"]:
248            logger.success(f"Matched: {matched}")
249        kwargs |= matched
250        if startswith_prefix(this_texts, prefix="/retry"):
251            await del_db(matched["db_key"])
252
253        if douyin and matched["platform"] == "douyin":
254            return await preview_douyin(client, message, **kwargs)
255        if tiktok and matched["platform"] == "tiktok":
256            return await preview_douyin(client, message, **kwargs)
257        if instagram and matched["platform"] == "instagram":
258            return await preview_instagram(client, message, **kwargs)
259        if twitter and matched["platform"] in ["x", "twitter", "fxtwitter", "fixupx"]:
260            return await preview_twitter(client, message, **kwargs)
261        if weibo and matched["platform"] == "weibo":
262            return await preview_weibo(client, message, **kwargs)
263        if xhs and matched["platform"] == "xiaohongshu":
264            return await preview_xhs(client, message, **kwargs)
265        if xhs and matched["platform"] == "wechat":
266            return await preview_wechat(client, message, **kwargs)
267        if github and matched["platform"] == "github":
268            return await preview_github(client, message, **kwargs)
269        if reddit and matched["platform"] == "reddit":
270            return await preview_reddit(client, message, **kwargs)
271        if music163 and matched["platform"] == "music163":
272            return await preview_music163(client, message, **kwargs)
273        if spotify and matched["platform"] == "spotify":
274            return await preview_spotify(client, message, **kwargs)
275        if v2ex and matched["platform"] == "v2ex":
276            return await preview_v2ex(client, message, **kwargs)
277        if arxiv and matched["platform"] == "arxiv":
278            return await preview_arxiv(client, message, **kwargs)
279        if matched["platform"].startswith("bilibili-"):  # this is not bilibili video, for videos, use yt-dlp
280            return await preview_bilibili(client, message, **kwargs)
281
282        sent_messages = []
283        try:
284            if ytdlp_bilibili and matched["platform"] == "bilibili":
285                sent_messages = await preview_ytdlp(client, message, **kwargs)
286            if ytdlp_youtube and matched["platform"] == "youtube":
287                sent_messages = await preview_ytdlp(client, message, **kwargs)
288            if ytdlp and matched["platform"] == "ytdlp":
289                sent_messages = await preview_ytdlp(client, message, **kwargs)
290        except ProxyError:
291            logger.error(f"🚫{matched['platform']}代理错误")
292            if PROXY.YTDLP_FALLBACK:
293                logger.warning(f"🔄使用备用代理{PROXY.YTDLP_FALLBACK}")
294                sent_messages = await preview_ytdlp(client, message, proxy=PROXY.YTDLP_FALLBACK, **kwargs)
295
296        if not sent_messages and any(keyword in info["text"] for keyword in ["facebook.com", "threads.com", "youtube.com", "bilibili.com"]):
297            if kwargs.get("show_progress"):
298                warn_msg = await message.reply_text("⚠️该链接不被支持", quote=True)
299                await asyncio.sleep(10)
300                await delete_message(warn_msg)
301        # if ytdlp failed, try to download directly
302        elif (
303            not sent_messages
304            and startswith_prefix(this_texts, prefix=PREFIX.SOCIAL_MEDIA)
305            and matched["platform"]
306            not in [
307                "bilibili",
308                "douyin",
309                "github",
310                "instagram",
311                "music163",
312                "reddit",
313                "spotify",
314                "tiktok",
315                "v2ex",
316                "weibo",
317                "x",
318                "xiaohongshu",
319                "youtube",
320            ]
321        ):
322            if kwargs.get("show_progress"):
323                warn_msg = await client.send_message(info["cid"], text="⚠️暂时不支持解析链接, 尝试直接下载该网页")
324            await download_url_in_message(client, this_msg, extra_prefix=PREFIX.SOCIAL_MEDIA, **kwargs)
325
326    except Exception as e:
327        logger.exception(e)
328    finally:
329        await delete_message(warn_msg)
330    return sent_messages