Merge pull request #3 from richardARPANET/misc-improvements
Misc improvements
This commit is contained in:
commit
5888ce13fc
7 changed files with 176 additions and 142 deletions
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
~~~~~~~~~~~~~
|
||||
|
|
|
@ -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
|
||||
|
|
263
src/mega/mega.py
263
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
|
||||
|
|
42
src/tests/test_crypto.py
Normal file
42
src/tests/test_crypto.py
Normal 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
|
2
tox.ini
2
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
|
||||
|
|
Loading…
Reference in a new issue