diff --git a/HISTORY.rst b/HISTORY.rst index b7acbaf..9be4496 100644 --- a/HISTORY.rst +++ b/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) diff --git a/README.rst b/README.rst index 7af9c2e..3bf001c 100644 --- a/README.rst +++ b/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 ~~~~~~~~~~~~~ diff --git a/requirements.txt b/requirements.txt index 1bb8c2f..d743837 100644 --- a/requirements.txt +++ b/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 diff --git a/src/mega/mega.py b/src/mega/mega.py index d544c8e..f509f59 100644 --- a/src/mega/mega.py +++ b/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 = ( - (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 + 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) - return self._api_request(data=data) + msg = 'Request failed, retrying' + 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): - """ - Get all files in account - """ + logger.info('Getting all files...') 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( - self.schema, self.domain, public_handle, decrypted_key + return ( + 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( - self.schema, self.domain, public_handle, decrypted_key + return ( + 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( - self.schema, self.domain, public_handle, decrypted_key + return ( + 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, - None, + file_handle=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(( - '{0} of {1} downloaded'.format( - file_info.st_size, file_size - ) - )) + 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) @@ -768,104 +771,101 @@ class Mega(object): dest = self.root_id # request upload url, call 'u' method - input_file = open(filename, 'rb') - file_size = os.path.getsize(filename) - ul_url = self._api_request({'a': 'u', 's': file_size})['p'] + with open(filename, 'rb') as input_file: + file_size = os.path.getsize(filename) + ul_url = self._api_request({'a': 'u', 's': file_size})['p'] - # generate random aes key (128) for file - ul_key = [random.randint(0, 0xFFFFFFFF) for _ in range(6)] - k_str = a32_to_str(ul_key[:4]) - count = Counter.new( - 128, initial_value=((ul_key[4] << 32) + ul_key[5]) << 64 - ) - aes = AES.new(k_str, AES.MODE_CTR, counter=count) + # generate random aes key (128) for file + ul_key = [random.randint(0, 0xFFFFFFFF) for _ in range(6)] + k_str = a32_to_str(ul_key[:4]) + count = Counter.new( + 128, initial_value=((ul_key[4] << 32) + ul_key[5]) << 64 + ) + aes = AES.new(k_str, AES.MODE_CTR, counter=count) - upload_progress = 0 - completion_file_handle = None + upload_progress = 0 + completion_file_handle = None - mac_str = '\0' * 16 - 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]]) - if file_size > 0: - for chunk_start, chunk_size in get_chunks(file_size): - chunk = input_file.read(chunk_size) - upload_progress += len(chunk) + mac_str = '\0' * 16 + 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]]) + if file_size > 0: + for chunk_start, chunk_size in get_chunks(file_size): + chunk = input_file.read(chunk_size) + 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 - if file_size > 16: - i += 16 - else: - i = 0 - - block = chunk[i:i + 16] - if len(block) % 16: - block += makebyte('\0' * (16 - len(block) % 16)) - mac_str = mac_encryptor.encrypt(encryptor.encrypt(block)) - - # encrypt file and upload - chunk = aes.encrypt(chunk) + # encrypt file and upload + chunk = aes.encrypt(chunk) + output_file = requests.post( + ul_url + "/" + str(chunk_start), + data=chunk, + timeout=self.timeout + ) + completion_file_handle = output_file.text + logger.info( + '%s of %s uploaded', upload_progress, file_size + ) + else: output_file = requests.post( - ul_url + "/" + str(chunk_start), - data=chunk, - timeout=self.timeout + ul_url + "/0", data='', timeout=self.timeout ) completion_file_handle = output_file.text - if self.options.get('verbose') is True: - # 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) - 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 - 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])) - 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)) - # update attributes - data = self._api_request( - { - 'a': - 'p', - 't': - dest, - 'n': [ - { - 'h': completion_file_handle, - 't': 0, - 'a': encrypt_attribs, - 'k': encrypted_key - } - ] - } - ) - # close input file and return API msg - input_file.close() - return data + 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': 'p', + 't': dest, + 'i': self.request_id, + 'n': [ + { + 'h': completion_file_handle, + 't': 0, + 'a': encrypt_attribs, + 'k': encrypted_key + } + ] + } + ) + 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)) - - data = self._api_request( + return self._api_request( { - 'a': - 'p', - 't': - dest_node['h'], + 'a': 'p', + 't': dest_node['h'], 'n': [ { 'ph': file_handle, @@ -1099,4 +1091,3 @@ class Mega(object): ] } ) - return data diff --git a/src/tests/test_crypto.py b/src/tests/test_crypto.py new file mode 100644 index 0000000..9b8cc77 --- /dev/null +++ b/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 diff --git a/src/tests/tests.py b/src/tests/test_mega.py similarity index 100% rename from src/tests/tests.py rename to src/tests/test_mega.py diff --git a/tox.ini b/tox.ini index fd4b4c3..52c8da7 100644 --- a/tox.ini +++ b/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