Update: Add some progress for mega downloader

All credit reserved to @adekmaulana

Signed-off-by: MoveAngel <[email protected]>

requirements.txt

@@ -43,5 +43,4 @@ telegraph
 urbandict>=0.5
 wikipedia>=1.4.0
 coffeehouse
-wget
 python-dateutil

userbot/modules/mega_downloader.py

@@ -18,17 +18,18 @@
 #  OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 #  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from humanize import naturalsize
 from subprocess import PIPE, Popen
 
 import re
 import json
-import wget
 import os
+import time
 
+from pySmartDL import SmartDL
 from os.path import exists
+from urllib.error import HTTPError
 
-from userbot import CMD_HELP
+from userbot import CMD_HELP, LOGS
 from userbot.events import register
 
 
@@ -48,72 +49,121 @@ def subprocess_run(cmd):
 
 
 @register(outgoing=True, pattern=r"^.mega(?: |$)(.*)")
-async def mega_downloader(event):
-    await event.edit("`Processing...`")
-    textx = await event.get_reply_message()
-    link = event.pattern_match.group(1)
+async def mega_downloader(megadl):
+    await megadl.edit("`Processing...`")
+    textx = await megadl.get_reply_message()
+    link = megadl.pattern_match.group(1)
     if link:
         pass
     elif textx:
         link = textx.text
     else:
-        await event.edit("`Usage: .mega <mega url>`")
+        await megadl.edit("`Usage: .mega <mega url>`")
         return
-    reply = ''
     if not link:
-        reply = "`No MEGA.nz link found!`"
-        await event.edit(reply)
-    await event.edit("`Downloading...`")
-    reply += mega_download(link)
-    await event.edit(reply)
+        await megadl.edit("`No MEGA.nz link found!`")
+    await mega_download(link, megadl)
 
 
-def mega_download(url: str) -> str:
-    reply = ''
+async def mega_download(url, megadl):
     try:
         link = re.findall(r'\bhttps?://.*mega.*\.nz\S+', url)[0]
     except IndexError:
-        reply = "`No MEGA.nz link found`\n"
-        return reply
+        await megadl.edit("`No MEGA.nz link found`\n")
+        return
     cmd = f'bin/megadirect {link}'
     result = subprocess_run(cmd)
     try:
         data = json.loads(result[0])
     except json.JSONDecodeError:
-        reply += "`Error: Can't extract the link`\n"
-        return reply
+        await megadl.edit("`Error: Can't extract the link`\n")
+        return
     file_name = data['file_name']
-    file_size = naturalsize(int(data['file_size']))
     file_url = data['url']
     file_hex = data['hex']
     file_raw_hex = data['raw_hex']
     if exists(file_name):
         os.remove(file_name)
     if not exists(file_name):
-        wget.download(file_url, out=file_name)
-        if exists(file_name):
-            encrypt_file(file_name, file_hex, file_raw_hex)
-            reply += (f"`{file_name}`\n"
-                      f"Size: {file_size}\n\n"
-                      "Successfully downloaded...")
+        temp_file_name = file_name + ".temp"
+        downloaded_file_name = "./" + "" + temp_file_name
+        downloader = SmartDL(file_url, downloaded_file_name, progress_bar=False)
+        display_message = None
+        try:
+            downloader.start(blocking=False)
+        except HTTPError as e:
+            await megadl.edit("`" + str(e) + "`")
+            LOGS.info(str(e))
+            return
+        while not downloader.isFinished():
+            status = downloader.get_status().capitalize()
+            total_length = downloader.filesize if downloader.filesize else None
+            downloaded = downloader.get_dl_size()
+            percentage = int(downloader.get_progress() * 100)
+            progress = downloader.get_progress_bar()
+            speed = downloader.get_speed(human=True)
+            estimated_total_time = downloader.get_eta(human=True)
+            try:
+                current_message = (
+                    f"**{status}**..."
+                    f"\nFile Name: `{file_name}`\n"
+                    f"\n{progress} `{percentage}%`"
+                    f"\n{humanbytes(downloaded)} of {humanbytes(total_length)}"
+                    f" @ {speed}"
+                    f"\nETA: {estimated_total_time}"
+                )
+                if status == "Downloading":
+                    await megadl.edit(current_message)
+                    time.sleep(0.2)
+                elif status == "Combining":
+                    if display_message != current_message:
+                        await megadl.edit(current_message)
+                        display_message = current_message
+            except Exception as e:
+                LOGS.info(str(e))
+        if downloader.isSuccessful():
+            download_time = downloader.get_dl_time(human=True)
+            if exists(temp_file_name):
+                await megadl.edit("Decrypting file...")
+                decrypt_file(file_name, temp_file_name, file_hex, file_raw_hex)
+                await megadl.edit(f"`{file_name}`\n\n"
+                                  "Successfully downloaded\n"
+                                  f"Download took: {download_time}")
         else:
-            reply += "Failed to download..."
-    return reply
+            await megadl.edit("Failed to download...")
+            for e in downloader.get_errors():
+                LOGS.info(str(e))
+    return
 
 
-def encrypt_file(file_name, file_hex, file_raw_hex):
-    os.rename(file_name, r"old_{}".format(file_name))
-    cmd = ("cat 'old_{}' | openssl enc -d -aes-128-ctr -K {} -iv {} > '{}'"
-           .format(file_name, file_hex, file_raw_hex, file_name))
+def decrypt_file(file_name, temp_file_name, file_hex, file_raw_hex):
+    cmd = ("cat '{}' | openssl enc -d -aes-128-ctr -K {} -iv {} > '{}'"
+           .format(temp_file_name, file_hex, file_raw_hex, file_name))
     subprocess_run(cmd)
-    os.remove(r"old_{}".format(file_name))
+    os.remove(r"{}".format(temp_file_name))
     return
 
 
+def humanbytes(size):
+    """Input size in bytes,
+    outputs in a human readable format"""
+    # https://stackoverflow.com/a/49361727/4723940
+    if not size:
+        return ""
+    # 2 ** 10 = 1024
+    power = 2**10
+    raised_to_pow = 0
+    dict_power_n = {0: "", 1: "Ki", 2: "Mi", 3: "Gi", 4: "Ti"}
+    while size > power:
+        size /= power
+        raised_to_pow += 1
+    return str(round(size, 2)) + " " + dict_power_n[raised_to_pow] + "B"
+
+
 CMD_HELP.update({
     "mega":
     ".mega <mega url>\n"
     "Usage: Reply to a mega link or paste your mega link to\n"
     "download the file into your userbot server\n\n"
-    "Only support for *FILE* only.\n"
+    "Only support for **FILE** only.\n"
 })