Merge pull request #3 from richardARPANET/misc-improvements

Misc improvements

HISTORY.rst

@@ -6,7 +6,10 @@ Release History
 1.0.1 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- When a request fails due to EAGAIN response, retry with exp backoff up to 20 seconds.
+- Adds logging, removes print statements.
+- Replace pycrypto dependency with pycryptodome.
+- Removes Python 2 specific code.
 
 
 1.0.0 (2019-10-31)

README.rst

@@ -56,8 +56,6 @@ Create an instance of Mega.py
 .. code:: python
 
    mega = Mega()
-   # add the verbose option for print output on some functions
-   mega = Mega({'verbose': True})
 
 Login to Mega
 ~~~~~~~~~~~~~

requirements.txt

@@ -1,4 +1,4 @@
 requests>=0.10
-pycrypto
+pycryptodome>=3.9.1,<4.0.0
 pathlib==1.0.1
-python2-secrets==1.0.5
+tenacity>=5.1.5,<6.0.0

src/mega/mega.py

@@ -1,6 +1,6 @@
 import re
-import time
 import json
+import logging
 import secrets
 from pathlib import Path
 import hashlib
@@ -8,23 +8,25 @@ from Crypto.Cipher import AES
 from Crypto.PublicKey import RSA
 from Crypto.Util import Counter
 import os
-import sys
 import random
 import binascii
-import requests
+import tempfile
 import shutil
+
+import requests
+from tenacity import retry, wait_exponential, retry_if_exception_type
+
 from .errors import ValidationError, RequestError
 from .crypto import (
     a32_to_base64, encrypt_key, base64_url_encode, encrypt_attr, base64_to_a32,
     base64_url_decode, decrypt_attr, a32_to_str, get_chunks, str_to_a32,
     decrypt_key, mpi_to_int, stringhash, prepare_key, make_id, makebyte
 )
-import tempfile
 
-PYTHON2 = sys.version_info < (3, )
+logger = logging.getLogger(__name__)
 
 
-class Mega(object):
+class Mega:
     def __init__(self, options=None):
         self.schema = 'https'
         self.domain = 'mega.co.nz'
@@ -44,9 +46,11 @@ class Mega(object):
         else:
             self.login_anonymous()
         self._trash_folder_node_id = self.get_node_by_type(4)[0]
+        logger.info('Login complete')
         return self
 
     def _login_user(self, email, password):
+        logger.info('Logging in user...')
         email = email.lower()
         get_user_salt_resp = self._api_request({'a': 'us0', 'user': email})
         user_salt = None
@@ -73,6 +77,7 @@ class Mega(object):
         self._login_process(resp, password_aes)
 
     def login_anonymous(self):
+        logger.info('Logging in anonymous temporary user...')
         master_key = [random.randint(0, 0xFFFFFFFF)] * 4
         password_key = [random.randint(0, 0xFFFFFFFF)] * 4
         session_self_challenge = [random.randint(0, 0xFFFFFFFF)] * 4
@@ -116,12 +121,6 @@ class Mega(object):
             private_key = a32_to_str(rsa_private_key)
             self.rsa_private_key = [0, 0, 0, 0]
             for i in range(4):
-                if PYTHON2:
-                    l = (
-                        (ord(private_key[0]) * 256 +
-                         ord(private_key[1]) + 7) / 8
-                    ) + 2
-                else:
                 l = int(
                     ((private_key[0]) * 256 + (private_key[1]) + 7) / 8
                 ) + 2
@@ -140,6 +139,10 @@ class Mega(object):
             sid = binascii.unhexlify('0' + sid if len(sid) % 2 else sid)
             self.sid = base64_url_encode(sid[:43])
 
+    @retry(
+        retry=retry_if_exception_type(RuntimeError),
+        wait=wait_exponential(multiplier=1, min=1, max=20)
+    )
     def _api_request(self, data):
         params = {'id': self.sequence_num}
         self.sequence_num += 1
@@ -151,7 +154,7 @@ class Mega(object):
         if not isinstance(data, list):
             data = [data]
 
-        url = '{0}://g.api.{1}/cs'.format(self.schema, self.domain)
+        url = f'{self.schema}://g.api.{self.domain}/cs'
         req = requests.post(
             url,
             params=params,
@@ -161,8 +164,9 @@ class Mega(object):
         json_resp = json.loads(req.text)
         if isinstance(json_resp, int):
             if json_resp == -3:
-                time.sleep(0.2)
+                msg = 'Request failed, retrying'
-                return self._api_request(data=data)
+                logger.info(msg)
+                raise RuntimeError(msg)
             raise RequestError(json_resp)
         return json_resp[0]
 
@@ -332,9 +336,7 @@ class Mega(object):
                 return file
 
     def get_files(self):
-        """
+        logger.info('Getting all files...')
-        Get all files in account
-        """
         files = self._api_request({'a': 'f', 'c': 1, 'r': 1})
         files_dict = {}
         shared_keys = {}
@@ -358,8 +360,9 @@ class Mega(object):
             decrypted_key = a32_to_base64(
                 decrypt_key(base64_to_a32(file_key), self.master_key)
             )
-            return '{0}://{1}/#!{2}!{3}'.format(
+            return (
-                self.schema, self.domain, public_handle, decrypted_key
+                f'{self.schema}://{self.domain}'
+                f'/#!{public_handle}!{decrypted_key}'
             )
         else:
             raise ValueError(
@@ -380,8 +383,9 @@ class Mega(object):
                     "(is this a shared file?)"
                 )
             decrypted_key = a32_to_base64(file['key'])
-            return '{0}://{1}/#!{2}!{3}'.format(
+            return (
-                self.schema, self.domain, public_handle, decrypted_key
+                f'{self.schema}://{self.domain}'
+                f'/#!{public_handle}!{decrypted_key}'
             )
         else:
             raise ValidationError('File id and key must be present')
@@ -405,8 +409,9 @@ class Mega(object):
                     "(is this a shared file?)"
                 )
             decrypted_key = a32_to_base64(file['shared_folder_key'])
-            return '{0}://{1}/#F!{2}!{3}'.format(
+            return (
-                self.schema, self.domain, public_handle, decrypted_key
+                f'{self.schema}://{self.domain}'
+                f'/#F!{public_handle}!{decrypted_key}'
             )
         else:
             raise ValidationError('File id and key must be present')
@@ -564,8 +569,8 @@ class Mega(object):
         Download a file by it's file object
         """
         self._download_file(
-            None,
+            file_handle=None,
-            None,
+            file_key=None,
             file=file[1],
             dest_path=dest_path,
             dest_filename=dest_filename,
@@ -642,7 +647,11 @@ class Mega(object):
         file_id = path[0]
         file_key = path[1]
         self._download_file(
-            file_id, file_key, dest_path, dest_filename, is_public=True
+            file_handle=file_id,
+            file_key=file_key,
+            dest_path=dest_path,
+            dest_filename=dest_filename,
+            is_public=True,
         )
 
     def _download_file(
@@ -740,14 +749,8 @@ class Mega(object):
                 block += '\0' * (16 - (len(block) % 16))
             mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
 
-            if self.options.get('verbose') is True:
-                # temp file size
             file_info = os.stat(temp_output_file.name)
-                print((
+            logger.info('%s of %s downloaded', file_info.st_size, file_size)
-                    '{0} of {1} downloaded'.format(
-                        file_info.st_size, file_size
-                    )
-                ))
 
         file_mac = str_to_a32(mac_str)
 
@@ -768,7 +771,7 @@ class Mega(object):
             dest = self.root_id
 
         # request upload url, call 'u' method
-        input_file = open(filename, 'rb')
+        with open(filename, 'rb') as input_file:
             file_size = os.path.getsize(filename)
             ul_url = self._api_request({'a': 'u', 's': file_size})['p']
 
@@ -815,44 +818,42 @@ class Mega(object):
                         timeout=self.timeout
                     )
                     completion_file_handle = output_file.text
-
+                    logger.info(
-                if self.options.get('verbose') is True:
+                        '%s of %s uploaded', upload_progress, file_size
-                    # upload progress
-                    print((
-                        '{0} of {1} uploaded'.format(
-                            upload_progress, file_size
                     )
-                    ))
             else:
                 output_file = requests.post(
                     ul_url + "/0", data='', timeout=self.timeout
                 )
                 completion_file_handle = output_file.text
 
+            logger.info('Chunks uploaded')
+            logger.info('Setting attributes to complete upload')
+            logger.info('Computing attributes')
             file_mac = str_to_a32(mac_str)
 
             # determine meta mac
             meta_mac = (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3])
 
-        if dest_filename is not None:
+            dest_filename = dest_filename or os.path.basename(filename)
             attribs = {'n': dest_filename}
-        else:
-            attribs = {'n': os.path.basename(filename)}
 
-        encrypt_attribs = base64_url_encode(encrypt_attr(attribs, ul_key[:4]))
+            encrypt_attribs = base64_url_encode(
+                encrypt_attr(attribs, ul_key[:4])
+            )
             key = [
                 ul_key[0] ^ ul_key[4], ul_key[1] ^ ul_key[5],
                 ul_key[2] ^ meta_mac[0], ul_key[3] ^ meta_mac[1], ul_key[4],
                 ul_key[5], meta_mac[0], meta_mac[1]
             ]
             encrypted_key = a32_to_base64(encrypt_key(key, self.master_key))
+            logger.info('Sending request to update attributes')
             # update attributes
             data = self._api_request(
                 {
-                'a':
+                    'a': 'p',
-                'p',
+                    't': dest,
-                't':
+                    'i': self.request_id,
-                dest,
                     'n': [
                         {
                             'h': completion_file_handle,
@@ -863,8 +864,7 @@ class Mega(object):
                     ]
                 }
             )
-        # close input file and return API msg
+            logger.info('Upload complete')
-        input_file.close()
             return data
 
     def _mkdir(self, name, parent_node_id):
@@ -930,9 +930,8 @@ class Mega(object):
         encrypted_key = a32_to_base64(
             encrypt_key(file['key'], self.master_key)
         )
-
         # update attributes
-        data = self._api_request(
+        return self._api_request(
             [
                 {
                     'a': 'a',
@@ -943,7 +942,6 @@ class Mega(object):
                 }
             ]
         )
-        return data
 
     def move(self, file_id, target):
         """
@@ -968,7 +966,7 @@ class Mega(object):
         # determine target_node_id
         if type(target) == int:
             target_node_id = str(self.get_node_by_type(target)[0])
-        elif type(target) in (str, str):
+        elif type(target) in (str, ):
             target_node_id = target
         else:
             file = target[1]
@@ -1053,9 +1051,7 @@ class Mega(object):
         unencrypted_attrs = decrypt_attr(base64_url_decode(data['at']), k)
         if not unencrypted_attrs:
             return None
-
         result = {'size': size, 'name': unencrypted_attrs['n']}
-
         return result
 
     def import_public_file(
@@ -1064,7 +1060,6 @@ class Mega(object):
         """
         Import the public file into user account
         """
-
         # Providing dest_node spare an API call to retrieve it.
         if dest_node is None:
             # Get '/Cloud Drive' folder no dest node specified
@@ -1082,13 +1077,10 @@ class Mega(object):
 
         encrypted_key = a32_to_base64(encrypt_key(key, self.master_key))
         encrypted_name = base64_url_encode(encrypt_attr({'n': dest_name}, k))
-
+        return self._api_request(
-        data = self._api_request(
             {
-                'a':
+                'a': 'p',
-                'p',
+                't': dest_node['h'],
-                't':
-                dest_node['h'],
                 'n': [
                     {
                         'ph': file_handle,
@@ -1099,4 +1091,3 @@ class Mega(object):
                 ]
             }
         )
-        return data

src/tests/test_crypto.py

@@ -0,0 +1,42 @@
+import pytest
+
+from mega.crypto import get_chunks
+
+
+@pytest.mark.parametrize(
+    'file_size, exp_result', [
+        (
+            10,
+            (
+                (0, 10),
+            )
+        ),
+        (
+            1000,
+            (
+                (0, 1000),
+            )
+        ),
+        (
+            1000000,
+            (
+                (0, 131072), (131072, 262144), (393216, 393216),
+                (786432, 213568)
+            )
+        ),
+        (
+            10000000,
+            (
+                (0, 131072), (131072, 262144), (393216, 393216),
+                (786432, 524288), (1310720, 655360), (1966080, 786432),
+                (2752512, 917504), (3670016, 1048576), (4718592, 1048576),
+                (5767168, 1048576), (6815744, 1048576), (7864320, 1048576),
+                (8912896, 1048576), (9961472, 38528)
+            )
+        ),
+    ]
+)
+def test_get_chunks(file_size, exp_result):
+    result = tuple(get_chunks(file_size))
+
+    assert result == exp_result

tox.ini

@@ -6,7 +6,7 @@ commands =
     flake8 {toxinidir}/src/
     coverage erase
     python setup.py install
-    pytest {toxinidir}/src/tests/tests.py
+    pytest {toxinidir}/src/tests/
 passenv =
     EMAIL
     PASS