From c453a88b28256527c52cb4fa5129ed0b2ecd41bf Mon Sep 17 00:00:00 2001 From: gissehel Date: Sun, 5 May 2013 19:48:23 +0200 Subject: [PATCH 1/2] Big optimisation on download code for mac checking --- mega/crypto.py | 28 ++++++++-------------------- mega/mega.py | 40 ++++++++++++++++++---------------------- 2 files changed, 26 insertions(+), 42 deletions(-) diff --git a/mega/crypto.py b/mega/crypto.py index 8a91356..c439971 100644 --- a/mega/crypto.py +++ b/mega/crypto.py @@ -108,26 +108,14 @@ def a32_to_base64(a): def get_chunks(size): - chunks = {} - p = pp = 0 - i = 1 - - while i <= 8 and p < size - i * 0x20000: - chunks[p] = i * 0x20000 - pp = p - p += chunks[p] - i += 1 - - while p < size: - chunks[p] = 0x100000 - pp = p - p += chunks[p] - - chunks[pp] = size - pp - if not chunks[pp]: - del chunks[pp] - - return chunks + p = 0 + s = 0x20000 + while p+s < size: + yield(p, s) + p += s + if s < 0x100000: + s += 0x20000 + yield(p, size-p) # more general functions def make_id(length): diff --git a/mega/mega.py b/mega/mega.py index d6ac841..24ef1ae 100644 --- a/mega/mega.py +++ b/mega/mega.py @@ -447,41 +447,37 @@ class Mega(object): temp_output_file = tempfile.NamedTemporaryFile(mode='w+b', prefix='megapy_', delete=False) + k_str = a32_to_str(k) counter = Counter.new( 128, initial_value=((iv[0] << 32) + iv[1]) << 64) - aes = AES.new(a32_to_str(k), AES.MODE_CTR, counter=counter) + aes = AES.new(k_str, AES.MODE_CTR, counter=counter) - file_mac = (0, 0, 0, 0) - for chunk_start, chunk_size in sorted(get_chunks(file_size).items()): + mac_str = '\0' * 16 + mac_encryptor = AES.new(k_str, AES.MODE_CBC, mac_str) + iv_str = a32_to_str([iv[0], iv[1], iv[0], iv[1]]) + + for chunk_start, chunk_size in get_chunks(file_size): chunk = input_file.read(chunk_size) chunk = aes.decrypt(chunk) temp_output_file.write(chunk) - chunk_mac = [iv[0], iv[1], iv[0], iv[1]] - for i in range(0, len(chunk), 16): + encryptor = AES.new(k_str, AES.MODE_CBC, iv_str) + for i in range(0, len(chunk)-16, 16): block = chunk[i:i + 16] - if len(block) % 16: - block += '\0' * (16 - (len(block) % 16)) - block = str_to_a32(block) - chunk_mac = [ - chunk_mac[0] ^ block[0], - chunk_mac[1] ^ block[1], - chunk_mac[2] ^ block[2], - chunk_mac[3] ^ block[3]] - chunk_mac = aes_cbc_encrypt_a32(chunk_mac, k) - - file_mac = [ - file_mac[0] ^ chunk_mac[0], - file_mac[1] ^ chunk_mac[1], - file_mac[2] ^ chunk_mac[2], - file_mac[3] ^ chunk_mac[3]] - file_mac = aes_cbc_encrypt_a32(file_mac, k) + encryptor.encrypt(block) + i += 16 + block = chunk[i:i + 16] + if len(block) % 16: + 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_mac = str_to_a32(mac_str) + temp_output_file.close() # check mac integrity @@ -513,7 +509,7 @@ class Mega(object): file_mac = [0, 0, 0, 0] upload_progress = 0 completion_file_handle = None - for chunk_start, chunk_size in sorted(get_chunks(file_size).items()): + for chunk_start, chunk_size in get_chunks(file_size): chunk = input_file.read(chunk_size) upload_progress += len(chunk) From c2befcb87ff4d7775270e1a007087df93c28ff0a Mon Sep 17 00:00:00 2001 From: gissehel Date: Sun, 5 May 2013 20:29:41 +0200 Subject: [PATCH 2/2] Speedup the mac calculation on upload --- mega/mega.py | 35 +++++++++++++++++------------------ 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/mega/mega.py b/mega/mega.py index 24ef1ae..adb148b 100644 --- a/mega/mega.py +++ b/mega/mega.py @@ -503,33 +503,30 @@ class Mega(object): #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(a32_to_str(ul_key[:4]), AES.MODE_CTR, counter=count) + aes = AES.new(k_str, AES.MODE_CTR, counter=count) - file_mac = [0, 0, 0, 0] 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]]) + for chunk_start, chunk_size in get_chunks(file_size): chunk = input_file.read(chunk_size) upload_progress += len(chunk) - #determine chunks mac - chunk_mac = [ul_key[4], ul_key[5], ul_key[4], ul_key[5]] - for i in range(0, len(chunk), 16): + encryptor = AES.new(k_str, AES.MODE_CBC, iv_str) + for i in range(0, len(chunk)-16, 16): block = chunk[i:i + 16] - if len(block) % 16: - block += '\0' * (16 - len(block) % 16) - block = str_to_a32(block) - chunk_mac = [chunk_mac[0] ^ block[0], chunk_mac[1] ^ block[1], - chunk_mac[2] ^ block[2], - chunk_mac[3] ^ block[3]] - chunk_mac = aes_cbc_encrypt_a32(chunk_mac, ul_key[:4]) - - #our files mac - file_mac = [file_mac[0] ^ chunk_mac[0], file_mac[1] ^ chunk_mac[1], - file_mac[2] ^ chunk_mac[2], - file_mac[3] ^ chunk_mac[3]] - file_mac = aes_cbc_encrypt_a32(file_mac, ul_key[:4]) + encryptor.encrypt(block) + i += 16 + block = chunk[i:i + 16] + if len(block) % 16: + block += '\0' * (16 - len(block) % 16) + mac_str = mac_encryptor.encrypt(encryptor.encrypt(block)) #encrypt file and upload chunk = aes.encrypt(chunk) @@ -541,6 +538,8 @@ class Mega(object): # upload progress print('{0} of {1} uploaded'.format(upload_progress, file_size)) + file_mac = str_to_a32(mac_str) + #determine meta mac meta_mac = (file_mac[0] ^ file_mac[1], file_mac[2] ^ file_mac[3])