From 6262b26fe51f6e722d8bd7490f3406c4c3ed81e3 Mon Sep 17 00:00:00 2001
From: jaskaranSM <beerabassian32@gmail.com>
Date: Sat, 28 Mar 2020 17:37:17 +0530
Subject: [PATCH] youtube-dl: fixed downloading from sites that provide
 encrypted content (m3u8)

---
 bot/helper/ext_utils/fs_utils.py                       | 10 ++++++++++
 .../download_utils/youtube_dl_download_helper.py       |  6 ++++--
 bot/modules/mirror.py                                  |  8 +++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py
index 019cc77..8f31692 100644
--- a/bot/helper/ext_utils/fs_utils.py
+++ b/bot/helper/ext_utils/fs_utils.py
@@ -30,6 +30,16 @@ def exit_clean_up(signal, frame):
         LOGGER.warning("Force Exiting before the cleanup finishes!")
         sys.exit(1)
 
+def get_path_size(path):
+    if os.path.isfile(path):
+        return os.path.getsize(path)
+    total_size = 0
+    for root, dirs, files in os.walk(path):
+        for f in files:
+            abs_path = os.path.join(root, f)
+            total_size += os.path.getsize(abs_path)
+    return total_size
+
 
 def tar(org_path):
     tar_path = org_path + ".tar"
diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py
index 2a590e5..334858d 100644
--- a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py
+++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py
@@ -72,13 +72,15 @@ class YoutubeDLHelper(DownloadHelper):
         if 'entries' in result:
             video = result['entries'][0]
             for v in result['entries']:
-                self.size += float(v['filesize'])
+                if v.get('filesize'):
+                    self.size += float(v['filesize'])
             self.name = video.get('playlist_title')
             self.vid_id = video.get('id')
             self.is_playlist = True
         else:
             video = result
-            self.size = int(video.get('filesize'))
+            if video.get('filesize'):
+                self.size = int(video.get('filesize'))
             self.name = f"{video.get('title')}.{video.get('ext')}"
             self.vid_id = video.get('id')
         return video
diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py
index 4b28563..a92623b 100644
--- a/bot/modules/mirror.py
+++ b/bot/modules/mirror.py
@@ -59,10 +59,12 @@ class MirrorListener(listeners.MirrorListeners):
         else:
             path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}'
         up_name = pathlib.PurePath(path).name
+        LOGGER.info(f"Upload Name : {up_name}")
+        drive = gdriveTools.GoogleDriveHelper(up_name, self)
+        if size == 0:
+            size = fs_utils.get_path_size(m_path)
+        upload_status = UploadStatus(drive, size, self.uid)
         with download_dict_lock:
-            LOGGER.info(f"Upload Name : {up_name}")
-            drive = gdriveTools.GoogleDriveHelper(up_name, self)
-            upload_status = UploadStatus(drive, size, self.uid)
             download_dict[self.uid] = upload_status
         update_all_messages()
         drive.upload(up_name)