Fixed tiny bugs with authentication, added more unit tests

network/mtproto_plain_sender.py

@@ -1,6 +1,7 @@
 # This file is based on TLSharp
 # https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/MtProtoPlainSender.cs
 import time
+import random
 from utils import BinaryWriter, BinaryReader
 
 
@@ -36,7 +37,11 @@ class MtProtoPlainSender:
 
     def get_new_msg_id(self):
         """Generates a new message ID based on the current time (in ms) since epoch"""
-        new_msg_id = int(self._time_offset + time.time() * 1000)  # Multiply by 1000 to get milliseconds
+        # See https://core.telegram.org/mtproto/description#message-identifier-msg-id
+        ms_time = int(time.time() * 1000)
+        new_msg_id = (((ms_time // 1000) << 32) |  # "must approximately equal unixtime*2^32"
+                      ((ms_time % 1000) << 22) |  # "approximate moment in time the message was created"
+                      random.randint(0, 524288) << 2)  # "message identifiers are divisible by 4"
 
         # Ensure that we always return a message ID which is higher than the previous one
         if self._last_msg_id >= new_msg_id:

network/mtproto_sender.py

@@ -100,7 +100,7 @@ class MtProtoSender:
             remote_auth_key_id = reader.read_long()
             msg_key = reader.read(16)
 
-            key, iv = utils.calc_key(self.session.auth_key.data, msg_key, False)
+            key, iv = utils.calc_key(self.session.auth_key.key, msg_key, False)
             plain_text = AES.decrypt_ige(reader.read(len(body) - reader.tell_position()), key, iv)
 
             with BinaryReader(plain_text) as plain_text_reader:
@@ -262,6 +262,8 @@ class MtProtoSender:
 
         if request_id == mtproto_request.msg_id:
             mtproto_request.confirm_received = True
+        else:
+            print('We did not get the ID we expected. Got {}, but it should have been {}'.format(request_id, mtproto_request.msg_id))
 
         inner_code = reader.read_int(signed=False)
         if inner_code == 0x2144ca19:  # RPC Error

network/tcp_message.py

@@ -1,8 +1,7 @@
 # This file is based on TLSharp
 # https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpMessage.cs
-from zlib import crc32
-
 from utils import BinaryWriter, BinaryReader
+from binascii import crc32
 
 
 class TcpMessage:
@@ -31,10 +30,8 @@ class TcpMessage:
             writer.write_int(len(self.body) + 12)
             writer.write_int(self.sequence_number)
             writer.write(self.body)
-            writer.flush()  # Flush so we can get the buffer in the CRC
 
-            # Ensure it's unsigned (see http://stackoverflow.com/a/30092291/4759433)
-            crc = crc32(writer.get_bytes()[0:8 + len(self.body)]) & 0xFFFFFFFF
+            crc = crc32(writer.get_bytes())
             writer.write_int(crc, signed=False)
 
             return writer.get_bytes()
@@ -55,10 +52,9 @@ class TcpMessage:
 
             seq = reader.read_int()
             packet = reader.read(packet_len - 12)
-            checksum = reader.read_int()
+            checksum = reader.read_int(signed=False)
 
-            # Ensure it's unsigned (see http://stackoverflow.com/a/30092291/4759433)
-            valid_checksum = crc32(body[:packet_len - 4]) & 0xFFFFFFFF
+            valid_checksum = crc32(body[:packet_len - 4])
             if checksum != valid_checksum:
                 raise ValueError('Invalid checksum, skip')
 

network/tcp_transport.py

@@ -1,7 +1,7 @@
 # This file is based on TLSharp
 # https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpTransport.cs
-from zlib import crc32
 from network import TcpMessage, TcpClient
+from binascii import crc32
 
 
 class TcpTransport:
@@ -35,12 +35,11 @@ class TcpTransport:
 
         body = self._tcp_client.read(packet_length - 12)
 
-        checksum = int.from_bytes(self._tcp_client.read(4), byteorder='little')
+        checksum = int.from_bytes(self._tcp_client.read(4), byteorder='little', signed=False)
 
         # Then perform the checks
         rv = packet_length_bytes + seq_bytes + body
-        # Ensure it's unsigned (http://stackoverflow.com/a/30092291/4759433)
-        valid_checksum = crc32(rv) & 0xFFFFFFFF
+        valid_checksum = crc32(rv)
 
         if checksum != valid_checksum:
             raise ValueError('Invalid checksum, skip')