Merge pull request #36 from gissehel/speedimprovement
mac check changes
This commit is contained in:
Richard O'Dwyer
commit
3bd08be20e
2 changed files with 43 additions and 60 deletions
|
|
@ -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):
|
||||
|
|
|
|||
75
mega/mega.py
75
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)
|
||||
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
|
||||
|
|
@ -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
|
||||
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)
|
||||
|
|
@ -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])
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue