voussoir/megapy
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Make shared_keys an instance attr, it seems that way already.

Browse source 
I noticed that _init_shared_keys assigns the result to the instance,
so let's just keep it on the instance from the beginning and not
pass it around. If some individual keys get overwritten during runtime
I expect that would be ok.
This commit is contained in:
Ethan Dalool 2020-03-23 20:22:35 -07:00
parent bb8a5081c5
commit 995dafceea
1 changed files with 17 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
src/mega/mega.py
View file

@ -43,6 +43,7 @@ class Mega:
self.sequence_num = random.randint(0, 0xFFFFFFFF) self.sequence_num = random.randint(0, 0xFFFFFFFF)
self.request_id = make_id(10) self.request_id = make_id(10)
self._trash_folder_node_id = None self._trash_folder_node_id = None
self.shared_keys = {}
if options is None: if options is None:
options = {} options = {}
@ -235,7 +236,7 @@ class Mega:
(public_handle, decryption_key) = match[0] (public_handle, decryption_key) = match[0]
return (public_handle, decryption_key) return (public_handle, decryption_key)
def _process_file(self, file, shared_keys): def _process_file(self, file):
if file['t'] in [NODE_TYPE_FILE, NODE_TYPE_DIR]: if file['t'] in [NODE_TYPE_FILE, NODE_TYPE_DIR]:
keys = dict( keys = dict(
keypart.split(':', 1) keypart.split(':', 1)
@ -253,19 +254,19 @@ class Mega:
base64_to_a32(file['sk']), self.master_key base64_to_a32(file['sk']), self.master_key
) )
key = decrypt_key(base64_to_a32(keys[file['h']]), shared_key) key = decrypt_key(base64_to_a32(keys[file['h']]), shared_key)
if file['su'] not in shared_keys: if file['su'] not in self.shared_keys:
shared_keys[file['su']] = {} self.shared_keys[file['su']] = {}
shared_keys[file['su']][file['h']] = shared_key self.shared_keys[file['su']][file['h']] = shared_key
# shared files # shared files
elif file['u'] and file['u'] in shared_keys: elif file['u'] and file['u'] in self.shared_keys:
for hkey in shared_keys[file['u']]: for hkey in self.shared_keys[file['u']]:
shared_key = shared_keys[file['u']][hkey] shared_key = self.shared_keys[file['u']][hkey]
if hkey in keys: if hkey in keys:
key = keys[hkey] key = keys[hkey]
key = decrypt_key(base64_to_a32(key), shared_key) key = decrypt_key(base64_to_a32(key), shared_key)
break break
if file['h'] and file['h'] in shared_keys.get('EXP', ()): if file['h'] and file['h'] in self.shared_keys.get('EXP', ()):
shared_key = shared_keys['EXP'][file['h']] shared_key = self.shared_keys['EXP'][file['h']]
encrypted_key = str_to_a32( encrypted_key = str_to_a32(
base64_url_decode(file['k'].split(':')[-1]) base64_url_decode(file['k'].split(':')[-1])
) )
@ -300,7 +301,7 @@ class Mega:
file['a'] = {'n': 'Rubbish Bin'} file['a'] = {'n': 'Rubbish Bin'}
return file return file
def _init_shared_keys(self, files, shared_keys): def _init_shared_keys(self, files):
""" """
Init shared key not associated with a user. Init shared key not associated with a user.
Seems to happen when a folder is shared, Seems to happen when a folder is shared,
@ -315,11 +316,10 @@ class Mega:
) )
ok_dict[ok_item['h']] = shared_key ok_dict[ok_item['h']] = shared_key
for s_item in files['s']: for s_item in files['s']:
if s_item['u'] not in shared_keys: if s_item['u'] not in self.shared_keys:
shared_keys[s_item['u']] = {} self.shared_keys[s_item['u']] = {}
if s_item['h'] in ok_dict: if s_item['h'] in ok_dict:
shared_keys[s_item['u']][s_item['h']] = ok_dict[s_item['h']] self.shared_keys[s_item['u']][s_item['h']] = ok_dict[s_item['h']]
self.shared_keys = shared_keys
def find_path_descriptor(self, path, files=()): def find_path_descriptor(self, path, files=()):
""" """
@ -393,10 +393,9 @@ class Mega:
logger.info('Getting all files...') logger.info('Getting all files...')
files = self._api_request({'a': 'f', 'c': 1, 'r': 1}) files = self._api_request({'a': 'f', 'c': 1, 'r': 1})
files_dict = {} files_dict = {}
shared_keys = {} self._init_shared_keys(files)
self._init_shared_keys(files, shared_keys)
for file in files['f']: for file in files['f']:
processed_file = self._process_file(file, shared_keys) processed_file = self._process_file(file)
# ensure each file has a name before returning # ensure each file has a name before returning
if processed_file['a']: if processed_file['a']:
files_dict[file['h']] = processed_file files_dict[file['h']] = processed_file
@ -504,8 +503,7 @@ class Mega:
files = self._api_request({'a': 'f', 'c': 1}) files = self._api_request({'a': 'f', 'c': 1})
# MERGE COMMON CODE WITH GET_FILES # MERGE COMMON CODE WITH GET_FILES
files_dict = {} files_dict = {}
shared_keys = {} self._init_shared_keys(files)
self._init_shared_keys(files, shared_keys)
for file in files['f']: for file in files['f']:
processed_file = self._process_file(file, shared_keys) processed_file = self._process_file(file, shared_keys)
if processed_file['a'] and processed_file['p'] == node_id: if processed_file['a'] and processed_file['p'] == node_id: