Merge pull request #36 from gissehel/speedimprovement

mac check changes

mega/crypto.py

@@ -108,26 +108,14 @@ def a32_to_base64(a):
 
 
 def get_chunks(size):
-    chunks = {}
+    p = 0
-    p = pp = 0
+    s = 0x20000
-    i = 1
+    while p+s < size:
-
+        yield(p, s)
-    while i <= 8 and p < size - i * 0x20000:
+        p += s
-        chunks[p] = i * 0x20000
+        if s < 0x100000:
-        pp = p
+            s += 0x20000
-        p += chunks[p]
+    yield(p, size-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
 
 # more general functions
 def make_id(length):

mega/mega.py

@@ -461,41 +461,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)
+        mac_str = '\0' * 16
-        for chunk_start, chunk_size in sorted(get_chunks(file_size).items()):
+        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]]
+            encryptor = AES.new(k_str, AES.MODE_CBC, iv_str)
-            for i in range(0, len(chunk), 16):
+            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)
+            mac_str = mac_encryptor.encrypt(encryptor.encrypt(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)
 
             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
@@ -521,33 +517,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
-        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([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
+            encryptor = AES.new(k_str, AES.MODE_CBC, iv_str)
-            chunk_mac = [ul_key[4], ul_key[5], ul_key[4], ul_key[5]]
+            for i in range(0, len(chunk)-16, 16):
-            for i in range(0, len(chunk), 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)
+            mac_str = mac_encryptor.encrypt(encryptor.encrypt(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])
 
             #encrypt file and upload
             chunk = aes.encrypt(chunk)
@@ -559,6 +552,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])