Merge pull request #3 from richardARPANET/misc-improvements

Misc improvements
This commit is contained in:
Richard 2019-11-06 23:16:01 +00:00 committed by GitHub
commit 5888ce13fc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 176 additions and 142 deletions

View file

@ -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)

View file

@ -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
~~~~~~~~~~~~~

View file

@ -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

View file

@ -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

42
src/tests/test_crypto.py Normal file
View file

@ -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

View file

@ -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