From 77a192e0a830ca14f2205570240ff42e7b37c1b5 Mon Sep 17 00:00:00 2001
From: Anton Grigoryev <grig.anton@gmail.com>
Date: Tue, 24 Mar 2015 08:49:27 +0300
Subject: [PATCH] Naive TL message container implementation

---
 TL.py      | 28 ++++++++++++++++++++++------
 mtproto.py | 53 ++++++++++++++++++++++-------------------------------
 testing.py |  3 ++-
 3 files changed, 46 insertions(+), 38 deletions(-)

diff --git a/TL.py b/TL.py
index 385764d8..090c0639 100644
--- a/TL.py
+++ b/TL.py
@@ -16,6 +16,12 @@ class TlConstructor:
             if param['type'] == "Vector<long>":
                 param['type'] = "Vector t"
                 param['subtype'] = "long"
+            elif param['type'] == "vector<%Message>":
+                param['type'] = "vector"
+                param['subtype'] = "message"
+            elif param['type'] == "vector<future_salt>":
+                param['type'] = "vector"
+                param['subtype'] = "future_salt"
             else:
                 param['subtype'] = None
             self.params.append(param)
@@ -56,7 +62,8 @@ class TL:
 tl = TL(os.path.join(os.path.dirname(__file__), "TL_schema.JSON"))
 
 
-def serialize_obj(bytes_io, type_, **kwargs):
+def serialize_obj(type_, **kwargs):
+    bytes_io = io.BytesIO()
     try:
         tl_constructor = tl.constructor_type[type_]
     except KeyError:
@@ -64,9 +71,11 @@ def serialize_obj(bytes_io, type_, **kwargs):
     bytes_io.write(struct.pack('<i', tl_constructor.id))
     for arg in tl_constructor.params:
         serialize_param(bytes_io, type_=arg['type'],  value=kwargs[arg['name']])
+    return bytes_io.getvalue()
 
 
-def serialize_method(bytes_io, type_, **kwargs):
+def serialize_method(type_, **kwargs):
+    bytes_io = io.BytesIO()
     try:
         tl_method = tl.method_name[type_]
     except KeyError:
@@ -74,6 +83,7 @@ def serialize_method(bytes_io, type_, **kwargs):
     bytes_io.write(struct.pack('<i', tl_method.id))
     for arg in tl_method.params:
         serialize_param(bytes_io, type_=arg['type'], value=kwargs[arg['name']])
+    return bytes_io.getvalue()
 
 
 def serialize_param(bytes_io, type_, value):
@@ -130,12 +140,18 @@ def deserialize(bytes_io, type_=None, subtype=None):
         count = struct.unpack('<l', bytes_io.read(4))[0]
         x = [deserialize(bytes_io, type_=subtype) for i in range(count)]
     else:
-        # Boxed types
-        i = struct.unpack('<i', bytes_io.read(4))[0]  # read type ID
+        # known types
         try:
-            tl_elem = tl.constructor_id[i]
+        # Bare types
+            tl_elem = tl.constructor_type[type_]
         except KeyError:
-            raise Exception("Could not extract type: %s" % type_)
+            # Boxed types
+            i = struct.unpack('<i', bytes_io.read(4))[0]  # read type ID
+            try:
+                tl_elem = tl.constructor_id[i]
+            except KeyError:
+                # Unknown type
+                raise Exception("Could not extract type: %s" % type_)
         base_boxed_types = ["Vector t", "Int", "Long", "Double", "String", "Int128", "Int256"]
         if tl_elem.type in base_boxed_types:
             x = deserialize(bytes_io, type_=tl_elem.predicate, subtype=subtype)
diff --git a/mtproto.py b/mtproto.py
index 3fb02ac2..ff9e40b7 100644
--- a/mtproto.py
+++ b/mtproto.py
@@ -128,12 +128,8 @@ class Session:
             raise Exception("Got unknown auth_key id")
         return data
 
-
     def method_call(self, method, **kwargs):
-        z = io.BytesIO()
-        TL.serialize_method(z, method, **kwargs)
-        z_val = z.getvalue()
-        self.send_message(z_val)
+        self.send_message(TL.serialize_method(method, **kwargs))
         server_answer = self.recv_message()
         return TL.deserialize(io.BytesIO(server_answer))
 
@@ -161,18 +157,14 @@ class Session:
         f = open(os.path.join(os.path.dirname(__file__), "rsa.pub"))
         key = RSA.importKey(f.read())
 
-        z = io.BytesIO()
-
         new_nonce = os.urandom(32)
-
-        TL.serialize_obj(z, 'p_q_inner_data',
-                         pq=pq_bytes,
-                         p=p_bytes,
-                         q=q_bytes,
-                         nonce=nonce,
-                         server_nonce=server_nonce,
-                         new_nonce=new_nonce)
-        data = z.getvalue()
+        data = TL.serialize_obj('p_q_inner_data',
+                                pq=pq_bytes,
+                                p=p_bytes,
+                                q=q_bytes,
+                                nonce=nonce,
+                                server_nonce=server_nonce,
+                                new_nonce=new_nonce)
 
         sha_digest = SHA.new(data).digest()
         random_bytes = os.urandom(255-len(data)-len(sha_digest))
@@ -180,17 +172,17 @@ class Session:
         encrypted_data = key.encrypt(to_encrypt, 0)[0]
 
         print("Starting Diffie Hellman key exchange")
-        server_DH_params = self.method_call('req_DH_params',
-                                            nonce=nonce, # 16 bytes
+        server_dh_params = self.method_call('req_DH_params',
+                                            nonce=nonce,
                                             server_nonce=server_nonce,
                                             p=p_bytes,
                                             q=q_bytes,
                                             public_key_fingerprint=public_key_fingerprint,
                                             encrypted_data=encrypted_data)
-        assert nonce == server_DH_params['nonce']
-        assert server_nonce == server_DH_params['server_nonce']
+        assert nonce == server_dh_params['nonce']
+        assert server_nonce == server_dh_params['server_nonce']
 
-        encrypted_answer = server_DH_params['encrypted_answer']
+        encrypted_answer = server_dh_params['encrypted_answer']
 
         tmp_aes_key = SHA.new(new_nonce + server_nonce).digest() + SHA.new(server_nonce + new_nonce).digest()[0:12]
         tmp_aes_iv = SHA.new(server_nonce + new_nonce).digest()[12:20] + SHA.new(new_nonce + new_nonce).digest() + new_nonce[0:4]
@@ -213,6 +205,7 @@ class Session:
 
         dh_prime = bytes_to_long(dh_prime_str)
         g_a = bytes_to_long(g_a_str)
+
         assert prime.isprime(dh_prime)
         retry_id = 0
         b_str = os.urandom(256)
@@ -221,21 +214,19 @@ class Session:
 
         g_b_str = long_to_bytes(g_b)
 
-        z = io.BytesIO()
-        TL.serialize_obj(z, 'client_DH_inner_data',
-                              nonce=nonce,
-                              server_nonce=server_nonce,
-                              retry_id=retry_id,
-                              g_b=g_b_str)
-        data = z.getvalue()
+        data = TL.serialize_obj('client_DH_inner_data',
+                                nonce=nonce,
+                                server_nonce=server_nonce,
+                                retry_id=retry_id,
+                                g_b=g_b_str)
         data_with_sha = SHA.new(data).digest()+data
         data_with_sha_padded = data_with_sha + os.urandom(-len(data_with_sha) % 16)
         encrypted_data = crypt.ige_encrypt(data_with_sha_padded, tmp_aes_key, tmp_aes_iv)
 
         Set_client_DH_params_answer = self.method_call('set_client_DH_params',
-                                                           nonce=nonce,
-                                                           server_nonce=server_nonce,
-                                                           encrypted_data=encrypted_data)
+                                                       nonce=nonce,
+                                                       server_nonce=server_nonce,
+                                                       encrypted_data=encrypted_data)
 
         auth_key = pow(g_a, b, dh_prime)
         auth_key_str = long_to_bytes(auth_key)
diff --git a/testing.py b/testing.py
index 125d1dd6..707dfcea 100644
--- a/testing.py
+++ b/testing.py
@@ -22,4 +22,5 @@ port = config.getint('App data', 'port')
 Session = mtproto.Session(ip, port)
 Session.create_auth_key()
 
-future_salts = Session.method_call('get_future_salts', num=3)
\ No newline at end of file
+future_salts = Session.method_call('get_future_salts', num=3)
+print(future_salts)
\ No newline at end of file