Fixed tiny bugs with authentication, added more unit tests

crypto/aes.py

@@ -36,8 +36,9 @@ class AES:
     def encrypt_ige(plain_text, key, iv):
         """Encrypts the given text in 16-bytes blocks by using the given key and 32-bytes initialization vector"""
         # TODO: Random padding
-        padding = bytes(16 - len(plain_text) % 16)
+        if len(plain_text) % 16 != 0:  # Add padding if and only if it's not evenly divisible by 16 already
-        plain_text += padding
+            padding = bytes(16 - len(plain_text) % 16)
+            plain_text += padding
 
         iv1 = iv[:len(iv)//2]
         iv2 = iv[len(iv)//2:]

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())
-            crc = crc32(writer.get_bytes()[0:8 + len(self.body)]) & 0xFFFFFFFF
             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])
-            valid_checksum = crc32(body[:packet_len - 4]) & 0xFFFFFFFF
             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)
-        valid_checksum = crc32(rv) & 0xFFFFFFFF
 
         if checksum != valid_checksum:
             raise ValueError('Invalid checksum, skip')

tl/__init__.py

@@ -1,5 +1,9 @@
-from .all_tlobjects import tlobjects
+import os
-from .session import Session
+# Only import most stuff if the TLObjects were generated
-from .mtproto_request import MTProtoRequest
+if os.path.isfile('tl/all_tlobjects.py'):
-from .telegram_client import TelegramClient
+    from .all_tlobjects import tlobjects
+    from .session import Session
+    from .mtproto_request import MTProtoRequest
+    from .telegram_client import TelegramClient
+del os
 from .tlobject import TLObject, TLArg

tl/session.py

@@ -4,6 +4,7 @@ from os.path import isfile as file_exists
 import time
 import pickle
 import utils
+import random
 
 
 class Session:
@@ -38,9 +39,13 @@ class Session:
     def get_new_msg_id(self):
         """Generates a new message ID based on the current time (in ms) since epoch"""
         # Refer to mtproto_plain_sender.py for the original method, this is a simple copy
-        new_msg_id = int(self.time_offset + time.time() * 1000)
+        ms_time = int(time.time() * 1000)
-        if self.last_message_id >= new_msg_id:
+        new_msg_id = (((ms_time // 1000 + self.time_offset) << 32) |  # "must approximately equal unixtime*2^32"
-            new_msg_id = self._last_msg_id + 4
+                      ((ms_time % 1000) << 22) |  # "approximate moment in time the message was created"
+                      random.randint(0, 524288) << 2)  # "message identifiers are divisible by 4"
 
-        self._last_msg_id = new_msg_id
+        if self.last_message_id >= new_msg_id:
+            new_msg_id = self.last_message_id + 4
+
+        self.last_message_id = new_msg_id
         return new_msg_id

tl_generator.py

@@ -2,8 +2,8 @@ import os
 import re
 import shutil
 
-from parser import TLParser
+from parser.tl_parser import TLParser
-from parser import SourceBuilder
+from parser.source_builder import SourceBuilder
 
 
 def tlobjects_exist():

unit_test.py

@@ -5,16 +5,16 @@ import unittest
 import os
 import platform
 
-import utils
+from tl import Session
-import network.authenticator
+from tl.functions import InvokeWithLayerRequest, InitConnectionRequest
+from tl.functions.help import GetConfigRequest
 
 from crypto import AES, Factorizator
 from network import TcpTransport, TcpClient, MtProtoSender
 from utils import BinaryWriter, BinaryReader
 
-from tl import Session
+import utils
-from tl.functions import InvokeWithLayerRequest, InitConnectionRequest
+import network.authenticator
-from tl.functions.help import GetConfigRequest
 
 host = 'localhost'
 port = random.randint(50000, 60000)  # Arbitrary non-privileged port
@@ -219,6 +219,50 @@ class UnitTest(unittest.TestCase):
         assert  cipher_text == real, 'Decrypted text does not equal the real value (expected "{}", got "{}")'\
             .format(get_representation(real), get_representation(cipher_text))
 
+    @staticmethod
+    def test_calc_key():
+        shared_key = get_bytes('BC-D2-6D-B7-CA-76-F4-5D-5B-88-83-27-20-F3-11-8A-73-D0-34-94-31-AE-2A-4F-03-86-9A-2F-48-23-1A-8C-B5-6A-E9-24-E0-49-76-43-6D-5E-E7-30-1A-35-43-09-16-03-D2-9D-A9-89-D6-CE-08-50-0F-64-72-A0-B3-EB-FE-63-76-1A-DF-4A-14-96-98-16-A3-47-AB-04-14-21-5C-EB-0A-BC-6E-DF-C4-25-C6-09-B7-16-14-9C-27-81-15-3D-B0-AF-0E-0B-52-AA-04-36-36-73-F0-CF-B7-B8-3E-2C-44-94-78-D7-F8-E0-84-CB-25-D3-05-B2-E8-95-4D-72-3F-A2-E8-49-6E-F9-0B-5B-45-9B-AA-0C-58-7F-0E-69-DE-EE-64-1D-78-2F-4A-CE-EA-5E-7D-30-3B-A8-33-42-BB-52-A1-BF-65-04-B9-1E-A1-22-66-3D-A5-4D-40-9E-DD-81-80-C9-A5-FB-FC-67-DD-15-03-70-21-0F-66-44-16-89-32-EA-CA-B1-41-99-4F-A9-34-50-A9-A2-C6-3B-B2-43-39-1D-43-35-D2-0D-EC-4C-D9-AB-77-2D-03-0D-79-C2-76-17-5D-02-15-0C-42-61-97-CE-A5-B1-E4-5D-8E-E0-2C-CF-43-7B-6F-FA-99-66-A4-70-4D-00')
+        msg_keys = [
+            'BA-1A-CF-DA-A8-5E-43-62-6C-FA-B6-0C-3A-9B-B0-FC',
+            '86-6D-92-69-CF-8B-93-AA-86-4B-1F-69-D0-34-83-5D',
+            'A1-2F-C0-61-6F-60-1A-B0-33-8F-7D-27-08-C8-EA-15',
+            '3A-56-52-0F-B7-89-D7-80-F6-18-72-CD-09-B5-A8-8A',
+            '06-F7-84-F3-91-CC-8D-DC-7D-92-41-7A-7E-84-25-E4',
+            '78-4D-FC-AE-F4-C4-55-81-6D-DD-99-A7-DB-B8-A3-88'
+        ]
+        are_client = [
+            True,
+            False,
+            True,
+            True,
+            False,
+            False
+        ]
+
+        valid_keys = [
+            'AF-E3-84-51-6D-E0-21-0C-D9-31-E4-9A-A0-76-5F-67-63-78-A1-B0-C9-BC-16-27-76-CF-2C-9D-4D-AE-C6-A5',
+            'DD-30-58-B6-93-8E-C9-79-EF-83-F8-8C-6A-A7-68-03-E2-C6-B1-36-C5-BB-FC-E7-DF-D6-B1-67-F7-75-CF-6B',
+            '70-3F-66-AF-FE-0A-F3-1E-95-C3-25-48-8D-0F-A7-95-59-53-BF-DD-35-97-6E-7A-C0-5E-79-9C-9D-09-3B-7B',
+            '20-9F-82-E3-95-3F-9D-1E-EE-3E-F3-82-B0-8D-6E-76-26-5B-94-27-DD-7D-61-C3-AC-EB-FA-71-FF-0D-8F-08',
+            'AE-06-BF-F1-89-88-22-66-98-48-76-E6-BD-D9-39-36-2D-36-4E-CF-FD-47-D2-87-D7-49-F8-93-22-2E-66-02',
+            'D9-6B-43-D0-89-F7-C5-75-A2-4A-F2-2F-F0-17-5C-95-AE-FA-46-A5-95-AA-C3-B9-76-B0-3A-A8-0E-7B-EA-5D'
+        ]
+
+        valid_ivs = [
+            'B8-51-F3-C5-A3-5D-C6-DF-9E-E0-51-BD-22-8D-13-09-0E-9A-9D-5E-38-A2-F8-E7-00-77-D9-C1-A7-A0-F7-0F',
+            'DC-4C-C2-18-01-4A-22-58-86-6C-62-B6-B5-34-37-FD-E2-61-34-B6-AF-7D-46-53-D7-5B-E0-4E-0D-19-FB-BC',
+            '68-BB-BA-7F-55-B9-EF-86-EE-20-5A-1A-45-4E-70-C3-48-56-A2-E9-2F-91-AD-74-23-FE-54-06-E5-68-04-E9',
+            '79-31-8D-F0-DC-31-60-D7-09-BC-66-F1-AB-0D-7C-CB-2A-AF-74-32-64-C5-B3-18-C4-ED-55-D9-F6-39-DD-F3',
+            '1F-51-66-06-05-54-B9-E0-52-6A-88-6A-70-0C-DA-8D-2B-CD-BF-8A-67-5E-1A-A7-DD-EA-1C-CE-4C-D4-34-3D',
+            '8E-00-97-5E-B7-A1-F7-3D-1C-16-03-CA-B3-ED-EA-80-64-8F-77-A6-C4-34-5B-B5-DC-5D-C9-EC-B7-F8-F4-76'
+        ]
+
+        for msg_key, is_client, valid_key, valid_iv in zip(msg_keys, are_client, valid_keys, valid_ivs):
+            msg_key = get_bytes(msg_key)
+            key, iv = utils.calc_key(shared_key, msg_key, is_client)
+            assert get_representation(key) == valid_key
+            assert get_representation(iv) == valid_iv
+
     @staticmethod
     def test_authenticator():
         transport = TcpTransport('149.154.167.91', 443)

utils/__init__.py

@@ -1,3 +1,3 @@
-from .binary_writer import BinaryWriter
 from .binary_reader import BinaryReader
+from .binary_writer import BinaryWriter
 from .helpers import *

utils/binary_reader.py

@@ -84,6 +84,16 @@ class BinaryReader:
         """Reads a Telegram-encoded string"""
         return str(self.tgread_bytes(), encoding='utf-8')
 
+    def tgread_bool(self):
+        """Reads a Telegram boolean value"""
+        value = self.read_int(signed=False)
+        if value == 0x997275b5:  # boolTrue
+            return True
+        elif value == 0xbc799737:  # boolFalse
+            return False
+        else:
+            raise ValueError('Invalid boolean code {}'.format(hex(value)))
+
     def tgread_object(self):
         """Reads a Telegram object"""
         constructor_id = self.read_int()