From e060278dcae17d2a58ff962d02ccc020df3918df Mon Sep 17 00:00:00 2001
From: tangyoha
-
-
\ No newline at end of file + diff --git a/docker-compose.yaml b/docker-compose.yaml index cf6fe5ca..85522996 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -7,10 +7,16 @@ services: ports: # Here is what you need to edit - "5000:5000" + #environment: + # - http_proxy=http://192.168.101.30:10811 + # - https_proxy=http://192.168.101.30:10811 volumes: # Here is what you need to edit - "./downloads/:/app/downloads/" + # Rclone + - "$HOME/.config/rclone/:$HOME/.config/rclone/" + # The following is what you don't need to edit - "./config.yaml:/app/config.yaml" - "./data.yaml:/app/data.yaml" diff --git a/media_downloader.py b/media_downloader.py index fc109fcd..9c633a09 100644 --- a/media_downloader.py +++ b/media_downloader.py @@ -287,6 +287,14 @@ async def download_task( download_status, ) + # rclone upload + if ( + not node.upload_telegram_chat_id + and download_status is DownloadStatus.SuccessDownload + ): + if await app.upload_file(file_name): + node.upload_success_count += 1 + await report_bot_download_status( node.bot, node, @@ -294,13 +302,6 @@ async def download_task( file_size, ) - # rclone upload - if ( - not node.upload_telegram_chat_id - and download_status is DownloadStatus.SuccessDownload - ): - await app.upload_file(file_name) - # pylint: disable = R0915,R0914 @@ -480,7 +481,10 @@ async def worker(client: pyrogram.client.Client): message = item[0] node: TaskNode = item[1] - await download_task(client, message, node) + 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}") diff --git a/module/app.py b/module/app.py index 68d2dc43..b7f122dd 100644 --- a/module/app.py +++ b/module/app.py @@ -89,6 +89,8 @@ def __init__( self.failed_forward_task: int = 0 self.skip_forward_task: int = 0 self.is_running: bool = False + self.client = None + self.upload_success_count: int = 0 def is_finish(self): """If is finish""" diff --git a/module/bot.py b/module/bot.py index e22c1060..7a5ccd1e 100644 --- a/module/bot.py +++ b/module/bot.py @@ -350,6 +350,7 @@ async def direct_download( chat_id: Union[str, int], message: pyrogram.types.Message, download_message: pyrogram.types.Message, + client: pyrogram.Client = None, ): """Direct Download""" @@ -368,6 +369,8 @@ async def direct_download( task_id=_bot.gen_task_id(), ) + node.client = client + _bot.add_task_node(node) await _bot.add_download_task( @@ -393,7 +396,7 @@ async def download_forward_media( """ if message.media and getattr(message, message.media.value): - await direct_download(_bot, message.from_user.id, message, message) + await direct_download(_bot, message.from_user.id, message, message, client) return await client.send_message( diff --git a/module/cloud_drive.py b/module/cloud_drive.py index 339a9c33..6e80ab4b 100644 --- a/module/cloud_drive.py +++ b/module/cloud_drive.py @@ -86,8 +86,9 @@ def zip_file(local_file_path: str) -> str: @staticmethod async def rclone_upload_file( drive_config: CloudDriveConfig, save_path: str, local_file_path: str - ): + ) -> bool: """Use Rclone upload file""" + upload_status: bool = False try: remote_dir = ( drive_config.remote_dir @@ -125,19 +126,24 @@ async def rclone_upload_file( os.remove(local_file_path) if drive_config.before_upload_file_zip: os.remove(zip_file_path) + upload_status = True await proc.wait() except Exception as e: logger.error(f"{e.__class__} {e}") + return False + + return upload_status @staticmethod def aligo_upload_file( drive_config: CloudDriveConfig, save_path: str, local_file_path: str ): """aliyun upload file""" + upload_status: bool = False if not drive_config.aligo: logger.warning("please config aligo! see README.md") - return + return False try: remote_dir = ( @@ -174,13 +180,18 @@ def aligo_upload_file( if drive_config.before_upload_file_zip: os.remove(zip_file_path) + upload_status = True + except Exception as e: logger.error(f"{e.__class__} {e}") + return False + + return upload_status @staticmethod async def upload_file( drive_config: CloudDriveConfig, save_path: str, local_file_path: str - ): + ) -> bool: """Upload file Parameters ---------- @@ -192,13 +203,21 @@ async def upload_file( local_file_path: str Local file path + + Returns + ------- + bool + True or False """ if not drive_config.enable_upload_file: - return + return False + ret: bool = False if drive_config.upload_adapter == "rclone": - await CloudDrive.rclone_upload_file( + ret = await CloudDrive.rclone_upload_file( drive_config, save_path, local_file_path ) elif drive_config.upload_adapter == "aligo": - CloudDrive.aligo_upload_file(drive_config, save_path, local_file_path) + ret = CloudDrive.aligo_upload_file(drive_config, save_path, local_file_path) + + return ret diff --git a/module/pyrogram_extension.py b/module/pyrogram_extension.py index f1d73985..8397fe35 100644 --- a/module/pyrogram_extension.py +++ b/module/pyrogram_extension.py @@ -408,6 +408,14 @@ async def report_bot_status( f"└─ ⏩ {_t('Skipped')}: {node.skip_forward_task}\n" ) + upload_msg_detail_str: str = "" + + if node.upload_success_count: + upload_msg_detail_str = ( + f"\n📥 {_t('Upload')}\n" + f"└─ ✅ {_t('Success')}: {node.upload_success_count}\n" + ) + download_result_str = "" download_result = get_download_result() if node.chat_id in download_result: @@ -441,6 +449,7 @@ async def report_bot_status( f"├─ ❌ {_t('Failed')}: {node.failed_download_task}\n" f"└─ ⏩ {_t('Skipped')}: {node.skip_download_task}\n" f"{node.forward_msg_detail_str}" + f"{upload_msg_detail_str}" f"{download_result_str}\n```" )