From 11be96f15ad83c2f339abe88b07e5df423e947b6 Mon Sep 17 00:00:00 2001
From: thedragonsinn <98635854+thedragonsinn@users.noreply.github.com>
Date: Sun, 12 Nov 2023 09:22:57 +0530
Subject: [PATCH] fix input bug and catch exceptions up/download.

---
 app/plugins/admin/fbans.py   |  4 ++--
 app/plugins/misc/download.py | 30 ++++++++++++++++----------
 app/plugins/misc/upload.py   | 42 +++++++++++++++++++++++++-----------
 3 files changed, 51 insertions(+), 25 deletions(-)

diff --git a/app/plugins/admin/fbans.py b/app/plugins/admin/fbans.py
index f5ff704..0594308 100644
--- a/app/plugins/admin/fbans.py
+++ b/app/plugins/admin/fbans.py
@@ -105,7 +105,7 @@ async def fed_ban(bot: bot, message: Message):
             chat_id=chat_id, text=fban_cmd, disable_web_page_preview=True
         )
         response: Message | None = await cmd.get_response(filters=FILTERS, timeout=8)
-        if not response or not (await FBAN_REGEX(bot, response)):
+        if not response or not (await FBAN_REGEX(bot, response)):  # NOQA
             failed.append(fed["name"])
         elif "Would you like to update this reason" in response.text:
             await response.click("Update reason")
@@ -113,7 +113,7 @@ async def fed_ban(bot: bot, message: Message):
     if not total:
         await progress.edit("You Don't have any feds connected!")
         return
-    resp_str = f"❯❯❯ <b>FBanned</b> {user.mention}\n<b>ID</b>: {user.id}\n<b>Reason</b>: {reason}\n<b>Initiated in</b>:{message.chat.title or 'PM'}"
+    resp_str = f"❯❯❯ <b>FBanned</b> {user.mention}\n<b>ID</b>: {user.id}\n<b>Reason</b>: {reason}\n<b>Initiated in</b>: {message.chat.title or 'PM'}"
     if failed:
         resp_str += f"\n<b>Failed</b> in: {len(failed)}/{total}\n• " + "\n• ".join(
             failed
diff --git a/app/plugins/misc/download.py b/app/plugins/misc/download.py
index b5a59eb..3f2d8fd 100644
--- a/app/plugins/misc/download.py
+++ b/app/plugins/misc/download.py
@@ -19,23 +19,31 @@ async def down_load(bot: BOT, message: Message):
     dl_path = os.path.join("downloads", str(time.time()))
     await response.edit("Input verified....Starting Download...")
     if message.replied and message.replied.media:
-        downloaded_file: DownloadedFile = await telegram_download(
+        download_coro = telegram_download(
             message=message.replied, response=response, path=dl_path
         )
     else:
         dl_obj: Download = await Download.setup(
             url=message.input, path=dl_path, message_to_edit=response
         )
-        downloaded_file: DownloadedFile = await dl_obj.download()
+        download_coro = dl_obj.download()
+    try:
+        downloaded_file: DownloadedFile = await download_coro
+        await response.edit(
+            f"<b>Download Completed</b>"
+            f"\n<pre language=bash>"
+            f"\nfile={downloaded_file.name}"
+            f"\npath={downloaded_file.full_path}"
+            f"\nsize={downloaded_file.size}mb</pre>"
+        )
+        return downloaded_file
 
-    await response.edit(
-        f"<b>Download Completed</b>"
-        f"\n<pre language=bash>"
-        f"\nfile={downloaded_file.name}"
-        f"\npath={downloaded_file.full_path}"
-        f"\nsize={downloaded_file.size}mb</pre>"
-    )
-    return downloaded_file
+    except asyncio.exceptions.CancelledError:
+        await response.edit("Cancelled....")
+    except TimeoutError:
+        await response.edit("Download Timeout...")
+    except Exception as e:
+        await response.edit(str(e))
 
 
 async def telegram_download(
@@ -54,6 +62,6 @@ async def telegram_download(
         media_obj.name,
         media_obj.full_path)
     await message.download(
-        file_name=media_obj.full_path, progress=progress, progress_args=progress_args
+        file_name=media_obj.full_path, progress=progress, progress_args=progress_args,
     )
     return media_obj
diff --git a/app/plugins/misc/upload.py b/app/plugins/misc/upload.py
index 4c24611..0ed0460 100644
--- a/app/plugins/misc/upload.py
+++ b/app/plugins/misc/upload.py
@@ -91,11 +91,18 @@ async def upload(bot: BOT, message: Message):
         return
     elif input.startswith("http") and not file_check(input):
         dl_obj: Download = await Download.setup(
-            url=message.input,
+            url=input,
             path=os.path.join("downloads", str(time.time())),
             message_to_edit=response,
         )
-        file: DownloadedFile = await dl_obj.download()
+        try:
+            file: DownloadedFile = await dl_obj.download()
+        except asyncio.exceptions.CancelledError:
+            await response.edit("Cancelled...")
+            return
+        except TimeoutError:
+            await response.edit("Download Timeout...")
+            return
     elif file_check(input):
         file = DownloadedFile(
             name=input,
@@ -107,21 +114,32 @@ async def upload(bot: BOT, message: Message):
         await response.edit("invalid `cmd` | `url` | `file path`!!!")
         return
     await response.edit("uploading....")
-    progress_args = (response, "Uploading...", file.name, file.full_path)
+    progress_args = (
+        response,
+        "Uploading...",
+        file.name,
+        file.full_path,
+    )
     if "-d" in message.flags:
         media: dict = {
             "method": bot.send_document,
-            "kwargs": {"document": file.full_path, "force_document": True},
+            "kwargs": {
+                "document": file.full_path,
+                "force_document": True,
+            },
         }
     else:
         media: dict = await FILE_TYPE_MAP[file.type](
             file, has_spoiler="-s" in message.flags
         )
-    await media["method"](
-        chat_id=message.chat.id,
-        reply_to_message_id=message.reply_id,
-        progress=progress,
-        progress_args=progress_args,
-        **media["kwargs"]
-    )
-    await response.delete()
+    try:
+        await media["method"](
+            chat_id=message.chat.id,
+            reply_to_message_id=message.reply_id,
+            progress=progress,
+            progress_args=progress_args,
+            **media["kwargs"]
+        )
+        await response.delete()
+    except asyncio.exceptions.CancelledError:
+        await response.edit("Cancelled....")