From d208e5697b129156392b2201697e28369ee61a3f Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Fri, 3 Jun 2022 18:52:38 -0700
Subject: [PATCH] Use voussoirkit.progressbars.

---
 src/mega/mega.py | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/mega/mega.py b/src/mega/mega.py
index 7460462..c7bd38a 100644
--- a/src/mega/mega.py
+++ b/src/mega/mega.py
@@ -19,6 +19,8 @@ from Crypto.Util import Counter
 from . import crypto
 from . import errors
 
+from voussoirkit import progressbars
+
 logger = logging.getLogger(__name__)
 
 NODE_TYPE_FILE = 0
@@ -547,6 +549,7 @@ class Mega:
             dest_filename=None,
             is_public=False,
             file=None,
+            progressbar=None,
         ):
         if file is None:
             if is_public:
@@ -607,8 +610,12 @@ class Mega:
             mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str.encode("utf8"))
             iv_str = crypto.a32_to_str([iv[0], iv[1], iv[0], iv[1]])
 
+            progressbar = progressbars.normalize_progressbar(progressbar, total=file_size)
+            download_progress = 0
+
             for chunk_start, chunk_size in crypto.get_chunks(file_size):
                 chunk = input_file.read(chunk_size)
+                download_progress += len(chunk)
                 chunk = aes.decrypt(chunk)
                 temp_output_file.write(chunk)
 
@@ -632,6 +639,10 @@ class Mega:
                 logger.debug(
                     '%s of %s downloaded', file_info.st_size, file_size
                 )
+                if progressbar is not None:
+                    progressbar.step(download_progress)
+            if progressbar is not None:
+                progressbar.done()
             file_mac = crypto.str_to_a32(mac_str)
             # check mac integrity
             if (file_mac[0] ^ file_mac[1],
@@ -641,7 +652,7 @@ class Mega:
             shutil.move(temp_output_file.name, output_path)
             return output_path
 
-    def download_file(self, file, dest_path=None, dest_filename=None):
+    def download_file(self, file, dest_path=None, dest_filename=None, progressbar=None):
         '''
         Download a file by its file object
         '''
@@ -652,9 +663,10 @@ class Mega:
             dest_path=dest_path,
             dest_filename=dest_filename,
             is_public=False,
+            progressbar=progressbar,
         )
 
-    def download_url(self, url, dest_path=None, dest_filename=None):
+    def download_url(self, url, dest_path=None, dest_filename=None, progressbar=None):
         '''
         Download a file by its public url
         '''
@@ -665,6 +677,7 @@ class Mega:
             dest_path=dest_path,
             dest_filename=dest_filename,
             is_public=True,
+            progressbar=progressbar,
         )
 
     # EMPTY TRASH ##################################################################################
@@ -1436,7 +1449,7 @@ class Mega:
 
     # UPLOAD #######################################################################################
 
-    def upload(self, filename, dest=None, dest_filename=None):
+    def upload(self, filename, dest=None, dest_filename=None, progressbar=None):
         # determine storage node
         if dest is None:
             # if none set, upload to cloud drive node
@@ -1462,6 +1475,7 @@ class Mega:
             mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str.encode("utf8"))
             iv_str = crypto.a32_to_str([ul_key[4], ul_key[5], ul_key[4], ul_key[5]])
             if file_size > 0:
+                progressbar = progressbars.normalize_progressbar(progressbar, total=file_size)
                 for chunk_start, chunk_size in crypto.get_chunks(file_size):
                     chunk = input_file.read(chunk_size)
                     upload_progress += len(chunk)
@@ -1491,6 +1505,10 @@ class Mega:
                     )
                     completion_file_handle = output_file.text
                     logger.debug('%s of %s uploaded', upload_progress, file_size)
+                    if progressbar is not None:
+                        progressbar.step(upload_progress)
+                if progressbar is not None:
+                    progressbar.done()
             else:
                 output_file = self.requests_session.post(
                     ul_url + "/0",