Reviewed authenticator.py

utils/auth_key.py

@@ -1,6 +1,6 @@
 # This file is based on TLSharp
 # https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/MTProto/Crypto/AuthKey.cs
-from hashlib import sha1
+import utils.helpers as utils
 from utils.binary_writer import BinaryWriter
 from utils.binary_reader import BinaryReader
 
@@ -8,23 +8,22 @@ from utils.binary_reader import BinaryReader
 class AuthKey:
     def __init__(self, gab=None, data=None):
         if gab:
-            self.key = gab.to_byte_array_unsigned()
+            self.key = utils.get_byte_array(gab, signed=False)
         elif data:
             self.key = data
         else:
             raise AssertionError('Either a gab integer or data bytes array must be provided')
 
-        with BinaryReader(sha1(self.key)) as reader:
+        with BinaryReader(utils.sha1(self.key)) as reader:
             self.aux_hash = reader.read_long(signed=False)
             reader.read(4)
             self.key_id = reader.read_long(signed=False)
 
-
     def calc_new_nonce_hash(self, new_nonce, number):
         with BinaryWriter() as writer:
             writer.write(new_nonce)
             writer.write_byte(number)
             writer.write_long(self.aux_hash, signed=False)
 
-            new_nonce_hash = sha1(writer.get_bytes())[4:20]
+            new_nonce_hash = utils.sha1(writer.get_bytes())[4:20]
             return new_nonce_hash

utils/helpers.py

@@ -1,11 +1,14 @@
 import os
 from utils.binary_writer import BinaryWriter
-from hashlib import sha1
+import hashlib
+
+
+bits_per_byte = 8
 
 
 def generate_random_long(signed=True):
     """Generates a random long integer (8 bytes), which is optionally signed"""
-    return int.from_bytes(os.urandom(8), signed=signed)
+    return int.from_bytes(os.urandom(8), signed=signed, byteorder='big')
 
 
 def generate_random_bytes(count):
@@ -13,27 +16,20 @@ def generate_random_bytes(count):
     return os.urandom(count)
 
 
+def get_byte_array(integer, signed):
+    bits = integer.bit_length()
+    byte_length = (bits + bits_per_byte - 1) // bits_per_byte
+    return int.to_bytes(integer, length=byte_length, byteorder='big', signed=signed)
+
+
 def calc_key(shared_key, msg_key, client):
     """Calculate the key based on Telegram guidelines, specifying whether it's the client or not"""
     x = 0 if client else 8
 
-    buffer = [0] * 48
-    buffer[0:16] = msg_key
-    buffer[16:48] = shared_key[x:x + 32]
-    sha1a = sha1(buffer)
-
-    buffer[0:16] = shared_key[x + 32:x + 48]
-    buffer[16:32] = msg_key
-    buffer[32:48] = shared_key[x + 48:x + 64]
-    sha1b = sha1(buffer)
-
-    buffer[0:32] = shared_key[x + 64:x + 96]
-    buffer[32:48] = msg_key
-    sha1c = sha1(buffer)
-
-    buffer[0:16] = msg_key
-    buffer[16:48] = shared_key[x + 96:x + 128]
-    sha1d = sha1(buffer)
+    sha1a = sha1(msg_key + shared_key[x:x + 32])
+    sha1b = sha1(shared_key[x + 32:x + 48] + msg_key + shared_key[x + 48:x + 64])
+    sha1c = sha1(shared_key[x + 64:x + 96] + msg_key)
+    sha1d = sha1(msg_key + shared_key[x + 96:x + 128])
 
     key = sha1a[0:8] + sha1b[8:20] + sha1c[4:16]
     iv = sha1a[8:20] + sha1b[0:8] + sha1c[16:20] + sha1d[0:8]
@@ -48,43 +44,27 @@ def calc_msg_key(data):
 
 def calc_msg_key_offset(data, offset, limit):
     """Calculates the message key from offset given data, with an optional offset and limit"""
-    # TODO untested, may not be offset like this
-    # In the original code it was as parameters for the sha function, not slicing the array
     return sha1(data[offset:offset + limit])[4:20]
 
 
-def generate_key_data_from_nonces(serverNonce, newNonce):
-    # TODO unsure that this works
-    nonces = [0] * 48
+def generate_key_data_from_nonces(server_nonce, new_nonce):
+    hash1 = sha1(bytes(new_nonce + server_nonce))
+    hash2 = sha1(bytes(server_nonce + new_nonce))
+    hash3 = sha1(bytes(new_nonce + new_nonce))
 
-    nonces[00:32] = newNonce
-    nonces[32:48] = serverNonce
-    hash1 = hash(bytes(nonces))
+    with BinaryWriter() as key_buffer:
+        with BinaryWriter() as iv_buffer:
+            key_buffer.write(hash1)
+            key_buffer.write(hash2[:12])
 
-    nonces[00:16] = serverNonce
-    nonces[16:32] = newNonce
-    hash2 = hash(bytes(nonces))
+            iv_buffer.write(hash2[12:20])
+            iv_buffer.write(hash3)
+            iv_buffer.write_byte(new_nonce[:4])
 
-    nonces = [0] * 64
-    nonces[00:32] = newNonce
-    nonces[32:64] = newNonce
-    hash2 = hash(bytes(nonces))
+            return key_buffer.get_bytes(), iv_buffer.get_bytes()
 
-    with BinaryWriter() as keyBuffer:
-        with BinaryWriter() as ivBuffer:
-            """
-                using (var keyBuffer = new MemoryStream(32))
-                using (var ivBuffer = new MemoryStream(32))
-                {
-                    keyBuffer.Write(hash1, 0, hash1.Length);
-                    keyBuffer.Write(hash2, 0, 12);
 
-                    ivBuffer.Write(hash2, 12, 8);
-                    ivBuffer.Write(hash3, 0, hash3.Length);
-                    ivBuffer.Write(newNonce, 0, 4);
-
-                    return new AESKeyData(keyBuffer.ToArray(), ivBuffer.ToArray());
-                }
-            """
-            # TODO implement
-            raise NotImplementedError()
+def sha1(data):
+    sha = hashlib.sha1
+    sha.update(data)
+    return sha.digest()