Fixes download, updated download methods to return downloded path instead of None

This commit is contained in:
Richard O'Dwyer 2019-11-12 09:49:57 +00:00
parent c0916b77bb
commit 92a86ff439
3 changed files with 83 additions and 57 deletions

View file

@ -6,7 +6,9 @@ Release History
1.0.3 (unreleased) 1.0.3 (unreleased)
------------------ ------------------
- Nothing changed yet. - Fixes broken ``download`` method.
- Changes ``download`` and ``download_url`` methods to return the path to the downloaded file, previously returned ``None``.
- Added LICENSE.
1.0.2 (2019-11-07) 1.0.2 (2019-11-07)

View file

@ -568,7 +568,7 @@ class Mega:
""" """
Download a file by it's file object Download a file by it's file object
""" """
self._download_file( return self._download_file(
file_handle=None, file_handle=None,
file_key=None, file_key=None,
file=file[1], file=file[1],
@ -646,7 +646,7 @@ class Mega:
path = self._parse_url(url).split('!') path = self._parse_url(url).split('!')
file_id = path[0] file_id = path[0]
file_key = path[1] file_key = path[1]
self._download_file( return self._download_file(
file_handle=file_id, file_handle=file_id,
file_key=file_key, file_key=file_key,
dest_path=dest_path, dest_path=dest_path,
@ -716,12 +716,13 @@ class Mega:
else: else:
dest_path += '/' dest_path += '/'
temp_output_file = tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(
mode='w+b', prefix='megapy_', delete=False mode='w+b', prefix='megapy_', delete=False
) ) as temp_output_file:
k_str = a32_to_str(k) k_str = a32_to_str(k)
counter = Counter.new(128, initial_value=((iv[0] << 32) + iv[1]) << 64) counter = Counter.new(
128, initial_value=((iv[0] << 32) + iv[1]) << 64
)
aes = AES.new(k_str, AES.MODE_CTR, counter=counter) aes = AES.new(k_str, AES.MODE_CTR, counter=counter)
mac_str = '\0' * 16 mac_str = '\0' * 16
@ -746,21 +747,22 @@ class Mega:
block = chunk[i:i + 16] block = chunk[i:i + 16]
if len(block) % 16: if len(block) % 16:
block += '\0' * (16 - (len(block) % 16)) block += b'\0' * (16 - (len(block) % 16))
mac_str = mac_encryptor.encrypt(encryptor.encrypt(block)) mac_str = mac_encryptor.encrypt(encryptor.encrypt(block))
file_info = os.stat(temp_output_file.name) file_info = os.stat(temp_output_file.name)
logger.info('%s of %s downloaded', file_info.st_size, file_size) logger.info(
'%s of %s downloaded', file_info.st_size, file_size
)
file_mac = str_to_a32(mac_str) file_mac = str_to_a32(mac_str)
temp_output_file.close()
# check mac integrity # check mac integrity
if (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3]) != meta_mac: if (
file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3]
) != meta_mac:
raise ValueError('Mismatched mac') raise ValueError('Mismatched mac')
output_path = Path(dest_path + file_name)
shutil.move(temp_output_file.name, dest_path + file_name) shutil.move(temp_output_file.name, output_path)
return output_path
def upload(self, filename, dest=None, dest_filename=None): def upload(self, filename, dest=None, dest_filename=None):
# determine storage node # determine storage node

View file

@ -28,6 +28,17 @@ def mega(folder_name):
mega_.destroy(node_id) mega_.destroy(node_id)
@pytest.fixture
def uploaded_file(mega, folder_name):
folder = mega.find(folder_name)
dest_node_id = folder[1]['h']
mega.upload(
__file__, dest=dest_node_id, dest_filename='test.py'
)
path = f'{folder_name}/test.py'
return mega.find(path)
def test_mega(mega): def test_mega(mega):
assert isinstance(mega, Mega) assert isinstance(mega, Mega)
@ -56,10 +67,8 @@ def test_get_files(mega):
assert isinstance(files, dict) assert isinstance(files, dict)
def test_get_link(mega): def test_get_link(mega, uploaded_file):
file = mega.find(TEST_FILE) link = mega.get_link(uploaded_file)
if file:
link = mega.get_link(file)
assert isinstance(link, str) assert isinstance(link, str)
@ -171,20 +180,33 @@ def test_delete_folder(mega, folder_name):
assert isinstance(resp, int) assert isinstance(resp, int)
def test_delete(mega): def test_delete(mega, uploaded_file):
file = mega.find(TEST_FILE) resp = mega.delete(uploaded_file[0])
if file:
resp = mega.delete(file[0])
assert isinstance(resp, int) assert isinstance(resp, int)
def test_destroy(mega): def test_destroy(mega, uploaded_file):
file = mega.find(TEST_FILE) resp = mega.destroy(uploaded_file[0])
if file:
resp = mega.destroy(file[0])
assert isinstance(resp, int) assert isinstance(resp, int)
def test_download(mega, tmpdir, folder_name):
# Upload a single file into a folder
folder = mega.find(folder_name)
dest_node_id = folder[1]['h']
mega.upload(
__file__, dest=dest_node_id, dest_filename='test.py'
)
path = f'{folder_name}/test.py'
file = mega.find(path)
output_path = mega.download(
file=file, dest_path=tmpdir, dest_filename='test.py'
)
assert output_path.exists()
def test_empty_trash(mega): def test_empty_trash(mega):
# resp None if already empty, else int # resp None if already empty, else int
resp = mega.empty_trash() resp = mega.empty_trash()