From 738a46c47fffc34a37da8c56e5bf4aac7a84dfde Mon Sep 17 00:00:00 2001 From: anasty17 Date: Fri, 27 Sep 2024 01:12:34 +0300 Subject: [PATCH] Select specific files or folders to download/copy from rclone using buttons - Fix minor issue in rclone list and gdrive list close #1668 Signed-off-by: anasty17 --- README.md | 7 +- .../download_utils/rclone_download.py | 90 ++++++++++++------- .../mirror_leech_utils/gdrive_utils/list.py | 1 + .../mirror_leech_utils/rclone_utils/list.py | 46 +++++++++- .../rclone_utils/transfer.py | 21 +++-- bot/modules/clone.py | 67 ++++++++------ 6 files changed, 162 insertions(+), 70 deletions(-) diff --git a/README.md b/README.md index bc037dc6acf..ca760d097af 100644 --- a/README.md +++ b/README.md @@ -124,9 +124,10 @@ programming in Python. ## Rclone -- Rclone transfer (download/upload/clone-server-side) without or with random service accounts (global and user option) -- Ability to choose config, remote and path from list with buttons (global, user and task option) -- Ability to set rclone flags for each task or globally from config (global, user and task option) +- Transfer (download/upload/clone-server-side) without or with random service accounts (global and user option) +- Ability to choose config, remote and path from list with or without buttons (global, user and task option) +- Ability to set flags for each task or globally from config (global, user and task option) +- Abitity to select specific files or folders to download/copy using buttons (task option) - Rclone.conf (global and user option) - Rclone serve for combine remote to use it as index from all remotes (global option) - Upload destination (global, user and task option) diff --git a/bot/helper/mirror_leech_utils/download_utils/rclone_download.py b/bot/helper/mirror_leech_utils/download_utils/rclone_download.py index cdf3b814c54..373e7fb725e 100644 --- a/bot/helper/mirror_leech_utils/download_utils/rclone_download.py +++ b/bot/helper/mirror_leech_utils/download_utils/rclone_download.py @@ -1,6 +1,7 @@ from asyncio import gather from json import loads from secrets import token_urlsafe +from aiofiles.os import remove from bot import task_dict, task_dict_lock, queue_dict_lock, non_queued_dl, LOGGER from ...ext_utils.bot_utils import cmd_exec @@ -20,6 +21,12 @@ async def add_rclone_download(listener, path): remote, listener.link = listener.link.split(":", 1) listener.link = listener.link.strip("/") + rclone_select = False + if listener.link.startswith("rclone_select"): + rclone_select = True + rpath = "" + else: + rpath = listener.link cmd1 = [ "rclone", @@ -30,7 +37,7 @@ async def add_rclone_download(listener, path): "--no-modtime", "--config", config_path, - f"{remote}:{listener.link}", + f"{remote}:{rpath}", ] cmd2 = [ "rclone", @@ -39,42 +46,63 @@ async def add_rclone_download(listener, path): "--json", "--config", config_path, - f"{remote}:{listener.link}", + f"{remote}:{rpath}", ] - res1, res2 = await gather(cmd_exec(cmd1), cmd_exec(cmd2)) - if res1[2] != res2[2] != 0: - if res1[2] != -9: - err = ( - res1[1] - or res2[1] - or "Use /shell cat rlog.txt to see more information" - ) - msg = f"Error: While getting rclone stat/size. Path: {remote}:{listener.link}. Stderr: {err[:4000]}" - await listener.on_download_error(msg) - return - try: - rstat = loads(res1[0]) - rsize = loads(res2[0]) - except Exception as err: - if not str(err): - err = "Use /shell cat rlog.txt to see more information" - await listener.on_download_error(f"RcloneDownload JsonLoad: {err}") - return - if rstat["IsDir"]: + if rclone_select: + cmd2.extend(("--files-from", listener.link)) + res = await cmd_exec(cmd2) + if res[2] != 0: + if res[2] != -9: + err = (res[1]or "Use /shell cat rlog.txt to see more information") + msg = f"Error: While getting rclone stat/size. Path: {remote}:{listener.link}. Stderr: {err[:4000]}" + await listener.on_download_error(msg) + return + try: + rsize = loads(res[0]) + except Exception as err: + if not str(err): + err = "Use /shell cat rlog.txt to see more information" + await listener.on_download_error(f"RcloneDownload JsonLoad: {err}") + return if not listener.name: - listener.name = ( - listener.link.rsplit("/", 1)[-1] if listener.link else remote - ) + listener.name = listener.link path += listener.name else: - listener.name = listener.link.rsplit("/", 1)[-1] + res1, res2 = await gather(cmd_exec(cmd1), cmd_exec(cmd2)) + if res1[2] != res2[2] != 0: + if res1[2] != -9: + err = ( + res1[1] + or res2[1] + or "Use /shell cat rlog.txt to see more information" + ) + msg = f"Error: While getting rclone stat/size. Path: {remote}:{listener.link}. Stderr: {err[:4000]}" + await listener.on_download_error(msg) + return + try: + rstat = loads(res1[0]) + rsize = loads(res2[0]) + except Exception as err: + if not str(err): + err = "Use /shell cat rlog.txt to see more information" + await listener.on_download_error(f"RcloneDownload JsonLoad: {err}") + return + if rstat["IsDir"]: + if not listener.name: + listener.name = ( + listener.link.rsplit("/", 1)[-1] if listener.link else remote + ) + path += listener.name + else: + listener.name = listener.link.rsplit("/", 1)[-1] listener.size = rsize["bytes"] gid = token_urlsafe(12) - msg, button = await stop_duplicate_check(listener) - if msg: - await listener.on_download_error(msg, button) - return + if not rclone_select: + msg, button = await stop_duplicate_check(listener) + if msg: + await listener.on_download_error(msg, button) + return add_to_queue, event = await check_running_tasks(listener) if add_to_queue: @@ -103,3 +131,5 @@ async def add_rclone_download(listener, path): LOGGER.info(f"Download with rclone: {listener.link}") await RCTransfer.download(remote, config_path, path) + if rclone_select: + await remove(listener.link) diff --git a/bot/helper/mirror_leech_utils/gdrive_utils/list.py b/bot/helper/mirror_leech_utils/gdrive_utils/list.py index cffc95b5679..5f456860f5d 100644 --- a/bot/helper/mirror_leech_utils/gdrive_utils/list.py +++ b/bot/helper/mirror_leech_utils/gdrive_utils/list.py @@ -68,6 +68,7 @@ async def id_updates(_, query, obj): obj.event.set() elif data[1] == "ps": if obj.page_step == int(data[2]): + obj.query_proc = False return obj.page_step = int(data[2]) await obj.get_items_buttons() diff --git a/bot/helper/mirror_leech_utils/rclone_utils/list.py b/bot/helper/mirror_leech_utils/rclone_utils/list.py index 0b547ddea61..f2fb2b3f4fa 100644 --- a/bot/helper/mirror_leech_utils/rclone_utils/list.py +++ b/bot/helper/mirror_leech_utils/rclone_utils/list.py @@ -43,6 +43,9 @@ async def path_updates(_, query, obj): elif data[1] == "nex": obj.iter_start += LIST_LIMIT * obj.page_step await obj.get_path_buttons() + elif data[1] == "select": + obj.select = not obj.select + await obj.get_path_buttons() elif data[1] == "back": if data[2] == "re": await obj.list_config() @@ -53,8 +56,31 @@ async def path_updates(_, query, obj): data = query.data.split(maxsplit=2) obj.remote = data[2] await obj.get_path() + elif data[1] == "clear": + obj.selected_pathes = set() + await obj.get_path_buttons() + elif data[1] == "ds": + obj.path = f"rclone_select_{time()}.txt" + async with aiopen(obj.path, "w") as txt_file: + for f in obj.selected_pathes: + await txt_file.write(f"{f}\n") + await delete_message(message) + obj.event.set() elif data[1] == "pa": index = int(data[3]) + if obj.select: + path = obj.path + ( + f"/{obj.path_list[index]['Path']}" + if obj.path + else obj.path_list[index]["Path"] + ) + if path in obj.selected_pathes: + obj.selected_pathes.remove(path) + else: + obj.selected_pathes.add(path) + await obj.get_path_buttons() + obj.query_proc = False + return obj.path += ( f"/{obj.path_list[index]['Path']}" if obj.path @@ -67,6 +93,7 @@ async def path_updates(_, query, obj): obj.event.set() elif data[1] == "ps": if obj.page_step == int(data[2]): + obj.query_proc = False return obj.page_step = int(data[2]) await obj.get_path_buttons() @@ -123,6 +150,8 @@ def __init__(self, listener): self.path_list = [] self.iter_start = 0 self.page_step = 1 + self.select = False + self.selected_pathes = set() async def _event_handler(self): pfunc = partial(path_updates, obj=self) @@ -162,12 +191,16 @@ async def get_path_buttons(self): self.path_list[self.iter_start : LIST_LIMIT + self.iter_start] ): orig_index = index + self.iter_start + name = idict["Path"] + if name in self.selected_pathes or any( + p.endswith(f"/{name}") for p in self.selected_pathes + ): + name = f"✅ {name}" if idict["IsDir"]: ptype = "fo" - name = idict["Path"] else: ptype = "fi" - name = f"[{get_readable_file_size(idict['Size'])}] {idict['Path']}" + name = f"[{get_readable_file_size(idict['Size'])}] {name}" buttons.data_button(name, f"rcq pa {ptype} {orig_index}") if items_no > LIST_LIMIT: for i in [1, 2, 4, 6, 10, 30, 50, 100]: @@ -185,6 +218,15 @@ async def get_path_buttons(self): ) if self.list_status == "rcu" or len(self.path_list) > 0: buttons.data_button("Choose Current Path", "rcq cur", position="footer") + if self.list_status == "rcd": + buttons.data_button( + f"Select: {'Enabled' if self.select else 'Disabled'}", + "rcq select", + position="footer", + ) + if len(self.selected_pathes) > 1: + buttons.data_button("Done With Selection", "rcq ds", position="footer") + buttons.data_button("Clear Selection", "rcq clear", position="footer") if self.list_status == "rcu": buttons.data_button("Set as Default Path", "rcq def", position="footer") if self.path or len(self._sections) > 1 or self._rc_user and self._rc_owner: diff --git a/bot/helper/mirror_leech_utils/rclone_utils/transfer.py b/bot/helper/mirror_leech_utils/rclone_utils/transfer.py index 141fd0a556f..c8bf82bad23 100644 --- a/bot/helper/mirror_leech_utils/rclone_utils/transfer.py +++ b/bot/helper/mirror_leech_utils/rclone_utils/transfer.py @@ -34,6 +34,7 @@ def __init__(self, listener): self._sa_index = 0 self._sa_number = 0 self._use_service_accounts = config_dict["USE_SERVICE_ACCOUNTS"] + self.rclone_select = False @property def transferred_size(self): @@ -120,7 +121,7 @@ async def _start_download(self, cmd, remote_type): if return_code == 0: await self._listener.on_download_complete() elif return_code != -9: - error = (await self._proc.stderr.read()).decode().strip() + error = (await self._proc.stderr.read()).decode().strip() or "Use /shell cat rlog.txt to see more information" if not error and remote_type == "drive" and self._use_service_accounts: error = "Mostly your service accounts don't have access to this drive!" elif not error: @@ -238,7 +239,7 @@ async def _start_upload(self, cmd, remote_type): if return_code == -9: return False elif return_code != 0: - error = (await self._proc.stderr.read()).decode().strip() + error = (await self._proc.stderr.read()).decode().strip() or "Use /shell cat rlog.txt to see more information" if not error and remote_type == "drive" and self._use_service_accounts: error = "Mostly your service accounts don't have access to this drive or RATE_LIMIT_EXCEEDED" elif not error: @@ -400,9 +401,7 @@ async def clone(self, config_path, src_remote, src_path, mime_type, method): if return_code == -9: return None, None elif return_code != 0: - error = ( - await self._proc.stderr.read() - ).decode().strip() or "Use /shell cat rlog.txt to see more information" + error = (await self._proc.stderr.read()).decode().strip() or "Use /shell cat rlog.txt to see more information" LOGGER.error(error) await self._listener.on_upload_error(error[:4000]) return None, None @@ -441,7 +440,11 @@ def _get_updated_command( ): if unwanted_files is None: unwanted_files = [] - ext = "*.{" + ",".join(self._listener.extension_filter) + "}" + if source.split(":")[-1].startswith("rclone_select"): + source = f"{source.split(":")[0]}:" + self.rclone_select = True + else: + ext = "*.{" + ",".join(self._listener.extension_filter) + "}" cmd = [ "rclone", method, @@ -451,8 +454,6 @@ def _get_updated_command( "-P", source, destination, - "--exclude", - ext, "--retries-sleep", "3s", "--ignore-case", @@ -464,6 +465,10 @@ def _get_updated_command( "--log-level", "DEBUG", ] + if self.rclone_select: + cmd.extend(("--files-from", self._listener.link)) + else: + cmd.extend(("--exclude", ext)) if rcflags := self._listener.rc_flags or config_dict["RCLONE_FLAGS"]: rcflags = rcflags.split("|") for flag in rcflags: diff --git a/bot/modules/clone.py b/bot/modules/clone.py index 4cbac647260..0d079dc7ef2 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -3,6 +3,7 @@ from pyrogram.filters import command from pyrogram.handlers import MessageHandler from secrets import token_urlsafe +from aiofiles.os import remove from bot import LOGGER, task_dict, task_dict_lock, bot, bot_loop from ..helper.ext_utils.bot_utils import ( @@ -70,6 +71,7 @@ async def new_event(self): "link": "", "-i": 0, "-b": False, + "-n": "", "-up": "", "-rcf": "", "-sync": False, @@ -85,6 +87,7 @@ async def new_event(self): self.up_dest = args["-up"] self.rc_flags = args["-rcf"] self.link = args["link"] + self.name = args["-n"] is_bulk = args["-b"] sync = args["-sync"] @@ -175,35 +178,43 @@ async def _proceed_to_clone(self, sync): config_path = "rclone.conf" remote, src_path = self.link.split(":", 1) - src_path = src_path.strip("/") - - cmd = [ - "rclone", - "lsjson", - "--fast-list", - "--stat", - "--no-modtime", - "--config", - config_path, - f"{remote}:{src_path}", - ] - res = await cmd_exec(cmd) - if res[2] != 0: - if res[2] != -9: - msg = f"Error: While getting rclone stat. Path: {remote}:{src_path}. Stderr: {res[1][:4000]}" - await send_message(self.message, msg) - return - rstat = loads(res[0]) - if rstat["IsDir"]: - self.name = src_path.rsplit("/", 1)[-1] if src_path else remote - self.up_dest += ( - self.name if self.up_dest.endswith(":") else f"/{self.name}" - ) - + self.link = src_path.strip("/") + if self.link.startswith("rclone_select"): mime_type = "Folder" + src_path = "" + if not self.name: + self.name = self.link else: - self.name = src_path.rsplit("/", 1)[-1] - mime_type = rstat["MimeType"] + src_path = self.link + cmd = [ + "rclone", + "lsjson", + "--fast-list", + "--stat", + "--no-modtime", + "--config", + config_path, + f"{remote}:{src_path}", + ] + res = await cmd_exec(cmd) + if res[2] != 0: + if res[2] != -9: + msg = f"Error: While getting rclone stat. Path: {remote}:{src_path}. Stderr: {res[1][:4000]}" + await send_message(self.message, msg) + return + rstat = loads(res[0]) + if rstat["IsDir"]: + if not self.name: + self.name = src_path.rsplit("/", 1)[-1] if src_path else remote + self.up_dest += ( + self.name if self.up_dest.endswith(":") else f"/{self.name}" + ) + + mime_type = "Folder" + else: + if not self.name: + self.name = src_path.rsplit("/", 1)[-1] + mime_type = rstat["MimeType"] await self.on_download_start() @@ -224,6 +235,8 @@ async def _proceed_to_clone(self, sync): mime_type, method, ) + if self.link.startswith("rclone_select"): + await remove(self.link) if not destination: return LOGGER.info(f"Cloning Done: {self.name}")