From 2de0e3219954f0198bd14bf9747a9a3f59c76605 Mon Sep 17 00:00:00 2001
From: Abhi <85984486+AbhiTheModder@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:04:37 +0000
Subject: [PATCH] - Better Downloads: Add Progress Args to `urldl` for better
 UI - File Deletion Safety: Before deleting the file, the code checks if the
 file exists to avoid any FileNotFoundError exceptions.

---
 modules/removebg.py |  9 +++++---
 modules/stickers.py |  9 +++++---
 modules/upl.py      |  3 ++-
 modules/url.py      | 53 ++++++++++++++++++++++++++++++---------------
 modules/vt.py       |  8 ++++---
 5 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/modules/removebg.py b/modules/removebg.py
index 3015a55..8afdd9d 100644
--- a/modules/removebg.py
+++ b/modules/removebg.py
@@ -111,7 +111,8 @@ async def rmbg(client: Client, message: Message):
         allow_redirects=True,
         stream=True,
     )
-    os.remove(cool)
+    if os.path.exists(cool):
+        os.remove(cool)
     output_file_name = r
     contentType = output_file_name.headers.get("content-type")
     if "image" in contentType:
@@ -126,7 +127,8 @@ async def rmbg(client: Client, message: Message):
         await pablo.edit(
             "<code>Removed image's Background in {} seconds, powered by </code> <b>@moonuserbot</b>".format(ms)
         )
-        os.remove("BG_rem.png")
+        if os.path.exists("BG_rem.png"):
+            os.remove("BG_rem.png")
     else:
         await pablo.edit(
             "ReMove.BG API returned Errors. Please report to @moonub_chat\n`{}".format(
@@ -159,7 +161,8 @@ async def rembg(client: Client, message: Message):
     except Exception as e:
         await message.reply_text(f"An error occurred: {format_exc(e)}")
     finally:
-        os.remove(photo_data)
+        if os.path.exists(photo_data):
+            os.remove(photo_data)
 
 
 modules_help["removebg"] = {
diff --git a/modules/stickers.py b/modules/stickers.py
index 53506ba..e10865a 100644
--- a/modules/stickers.py
+++ b/modules/stickers.py
@@ -73,7 +73,8 @@ async def kang(client: Client, message: types.Message):
         return
 
     resized = resize_image(path)
-    os.remove(path)
+    if os.path.exists(path):
+        os.remove(path)
 
     await interact_with(await client.send_document("@stickers", resized, parse_mode=enums.ParseMode.MARKDOWN))
     response = await interact_with(await client.send_message("@stickers", emoji, parse_mode=enums.ParseMode.MARKDOWN))
@@ -99,7 +100,8 @@ async def stick2png(client: Client, message: types.Message):
         path = await message.reply_to_message.download()
         with open(path, "rb") as f:
             content = f.read()
-        os.remove(path)
+        if os.path.exists(path):
+            os.remove(path)
 
         file_io = BytesIO(content)
         file_io.name = "sticker.png"
@@ -120,7 +122,8 @@ async def resize_cmd(client: Client, message: types.Message):
         path = await message.reply_to_message.download()
         resized = resize_image(path)
         resized.name = "image.png"
-        os.remove(path)
+        if os.path.exists(path):
+            os.remove(path)
 
         await client.send_document(message.chat.id, resized, parse_mode=enums.ParseMode.MARKDOWN)
     except Exception as e:
diff --git a/modules/upl.py b/modules/upl.py
index f43e65a..2aeea2e 100644
--- a/modules/upl.py
+++ b/modules/upl.py
@@ -62,7 +62,8 @@ async def uplr(client: Client, message: Message):
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)
    finally:
-       os.remove(link)
+       if os.path.exists(link):
+        os.remove(link)
 
 modules_help["uplud"] = {
    "upl": f"[filepath]/[reply to path]*",
diff --git a/modules/url.py b/modules/url.py
index 221ef6b..2a04f0d 100644
--- a/modules/url.py
+++ b/modules/url.py
@@ -14,15 +14,16 @@
 #  You should have received a copy of the GNU General Public License
 #  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 import os
+import requests
+import time
+import traceback
 import urllib3
 from pyrogram import Client, filters, enums
 from pyrogram.types import Message
-import requests
 from utils.misc import modules_help, prefix
-from utils.scripts import format_exc
+from utils.scripts import format_exc, progress, edit_or_reply
 
 from io import BytesIO
-import requests
 
 from utils.config import apiflash_key
 
@@ -53,8 +54,6 @@ async def short(_, message: Message):
 
   await message.edit(r.data.decode().replace("https://", "<b>Shortened Url:</b>"), disable_web_page_preview=True, parse_mode=enums.ParseMode.HTML)
 
-
-
 @Client.on_message(filters.command("urldl", prefix) & filters.me)
 async def urldl(client: Client, message: Message):
     if len(message.command) > 1:
@@ -72,21 +71,33 @@ async def urldl(client: Client, message: Message):
     file_name = "downloads/" + link.split("/")[-1]
 
     try:
-        resp = requests.get(link)
+        resp = requests.get(link, stream=True)
         resp.raise_for_status()
 
+        # Get the total file size
+        total_size = int(resp.headers.get('content-length', 0))
+
         with open(file_name, "wb") as f:
+            downloaded = 0
+            chunk_count = 0
+            update_frequency = 800
             for chunk in resp.iter_content(chunk_size=8192):
                 f.write(chunk)
+                downloaded += len(chunk)
+                chunk_count += 1
+                progress_dl = (downloaded / total_size) * 100
+                if chunk_count % update_frequency == 0:
+                    await message.edit_text(f"<b>Downloading...</b>\nProgress: {progress_dl:.2f}%")
 
-        await message.edit("<b>Uploading...</b>", parse_mode=enums.ParseMode.HTML)
-        await client.send_document(message.chat.id, file_name, parse_mode=enums.ParseMode.HTML)
+        ms_ = await message.edit("<b>Uploading...</b>", parse_mode=enums.ParseMode.HTML)
+        c_time = time.time()
+        await client.send_document(message.chat.id, file_name, progress=progress, progress_args=(ms_, c_time, '`Uploading...`'), parse_mode=enums.ParseMode.MARKDOWN)
         await message.delete()
     except Exception as e:
-        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)
+        await message.edit(f"<b>Error:</b> {traceback.format_exc(e)}", parse_mode=enums.ParseMode.HTML)
     finally:
-        os.remove(file_name)
-
+        if os.path.exists(file_name):
+            os.remove(file_name)
 
 @Client.on_message(filters.command("upload", prefix) & filters.me)
 async def upload_cmd(_, message: Message):
@@ -96,20 +107,28 @@ async def upload_cmd(_, message: Message):
     min_file_age = 31
     max_file_age = 180
 
-    await message.edit("<b>Downloading...</b>", parse_mode=enums.ParseMode.HTML)
+    ms_ = await message.edit("`Downloading...`", parse_mode=enums.ParseMode.MARKDOWN)
+    c_time = time.time()
 
     try:
-        file_name = await message.download()
+        file_name = await message.download(
+            progress=progress,
+            progress_args=(ms_, c_time, '`Downloading...`')
+            )
     except ValueError:
         try:
-            file_name = await message.reply_to_message.download()
+            file_name = await message.reply_to_message.download(
+                progress=progress,
+                progress_args=(ms_, c_time, '`Downloading...`')
+                )
         except ValueError:
             await message.edit("<b>File to upload not found</b>", parse_mode=enums.ParseMode.HTML)
             return
 
     if os.path.getsize(file_name) > max_size:
         await message.edit(f"<b>Files longer than {max_size_mb}MB isn't supported</b>", parse_mode=enums.ParseMode.HTML)
-        os.remove(file_name)
+        if os.path.exists(file_name):
+            os.remove(file_name)
         return
 
     await message.edit("<b>Uploading...</b>", parse_mode=enums.ParseMode.HTML)
@@ -135,8 +154,8 @@ async def upload_cmd(_, message: Message):
     else:
         await message.edit(f"<b>API returned an error!\n" f"{response.text}\n Not allowed</b>", parse_mode=enums.ParseMode.HTML)
         print(response.text)
-
-    os.remove(file_name)
+    if os.path.exists(file_name):
+        os.remove(file_name)
 
 
 
diff --git a/modules/vt.py b/modules/vt.py
index 427c5b9..bd0dfc4 100644
--- a/modules/vt.py
+++ b/modules/vt.py
@@ -14,7 +14,7 @@ from pyrogram import Client, filters, enums
 from pyrogram.types import Message
 
 from utils.misc import modules_help, prefix
-from utils.scripts import edit_or_reply, format_exc, time_formatter, humanbytes, edit_or_send_as_file, get_text, progress
+from utils.scripts import edit_or_reply, format_exc, progress
 from utils.config import vt_key as vak
 
 
@@ -48,7 +48,8 @@ async def scan_my_file(client, message):
     except Exception as e:
       return await ms_.edit(format_exc(e))
     await ms_.edit(f"<b><u>Scanned {message.reply_to_message.document.file_name}</b></u>. <b>You Can Visit :</b> <a href=\"https://www.virustotal.com/gui/file/{md5}\">Here</a> <b>In 5-10 Min To See File Report</b>")
-    os.remove(downloaded_file_name)
+    if os.path.exists(downloaded_file_name):
+      os.remove(downloaded_file_name)
 
 @Client.on_message(filters.command("vtl", prefix) & filters.me)
 async def scan_my_large_file(client, message):
@@ -119,7 +120,8 @@ async def scan_my_large_file(client, message):
     except Exception as e:
       return await ms_.edit(format_exc(e))
     await ms_.edit(f"<b><u>Scanned {message.reply_to_message.document.file_name}</b></u>. <b>You Can Visit :</b> <a href=\"https://www.virustotal.com/gui/file/{md5}\">Here</a> <b>In 5-10 Min To See File Report</b>")
-    os.remove(downloaded_file_name)
+    if os.path.exists(downloaded_file_name):
+      os.remove(downloaded_file_name)
 
 modules_help["virustotal"] = {
     "vt [reply to file]*": "Scan for viruses on Virus Total (for lower file size <32MB)",