update 26.7.25
This commit is contained in:
Executable
+687
@@ -0,0 +1,687 @@
|
||||
"""Downloads media from telegram."""
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import pyrogram
|
||||
from loguru import logger
|
||||
from pyrogram.types import Audio, Document, Photo, Video, VideoNote, Voice
|
||||
from rich.logging import RichHandler
|
||||
|
||||
from module.app import Application, ChatDownloadConfig, DownloadStatus, TaskNode
|
||||
from module.bot import start_download_bot, stop_download_bot
|
||||
from module.download_record import (
|
||||
build_download_record,
|
||||
get_database_path,
|
||||
upsert_download_record,
|
||||
)
|
||||
from module.download_stat import update_download_status
|
||||
from module.get_chat_history_v2 import get_chat_history_v2
|
||||
from module.language import _t
|
||||
from module.nfo import write_video_nfo
|
||||
from module.pyrogram_extension import (
|
||||
HookClient,
|
||||
fetch_message,
|
||||
get_extension,
|
||||
record_download_status,
|
||||
report_bot_download_status,
|
||||
set_max_concurrent_transmissions,
|
||||
set_meta_data,
|
||||
update_cloud_upload_stat,
|
||||
upload_telegram_chat,
|
||||
)
|
||||
from module.storage import build_download_paths
|
||||
from module.web import init_web
|
||||
from utils.format import truncate_filename, validate_title
|
||||
from utils.log import LogFilter
|
||||
from utils.meta import print_meta
|
||||
from utils.meta_data import MetaData
|
||||
from utils.updates import check_for_updates
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(message)s",
|
||||
datefmt="[%X]",
|
||||
handlers=[RichHandler()],
|
||||
)
|
||||
|
||||
CONFIG_NAME = "config.yaml"
|
||||
DATA_FILE_NAME = "data.yaml"
|
||||
APPLICATION_NAME = "media_downloader"
|
||||
app = Application(CONFIG_NAME, DATA_FILE_NAME, APPLICATION_NAME)
|
||||
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
RETRY_TIME_OUT = 3
|
||||
|
||||
logging.getLogger("pyrogram.session.session").addFilter(LogFilter())
|
||||
logging.getLogger("pyrogram.client").addFilter(LogFilter())
|
||||
|
||||
logging.getLogger("pyrogram").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
def _check_download_finish(media_size: int, download_path: str, ui_file_name: str):
|
||||
"""Check download task if finish
|
||||
|
||||
Parameters
|
||||
----------
|
||||
media_size: int
|
||||
The size of the downloaded resource
|
||||
download_path: str
|
||||
Resource download hold path
|
||||
ui_file_name: str
|
||||
Really show file name
|
||||
|
||||
"""
|
||||
download_size = os.path.getsize(download_path)
|
||||
if media_size == download_size:
|
||||
logger.success(f"{_t('Successfully downloaded')} - {ui_file_name}")
|
||||
else:
|
||||
logger.warning(
|
||||
f"{_t('Media downloaded with wrong size')}: "
|
||||
f"{download_size}, {_t('actual')}: "
|
||||
f"{media_size}, {_t('file name')}: {ui_file_name}"
|
||||
)
|
||||
os.remove(download_path)
|
||||
raise pyrogram.errors.exceptions.bad_request_400.BadRequest()
|
||||
|
||||
|
||||
def _move_to_download_path(temp_download_path: str, download_path: str):
|
||||
"""Move file to download path
|
||||
|
||||
Parameters
|
||||
----------
|
||||
temp_download_path: str
|
||||
Temporary download path
|
||||
|
||||
download_path: str
|
||||
Download path
|
||||
|
||||
"""
|
||||
|
||||
directory, _ = os.path.split(download_path)
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
shutil.move(temp_download_path, download_path)
|
||||
|
||||
|
||||
def _check_timeout(retry: int, _: int):
|
||||
"""Check if message download timeout, then add message id into failed_ids
|
||||
|
||||
Parameters
|
||||
----------
|
||||
retry: int
|
||||
Retry download message times
|
||||
|
||||
message_id: int
|
||||
Try to download message 's id
|
||||
|
||||
"""
|
||||
if retry == 2:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _can_download(_type: str, file_formats: dict, file_format: Optional[str]) -> bool:
|
||||
"""
|
||||
Check if the given file format can be downloaded.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
_type: str
|
||||
Type of media object.
|
||||
file_formats: dict
|
||||
Dictionary containing the list of file_formats
|
||||
to be downloaded for `audio`, `document` & `video`
|
||||
media types
|
||||
file_format: str
|
||||
Format of the current file to be downloaded.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if the file format can be downloaded else False.
|
||||
"""
|
||||
if _type in ["audio", "document", "video"]:
|
||||
allowed_formats: list = file_formats[_type]
|
||||
if not file_format in allowed_formats and allowed_formats[0] != "all":
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _is_exist(file_path: str) -> bool:
|
||||
"""
|
||||
Check if a file exists and it is not a directory.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
file_path: str
|
||||
Absolute path of the file to be checked.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if the file exists else False.
|
||||
"""
|
||||
return not os.path.isdir(file_path) and os.path.exists(file_path)
|
||||
|
||||
|
||||
# pylint: disable = R0912
|
||||
|
||||
|
||||
async def _get_media_meta(
|
||||
chat_id: Union[int, str],
|
||||
message: pyrogram.types.Message,
|
||||
media_obj: Union[Audio, Document, Photo, Video, VideoNote, Voice],
|
||||
_type: str,
|
||||
) -> Tuple[str, str, Optional[str]]:
|
||||
"""Extract file name and file id from media object.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
media_obj: Union[Audio, Document, Photo, Video, VideoNote, Voice]
|
||||
Media object to be extracted.
|
||||
_type: str
|
||||
Type of media object.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[str, str, Optional[str]]
|
||||
file_name, file_format
|
||||
"""
|
||||
caption = getattr(message, "caption", None)
|
||||
if caption:
|
||||
caption = validate_title(caption)
|
||||
app.set_caption_name(chat_id, message.media_group_id, caption)
|
||||
app.set_caption_entities(
|
||||
chat_id, message.media_group_id, message.caption_entities
|
||||
)
|
||||
else:
|
||||
app.get_caption_name(chat_id, message.media_group_id)
|
||||
|
||||
paths = build_download_paths(app, chat_id, message, media_obj, _type)
|
||||
return paths.final_path, paths.temp_path, paths.file_format
|
||||
|
||||
|
||||
async def add_download_task(
|
||||
message: pyrogram.types.Message,
|
||||
node: TaskNode,
|
||||
):
|
||||
"""Add Download task"""
|
||||
if message.empty:
|
||||
return False
|
||||
node.download_status[message.id] = DownloadStatus.Downloading
|
||||
await queue.put((message, node))
|
||||
node.total_task += 1
|
||||
return True
|
||||
|
||||
|
||||
async def save_msg_to_file(
|
||||
app, chat_id: Union[int, str], message: pyrogram.types.Message
|
||||
):
|
||||
"""Write message text into file"""
|
||||
dirname = validate_title(
|
||||
message.chat.title if message.chat and message.chat.title else str(chat_id)
|
||||
)
|
||||
datetime_dir_name = message.date.strftime(app.date_format) if message.date else "0"
|
||||
|
||||
file_save_path = app.get_file_save_path("msg", dirname, datetime_dir_name)
|
||||
stored_file_name = validate_title(f"{chat_id}_{message.id}.txt")
|
||||
file_name = os.path.join(file_save_path, stored_file_name)
|
||||
|
||||
os.makedirs(os.path.dirname(file_name), exist_ok=True)
|
||||
|
||||
if _is_exist(file_name):
|
||||
return DownloadStatus.SkipDownload, None
|
||||
|
||||
with open(file_name, "w", encoding="utf-8") as f:
|
||||
f.write(message.text or "")
|
||||
|
||||
return DownloadStatus.SuccessDownload, file_name
|
||||
|
||||
|
||||
def _record_download_result(
|
||||
message: pyrogram.types.Message,
|
||||
node: TaskNode,
|
||||
download_status: DownloadStatus,
|
||||
file_name: str = None,
|
||||
error_message: str = None,
|
||||
):
|
||||
"""Persist final download status and generate video NFOs when needed."""
|
||||
try:
|
||||
record = build_download_record(
|
||||
app,
|
||||
node,
|
||||
message,
|
||||
download_status,
|
||||
local_path=file_name,
|
||||
error_message=error_message,
|
||||
)
|
||||
|
||||
if (
|
||||
download_status is DownloadStatus.SuccessDownload
|
||||
and record.media_type == "video"
|
||||
and file_name
|
||||
):
|
||||
try:
|
||||
record.nfo_path = write_video_nfo(record)
|
||||
except Exception as nfo_error:
|
||||
logger.warning(f"generate nfo error: {nfo_error}")
|
||||
if not record.error_message:
|
||||
record.error_message = str(nfo_error)
|
||||
|
||||
upsert_download_record(get_database_path(app.save_path), record)
|
||||
except Exception as record_error:
|
||||
logger.warning(f"record download result error: {record_error}")
|
||||
|
||||
|
||||
async def download_task(
|
||||
client: pyrogram.Client, message: pyrogram.types.Message, node: TaskNode
|
||||
):
|
||||
"""Download and Forward media"""
|
||||
|
||||
download_status, file_name = await download_media(
|
||||
client, message, app.media_types, app.file_formats, node
|
||||
)
|
||||
|
||||
if app.enable_download_txt and message.text and not message.media:
|
||||
download_status, file_name = await save_msg_to_file(app, node.chat_id, message)
|
||||
|
||||
if not node.bot:
|
||||
app.set_download_id(node, message.id, download_status)
|
||||
|
||||
node.download_status[message.id] = download_status
|
||||
|
||||
file_size = os.path.getsize(file_name) if file_name else 0
|
||||
_record_download_result(message, node, download_status, file_name)
|
||||
|
||||
await upload_telegram_chat(
|
||||
client,
|
||||
node.upload_user if node.upload_user else client,
|
||||
app,
|
||||
node,
|
||||
message,
|
||||
download_status,
|
||||
file_name,
|
||||
)
|
||||
|
||||
# rclone upload
|
||||
if (
|
||||
not node.upload_telegram_chat_id
|
||||
and download_status is DownloadStatus.SuccessDownload
|
||||
):
|
||||
ui_file_name = file_name
|
||||
if app.hide_file_name:
|
||||
ui_file_name = f"****{os.path.splitext(file_name)[-1]}"
|
||||
if await app.upload_file(
|
||||
file_name, update_cloud_upload_stat, (node, message.id, ui_file_name)
|
||||
):
|
||||
node.upload_success_count += 1
|
||||
|
||||
await report_bot_download_status(
|
||||
node.bot,
|
||||
node,
|
||||
download_status,
|
||||
file_size,
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable = R0915,R0914
|
||||
|
||||
|
||||
@record_download_status
|
||||
async def download_media(
|
||||
client: pyrogram.client.Client,
|
||||
message: pyrogram.types.Message,
|
||||
media_types: List[str],
|
||||
file_formats: dict,
|
||||
node: TaskNode,
|
||||
):
|
||||
"""
|
||||
Download media from Telegram.
|
||||
|
||||
Each of the files to download are retried 3 times with a
|
||||
delay of 5 seconds each.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
client: pyrogram.client.Client
|
||||
Client to interact with Telegram APIs.
|
||||
message: pyrogram.types.Message
|
||||
Message object retrieved from telegram.
|
||||
media_types: list
|
||||
List of strings of media types to be downloaded.
|
||||
Ex : `["audio", "photo"]`
|
||||
Supported formats:
|
||||
* audio
|
||||
* document
|
||||
* photo
|
||||
* video
|
||||
* voice
|
||||
file_formats: dict
|
||||
Dictionary containing the list of file_formats
|
||||
to be downloaded for `audio`, `document` & `video`
|
||||
media types.
|
||||
|
||||
Returns
|
||||
-------
|
||||
int
|
||||
Current message id.
|
||||
"""
|
||||
|
||||
# pylint: disable = R0912
|
||||
|
||||
file_name: str = ""
|
||||
ui_file_name: str = ""
|
||||
task_start_time: float = time.time()
|
||||
media_size = 0
|
||||
_media = None
|
||||
message = await fetch_message(client, message)
|
||||
try:
|
||||
for _type in media_types:
|
||||
_media = getattr(message, _type, None)
|
||||
if _media is None:
|
||||
continue
|
||||
file_name, temp_file_name, file_format = await _get_media_meta(
|
||||
node.chat_id, message, _media, _type
|
||||
)
|
||||
media_size = getattr(_media, "file_size", 0)
|
||||
|
||||
ui_file_name = file_name
|
||||
if app.hide_file_name:
|
||||
ui_file_name = f"****{os.path.splitext(file_name)[-1]}"
|
||||
|
||||
if _can_download(_type, file_formats, file_format):
|
||||
if _is_exist(file_name):
|
||||
file_size = os.path.getsize(file_name)
|
||||
if file_size or file_size == media_size:
|
||||
logger.info(
|
||||
f"id={message.id} {ui_file_name} "
|
||||
f"{_t('already download,download skipped')}.\n"
|
||||
)
|
||||
|
||||
return DownloadStatus.SkipDownload, file_name
|
||||
else:
|
||||
return DownloadStatus.SkipDownload, None
|
||||
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Message[{message.id}]: "
|
||||
f"{_t('could not be downloaded due to following exception')}:\n[{e}].",
|
||||
exc_info=True,
|
||||
)
|
||||
return DownloadStatus.FailedDownload, None
|
||||
if _media is None:
|
||||
return DownloadStatus.SkipDownload, None
|
||||
|
||||
message_id = message.id
|
||||
|
||||
for retry in range(3):
|
||||
try:
|
||||
temp_download_path = await client.download_media(
|
||||
message,
|
||||
file_name=temp_file_name,
|
||||
progress=update_download_status,
|
||||
progress_args=(
|
||||
message_id,
|
||||
ui_file_name,
|
||||
task_start_time,
|
||||
node,
|
||||
client,
|
||||
),
|
||||
)
|
||||
|
||||
if temp_download_path and isinstance(temp_download_path, str):
|
||||
_check_download_finish(media_size, temp_download_path, ui_file_name)
|
||||
await asyncio.sleep(0.5)
|
||||
_move_to_download_path(temp_download_path, file_name)
|
||||
# TODO: if not exist file size or media
|
||||
return DownloadStatus.SuccessDownload, file_name
|
||||
except pyrogram.errors.exceptions.bad_request_400.BadRequest:
|
||||
logger.warning(
|
||||
f"Message[{message.id}]: {_t('file reference expired, refetching')}..."
|
||||
)
|
||||
await asyncio.sleep(RETRY_TIME_OUT)
|
||||
message = await fetch_message(client, message)
|
||||
if _check_timeout(retry, message.id):
|
||||
# pylint: disable = C0301
|
||||
logger.error(
|
||||
f"Message[{message.id}]: "
|
||||
f"{_t('file reference expired for 3 retries, download skipped.')}"
|
||||
)
|
||||
except pyrogram.errors.exceptions.flood_420.FloodWait as wait_err:
|
||||
await asyncio.sleep(wait_err.value)
|
||||
logger.warning("Message[{}]: FlowWait {}", message.id, wait_err.value)
|
||||
_check_timeout(retry, message.id)
|
||||
except TypeError:
|
||||
# pylint: disable = C0301
|
||||
logger.warning(
|
||||
f"{_t('Timeout Error occurred when downloading Message')}[{message.id}], "
|
||||
f"{_t('retrying after')} {RETRY_TIME_OUT} {_t('seconds')}"
|
||||
)
|
||||
await asyncio.sleep(RETRY_TIME_OUT)
|
||||
if _check_timeout(retry, message.id):
|
||||
logger.error(
|
||||
f"Message[{message.id}]: {_t('Timing out after 3 reties, download skipped.')}"
|
||||
)
|
||||
except Exception as e:
|
||||
# pylint: disable = C0301
|
||||
logger.error(
|
||||
f"Message[{message.id}]: "
|
||||
f"{_t('could not be downloaded due to following exception')}:\n[{e}].",
|
||||
exc_info=True,
|
||||
)
|
||||
break
|
||||
|
||||
return DownloadStatus.FailedDownload, None
|
||||
|
||||
|
||||
def _load_config():
|
||||
"""Load config"""
|
||||
app.load_config()
|
||||
|
||||
|
||||
def _check_config() -> bool:
|
||||
"""Check config"""
|
||||
print_meta(logger)
|
||||
try:
|
||||
_load_config()
|
||||
logger.add(
|
||||
os.path.join(app.log_file_path, "tdl.log"),
|
||||
rotation="10 MB",
|
||||
retention="10 days",
|
||||
level=app.log_level,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception(f"load config error: {e}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
async def worker(client: pyrogram.client.Client):
|
||||
"""Work for download task"""
|
||||
while app.is_running:
|
||||
try:
|
||||
item = await queue.get()
|
||||
message = item[0]
|
||||
node: TaskNode = item[1]
|
||||
|
||||
if node.is_stop_transmission:
|
||||
continue
|
||||
|
||||
if node.client:
|
||||
await download_task(node.client, message, node)
|
||||
else:
|
||||
await download_task(client, message, node)
|
||||
except Exception as e:
|
||||
logger.exception(f"{e}")
|
||||
|
||||
|
||||
async def download_chat_task(
|
||||
client: pyrogram.Client,
|
||||
chat_download_config: ChatDownloadConfig,
|
||||
node: TaskNode,
|
||||
):
|
||||
"""Download all task"""
|
||||
messages_iter = get_chat_history_v2(
|
||||
client,
|
||||
node.chat_id,
|
||||
limit=node.limit,
|
||||
max_id=node.end_offset_id,
|
||||
offset_id=chat_download_config.last_read_message_id,
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
chat_download_config.node = node
|
||||
|
||||
if chat_download_config.ids_to_retry:
|
||||
logger.info(f"{_t('Downloading files failed during last run')}...")
|
||||
skipped_messages: list = await client.get_messages( # type: ignore
|
||||
chat_id=node.chat_id, message_ids=chat_download_config.ids_to_retry
|
||||
)
|
||||
|
||||
for message in skipped_messages:
|
||||
await add_download_task(message, node)
|
||||
|
||||
async for message in messages_iter: # type: ignore
|
||||
meta_data = MetaData()
|
||||
|
||||
caption = message.caption
|
||||
if caption:
|
||||
caption = validate_title(caption)
|
||||
app.set_caption_name(node.chat_id, message.media_group_id, caption)
|
||||
app.set_caption_entities(
|
||||
node.chat_id, message.media_group_id, message.caption_entities
|
||||
)
|
||||
else:
|
||||
caption = app.get_caption_name(node.chat_id, message.media_group_id)
|
||||
set_meta_data(meta_data, message, caption)
|
||||
|
||||
if app.need_skip_message(chat_download_config, message.id):
|
||||
continue
|
||||
|
||||
if app.exec_filter(chat_download_config, meta_data):
|
||||
await add_download_task(message, node)
|
||||
else:
|
||||
node.download_status[message.id] = DownloadStatus.SkipDownload
|
||||
if message.media_group_id:
|
||||
await upload_telegram_chat(
|
||||
client,
|
||||
node.upload_user,
|
||||
app,
|
||||
node,
|
||||
message,
|
||||
DownloadStatus.SkipDownload,
|
||||
)
|
||||
|
||||
chat_download_config.need_check = True
|
||||
chat_download_config.total_task = node.total_task
|
||||
node.is_running = True
|
||||
|
||||
|
||||
async def download_all_chat(client: pyrogram.Client):
|
||||
"""Download All chat"""
|
||||
for key, value in app.chat_download_config.items():
|
||||
value.node = TaskNode(chat_id=key)
|
||||
try:
|
||||
await download_chat_task(client, value, value.node)
|
||||
except Exception as e:
|
||||
logger.warning(f"Download {key} error: {e}")
|
||||
finally:
|
||||
value.need_check = True
|
||||
|
||||
|
||||
async def run_until_all_task_finish():
|
||||
"""Normal download"""
|
||||
while True:
|
||||
finish: bool = True
|
||||
for _, value in app.chat_download_config.items():
|
||||
if not value.need_check or value.total_task != value.finish_task:
|
||||
finish = False
|
||||
|
||||
if (not app.bot_token and finish) or app.restart_program:
|
||||
break
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
def _exec_loop():
|
||||
"""Exec loop"""
|
||||
|
||||
app.loop.run_until_complete(run_until_all_task_finish())
|
||||
|
||||
|
||||
async def start_server(client: pyrogram.Client):
|
||||
"""
|
||||
Start the server using the provided client.
|
||||
"""
|
||||
await client.start()
|
||||
|
||||
|
||||
async def stop_server(client: pyrogram.Client):
|
||||
"""
|
||||
Stop the server using the provided client.
|
||||
"""
|
||||
await client.stop()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function of the downloader."""
|
||||
tasks = []
|
||||
client = HookClient(
|
||||
"media_downloader",
|
||||
api_id=app.api_id,
|
||||
api_hash=app.api_hash,
|
||||
proxy=app.proxy,
|
||||
workdir=app.session_file_path,
|
||||
start_timeout=app.start_timeout,
|
||||
no_updates=True,
|
||||
)
|
||||
try:
|
||||
app.pre_run()
|
||||
init_web(app)
|
||||
|
||||
set_max_concurrent_transmissions(client, app.max_concurrent_transmissions)
|
||||
|
||||
app.loop.run_until_complete(start_server(client))
|
||||
logger.success(_t("Successfully started (Press Ctrl+C to stop)"))
|
||||
|
||||
app.loop.create_task(download_all_chat(client))
|
||||
for _ in range(app.max_download_task):
|
||||
task = app.loop.create_task(worker(client))
|
||||
tasks.append(task)
|
||||
|
||||
if app.bot_token:
|
||||
app.loop.run_until_complete(
|
||||
start_download_bot(app, client, add_download_task, download_chat_task)
|
||||
)
|
||||
_exec_loop()
|
||||
except KeyboardInterrupt:
|
||||
logger.info(_t("KeyboardInterrupt"))
|
||||
except Exception as e:
|
||||
logger.exception("{}", e)
|
||||
finally:
|
||||
app.is_running = False
|
||||
if app.bot_token:
|
||||
app.loop.run_until_complete(stop_download_bot())
|
||||
app.loop.run_until_complete(stop_server(client))
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
logger.info(_t("Stopped!"))
|
||||
# check_for_updates(app.proxy)
|
||||
logger.info(f"{_t('update config')}......")
|
||||
app.update_config()
|
||||
logger.success(
|
||||
f"{_t('Updated last read message_id to config file')},"
|
||||
f"{_t('total download')} {app.total_download_task}, "
|
||||
f"{_t('total upload file')} "
|
||||
f"{app.cloud_drive_config.total_upload_success_file_count}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if _check_config():
|
||||
main()
|
||||
Reference in New Issue
Block a user