688 lines
20 KiB
Python
Executable File
688 lines
20 KiB
Python
Executable File
"""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()
|