Big optimisation on download code for mac checking

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

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]
+                encryptor.encrypt(block)
+            i += 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)
+            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)