1#!/usr/bin/env python
 2# -*- coding: utf-8 -*-
 3from pathlib import Path
 4
 5from loguru import logger
 6from pyrogram.client import Client
 7from pyrogram.types import Message
 8
 9from config import MAX_FILE_BYTES, PREFIX, PROXY
10from messages.parser import parse_msg
11from messages.progress import modify_progress
12from messages.sender import send2tg
13from messages.utils import equal_prefix, get_reply_to, startswith_prefix
14from multimedia import is_valid_video_or_audio, validate_img
15from networking import download_file
16from publish import publish_telegraph
17from utils import convert_md, find_url, guess_mime, readable_size, to_int
18
19HELP = f"""
20⏬**下载文件**
21使用说明:
22`{PREFIX.WGET}` + URL 下载文件
23文件大小不能超过Telegram的限制: {round(MAX_FILE_BYTES / 1024 / 1024)}MB
24"""
25AUDIO_FORMAT = [".aac", ".aiff", ".flac", ".m4a", ".mp3", ".oga", ".ogg", ".opus", ".pcm", ".wav"]
26VIDEO_FORMAT = [".3gp", ".avi", ".flv", ".m2ts", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".mts", ".ogg", ".ogv", ".rmvb", ".webm", ".wmv"]
27
28
29async def download_url_in_message(client: Client, message: Message, extra_prefix: str = "", **kwargs):
30    """Download the url from the message."""
31    info = parse_msg(message)
32    if not startswith_prefix(info["text"], prefix=[PREFIX.WGET, extra_prefix]):
33        return
34    # send docs if message == "/wget", without reply
35    if equal_prefix(message.text, prefix=[PREFIX.WGET]) and not message.reply_to_message:
36        await send2tg(client, message, texts=HELP, **kwargs)
37        return
38    # reply a message with /wget
39    if message.reply_to_message:
40        message = message.reply_to_message
41        info = parse_msg(message, silent=True)  # parse again
42
43    target_chat = kwargs["target_chat"] if kwargs.get("target_chat") else message.chat.id
44    target_chat = to_int(target_chat)
45    reply_msg_id = kwargs.get("reply_msg_id", 0)
46    reply_parameters = get_reply_to(message.id, reply_msg_id)
47
48    if not (url := find_url(info["text"])):
49        await message.reply_text("❌未找到URL")
50        return
51    caption = f"🔗[原始链接]({url})"
52    msg = f"⏬开始下载:\n{url}"
53    if kwargs.get("show_progress") and "progress" not in kwargs:
54        res = await send2tg(client, message, texts=msg, **kwargs)
55        kwargs["progress"] = res[0]
56    else:
57        await modify_progress(text=msg, force_update=True, **kwargs)
58    success = False
59    path = Path("non-exist")
60    try:
61        path = await download_file(url, proxy=PROXY.DOWNLOAD, stream=True, **kwargs)
62        path = Path(path)
63        suffix = path.suffix
64        mime = guess_mime(path)
65        logger.debug(f"File MIME: {mime}")
66        if not (mime.startswith("text") or suffix in AUDIO_FORMAT + VIDEO_FORMAT):
67            suffix = "." + mime.split("/")[-1]
68        if path.suffix != suffix:
69            path.rename(path.with_suffix(suffix))
70            path = path.with_suffix(suffix)
71        if img := validate_img(path, delete=False):
72            await modify_progress(text=f"🏞图片下载成功: {readable_size(path=img)}", force_update=True, **kwargs)
73            success = await send2tg(client, message, target_chat, reply_msg_id, texts=caption, media=[{"photo": img}])
74        elif path.suffix in AUDIO_FORMAT:
75            await modify_progress(text=f"🎧音频下载成功: {readable_size(path=path)}", force_update=True, **kwargs)
76            success = await client.send_audio(target_chat, path.as_posix(), caption=caption, reply_parameters=reply_parameters)
77        elif await is_valid_video_or_audio(path, delete=False):
78            await modify_progress(text=f"🎬视频下载成功: {readable_size(path=path)}", force_update=True, **kwargs)
79            success = await send2tg(client, message, target_chat, reply_msg_id, texts=caption, media=[{"video": path}])
80        elif path.stat().st_size < MAX_FILE_BYTES:
81            await modify_progress(text=f"💾文件下载成功: {readable_size(path=path)}", force_update=True, **kwargs)
82            if suffix == ".html":
83                markdown = convert_md(path)
84                markdown_path = path.with_suffix(".md")
85                markdown_path.write_text(markdown)
86                if telegraph_url := await publish_telegraph(title="全文内容", texts=markdown, author=info["full_name"], url=url):
87                    caption += f"\n⚡️[即时预览]({telegraph_url})"
88            success = await client.send_document(target_chat, path.as_posix(), caption=caption, reply_parameters=reply_parameters)
89        else:
90            await modify_progress(text=f"❌文件大小: {readable_size(path=path)} 超出限制\nTelegram只允许上传小于{round(MAX_FILE_BYTES / 1024 / 1024)}MB的文件", force_update=True, **kwargs)
91    except Exception as e:
92        logger.error(e)
93        await modify_progress(text=f"❌下载失败: {url}", force_update=True, **kwargs)
94
95    if success:
96        await modify_progress(del_status=True, **kwargs)
97        Path(path).unlink(missing_ok=True)