Merge branch 'master' of code.richard.do:richardARPANET/mega.py

HISTORY.rst

@@ -3,12 +3,20 @@
 Release History
 ===============
 
-1.0.3 (unreleased)
+1.0.4 (unreleased)
 ------------------
 
 - Increase the wait time in between failed API request retries.
 
 
+1.0.3 (2019-11-12)
+------------------
+
+- Fixes broken ``download`` method.
+- Changes ``download`` and ``download_url`` methods to return the path to the downloaded file, previously returned ``None``.
+- Added LICENSE.
+
+
 1.0.2 (2019-11-07)
 ------------------
 

src/mega/mega.py

@@ -568,7 +568,7 @@ class Mega:
         """
         Download a file by it's file object
         """
-        self._download_file(
+        return self._download_file(
             file_handle=None,
             file_key=None,
             file=file[1],
@@ -646,7 +646,7 @@ class Mega:
         path = self._parse_url(url).split('!')
         file_id = path[0]
         file_key = path[1]
-        self._download_file(
+        return self._download_file(
             file_handle=file_id,
             file_key=file_key,
             dest_path=dest_path,
@@ -694,7 +694,7 @@ class Mega:
             iv = file['iv']
             meta_mac = file['meta_mac']
 
-        # Seems to happens sometime... When  this occurs, files are
+        # Seems to happens sometime... When this occurs, files are
         # inaccessible also in the official also in the official web app.
         # Strangely, files can come back later.
         if 'g' not in file_data:
@@ -716,51 +716,53 @@ class Mega:
         else:
             dest_path += '/'
 
-        temp_output_file = tempfile.NamedTemporaryFile(
+        with tempfile.NamedTemporaryFile(
             mode='w+b', prefix='megapy_', delete=False
-        )
-
-        k_str = a32_to_str(k)
-        counter = Counter.new(128, initial_value=((iv[0] << 32) + iv[1]) << 64)
-        aes = AES.new(k_str, AES.MODE_CTR, counter=counter)
-
-        mac_str = '\0' * 16
-        mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str)
-        iv_str = a32_to_str([iv[0], iv[1], iv[0], iv[1]])
-
-        for chunk_start, chunk_size in get_chunks(file_size):
-            chunk = input_file.read(chunk_size)
-            chunk = aes.decrypt(chunk)
-            temp_output_file.write(chunk)
-
-            encryptor = AES.new(k_str, AES.MODE_CBC, iv_str)
-            for i in range(0, len(chunk) - 16, 16):
-                block = chunk[i:i + 16]
-                encryptor.encrypt(block)
-
-            # fix for files under 16 bytes failing
-            if file_size > 16:
-                i += 16
-            else:
-                i = 0
+        ) as temp_output_file:
+            k_str = a32_to_str(k)
+            counter = Counter.new(
+                128, initial_value=((iv[0] << 32) + iv[1]) << 64
+            )
+            aes = AES.new(k_str, AES.MODE_CTR, counter=counter)
 
-            block = chunk[i:i + 16]
-            if len(block) % 16:
-                block += '\0' * (16 - (len(block) % 16))
-            mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
+            mac_str = '\0' * 16
+            mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str)
+            iv_str = a32_to_str([iv[0], iv[1], iv[0], iv[1]])
 
-            file_info = os.stat(temp_output_file.name)
-            logger.info('%s of %s downloaded', file_info.st_size, file_size)
+            for chunk_start, chunk_size in get_chunks(file_size):
+                chunk = input_file.read(chunk_size)
+                chunk = aes.decrypt(chunk)
+                temp_output_file.write(chunk)
 
-        file_mac = str_to_a32(mac_str)
+                encryptor = AES.new(k_str, AES.MODE_CBC, iv_str)
+                for i in range(0, len(chunk) - 16, 16):
+                    block = chunk[i:i + 16]
+                    encryptor.encrypt(block)
 
-        temp_output_file.close()
+                # fix for files under 16 bytes failing
+                if file_size > 16:
+                    i += 16
+                else:
+                    i = 0
 
-        # check mac integrity
-        if (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3]) != meta_mac:
-            raise ValueError('Mismatched mac')
+                block = chunk[i:i + 16]
+                if len(block) % 16:
+                    block += b'\0' * (16 - (len(block) % 16))
+                mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
 
-        shutil.move(temp_output_file.name, dest_path + file_name)
+                file_info = os.stat(temp_output_file.name)
+                logger.info(
+                    '%s of %s downloaded', file_info.st_size, file_size
+                )
+            file_mac = str_to_a32(mac_str)
+            # check mac integrity
+            if (
+                file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3]
+            ) != meta_mac:
+                raise ValueError('Mismatched mac')
+            output_path = Path(dest_path + file_name)
+            shutil.move(temp_output_file.name, output_path)
+            return output_path
 
     def upload(self, filename, dest=None, dest_filename=None):
         # determine storage node

src/tests/test_mega.py

@@ -28,6 +28,17 @@ def mega(folder_name):
     mega_.destroy(node_id)
 
 
+@pytest.fixture
+def uploaded_file(mega, folder_name):
+    folder = mega.find(folder_name)
+    dest_node_id = folder[1]['h']
+    mega.upload(
+        __file__, dest=dest_node_id, dest_filename='test.py'
+    )
+    path = f'{folder_name}/test.py'
+    return mega.find(path)
+
+
 def test_mega(mega):
     assert isinstance(mega, Mega)
 
@@ -56,11 +67,9 @@ def test_get_files(mega):
     assert isinstance(files, dict)
 
 
-def test_get_link(mega):
-    file = mega.find(TEST_FILE)
-    if file:
-        link = mega.get_link(file)
-        assert isinstance(link, str)
+def test_get_link(mega, uploaded_file):
+    link = mega.get_link(uploaded_file)
+    assert isinstance(link, str)
 
 
 class TestExport:
@@ -171,18 +180,31 @@ def test_delete_folder(mega, folder_name):
     assert isinstance(resp, int)
 
 
-def test_delete(mega):
-    file = mega.find(TEST_FILE)
-    if file:
-        resp = mega.delete(file[0])
-        assert isinstance(resp, int)
+def test_delete(mega, uploaded_file):
+    resp = mega.delete(uploaded_file[0])
+    assert isinstance(resp, int)
 
 
-def test_destroy(mega):
-    file = mega.find(TEST_FILE)
-    if file:
-        resp = mega.destroy(file[0])
-        assert isinstance(resp, int)
+def test_destroy(mega, uploaded_file):
+    resp = mega.destroy(uploaded_file[0])
+    assert isinstance(resp, int)
+
+
+def test_download(mega, tmpdir, folder_name):
+    # Upload a single file into a folder
+    folder = mega.find(folder_name)
+    dest_node_id = folder[1]['h']
+    mega.upload(
+        __file__, dest=dest_node_id, dest_filename='test.py'
+    )
+    path = f'{folder_name}/test.py'
+    file = mega.find(path)
+
+    output_path = mega.download(
+        file=file, dest_path=tmpdir, dest_filename='test.py'
+    )
+
+    assert output_path.exists()
 
 
 def test_empty_trash(mega):