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,15 +121,9 @@ 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 = int(
-                    l = (
+                    ((private_key[0]) * 256 + (private_key[1]) + 7) / 8
-                        (ord(private_key[0]) * 256 +
+                ) + 2
-                         ord(private_key[1]) + 7) / 8
-                    ) + 2
-                else:
-                    l = int(
-                        ((private_key[0]) * 256 + (private_key[1]) + 7) / 8
-                    ) + 2
                 self.rsa_private_key[i] = mpi_to_int(private_key[:l])
                 private_key = private_key[l:]
 
@@ -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:
+            file_info = os.stat(temp_output_file.name)
-                # temp file size
+            logger.info('%s of %s downloaded', file_info.st_size, file_size)
-                file_info = os.stat(temp_output_file.name)
-                print((
-                    '{0} of {1} downloaded'.format(
-                        file_info.st_size, file_size
-                    )
-                ))
 
         file_mac = str_to_a32(mac_str)
 
@@ -768,104 +771,101 @@ 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)
+            file_size = os.path.getsize(filename)
-        ul_url = self._api_request({'a': 'u', 's': file_size})['p']
+            ul_url = self._api_request({'a': 'u', 's': file_size})['p']
 
-        # generate random aes key (128) for file
+            # generate random aes key (128) for file
-        ul_key = [random.randint(0, 0xFFFFFFFF) for _ in range(6)]
+            ul_key = [random.randint(0, 0xFFFFFFFF) for _ in range(6)]
-        k_str = a32_to_str(ul_key[:4])
+            k_str = a32_to_str(ul_key[:4])
-        count = Counter.new(
+            count = Counter.new(
-            128, initial_value=((ul_key[4] << 32) + ul_key[5]) << 64
+                128, initial_value=((ul_key[4] << 32) + ul_key[5]) << 64
-        )
+            )
-        aes = AES.new(k_str, AES.MODE_CTR, counter=count)
+            aes = AES.new(k_str, AES.MODE_CTR, counter=count)
 
-        upload_progress = 0
+            upload_progress = 0
-        completion_file_handle = None
+            completion_file_handle = None
 
-        mac_str = '\0' * 16
+            mac_str = '\0' * 16
-        mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str)
+            mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str)
-        iv_str = a32_to_str([ul_key[4], ul_key[5], ul_key[4], ul_key[5]])
+            iv_str = a32_to_str([ul_key[4], ul_key[5], ul_key[4], ul_key[5]])
-        if file_size > 0:
+            if file_size > 0:
-            for chunk_start, chunk_size in get_chunks(file_size):
+                for chunk_start, chunk_size in get_chunks(file_size):
-                chunk = input_file.read(chunk_size)
+                    chunk = input_file.read(chunk_size)
-                upload_progress += len(chunk)
+                    upload_progress += len(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
 
-                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)
+                    if len(block) % 16:
+                        block += makebyte('\0' * (16 - len(block) % 16))
+                    mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
 
-                # fix for files under 16 bytes failing
+                    # encrypt file and upload
-                if file_size > 16:
+                    chunk = aes.encrypt(chunk)
-                    i += 16
+                    output_file = requests.post(
-                else:
+                        ul_url + "/" + str(chunk_start),
-                    i = 0
+                        data=chunk,
-
+                        timeout=self.timeout
-                block = chunk[i:i + 16]
+                    )
-                if len(block) % 16:
+                    completion_file_handle = output_file.text
-                    block += makebyte('\0' * (16 - len(block) % 16))
+                    logger.info(
-                mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
+                        '%s of %s uploaded', upload_progress, file_size
-
+                    )
-                # encrypt file and upload
+            else:
-                chunk = aes.encrypt(chunk)
                 output_file = requests.post(
-                    ul_url + "/" + str(chunk_start),
+                    ul_url + "/0", data='', timeout=self.timeout
-                    data=chunk,
-                    timeout=self.timeout
                 )
                 completion_file_handle = output_file.text
 
-                if self.options.get('verbose') is True:
+            logger.info('Chunks uploaded')
-                    # upload progress
+            logger.info('Setting attributes to complete upload')
-                    print((
+            logger.info('Computing attributes')
-                        '{0} of {1} uploaded'.format(
+            file_mac = str_to_a32(mac_str)
-                            upload_progress, file_size
-                        )
-                    ))
-        else:
-            output_file = requests.post(
-                ul_url + "/0", data='', timeout=self.timeout
-            )
-            completion_file_handle = output_file.text
 
-        file_mac = str_to_a32(mac_str)
+            # determine meta mac
+            meta_mac = (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3])
 
-        # determine meta mac
+            dest_filename = dest_filename or os.path.basename(filename)
-        meta_mac = (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3])
-
-        if dest_filename is not None:
             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(
-        key = [
+                encrypt_attr(attribs, ul_key[:4])
-            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],
+            key = [
-            ul_key[5], meta_mac[0], meta_mac[1]
+                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],
-        encrypted_key = a32_to_base64(encrypt_key(key, self.master_key))
+                ul_key[5], meta_mac[0], meta_mac[1]
-        # update attributes
+            ]
-        data = self._api_request(
+            encrypted_key = a32_to_base64(encrypt_key(key, self.master_key))
-            {
+            logger.info('Sending request to update attributes')
-                'a':
+            # update attributes
-                'p',
+            data = self._api_request(
-                't':
+                {
-                dest,
+                    'a': 'p',
-                'n': [
+                    't': dest,
-                    {
+                    'i': self.request_id,
-                        'h': completion_file_handle,
+                    'n': [
-                        't': 0,
+                        {
-                        'a': encrypt_attribs,
+                            'h': completion_file_handle,
-                        'k': encrypted_key
+                            't': 0,
-                    }
+                            'a': encrypt_attribs,
-                ]
+                            'k': encrypted_key
-            }
+                        }
-        )
+                    ]
-        # close input file and return API msg
+                }
-        input_file.close()
+            )
-        return data
+            logger.info('Upload complete')
+            return data
 
     def _mkdir(self, name, parent_node_id):
         # generate random aes key (128) for folder
@@ -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