From 66a5577ab978969018af86e89d8e78e223dfd63b Mon Sep 17 00:00:00 2001
From: Anton Grigoryev <mail@griganton.ru>
Date: Thu, 12 Mar 2015 03:17:43 +0300
Subject: [PATCH] Some refactor

now testing.py looks almost just like i want.
---
 mtproto.py | 148 +++++++++++++++++++++++++++++------------------------
 testing.py |  60 +---------------------
 2 files changed, 81 insertions(+), 127 deletions(-)

diff --git a/mtproto.py b/mtproto.py
index c9facee2..6daab00e 100644
--- a/mtproto.py
+++ b/mtproto.py
@@ -49,8 +49,8 @@ class TlMethod:
 
 
 class TL:
-    def __init__(self):
-        with open("TL_schema.JSON", 'r') as f:
+    def __init__(self, filename):
+        with open(filename, 'r') as f:
            TL_dict = json.load(f)
 
         # Read constructors
@@ -71,79 +71,75 @@ class TL:
             self.method_id[z.id] = z
             self.method_name[z.method] = z
 
-    def method_call(self, method, **kwargs):
-        z = io.BytesIO()
-        tl_method = self.method_name[method]
-        z.write(struct.pack('<i', tl_method.id))
-        for param in tl_method.params:
-            self.serialize(bytes_io=z, value=kwargs[param['name']], type_=param['type'])
-        return z.getvalue()
 
-    def serialize(self, bytes_io, type_, value):
-        if  type_ == "int":
-            assert isinstance(value, int)
-            bytes_io.write(struct.pack('<i', value))
-        if  type_ == "int128":
-            assert isinstance(value, bytes)
-            bytes_io.write(struct.pack('<16s', value))
+## Loading TL_schema
+tl = TL("TL_schema.JSON")
 
-    def deserialize(self, bytes_io, type_=None, subtype=None):
-        assert isinstance(bytes_io, io.BytesIO)
+def serialize(bytes_io, type_, value):
+    if  type_ == "int":
+        assert isinstance(value, int)
+        bytes_io.write(struct.pack('<i', value))
+    if  type_ == "int128":
+        assert isinstance(value, bytes)
+        bytes_io.write(struct.pack('<16s', value))
 
-        # Built-in bare types
-        if type_ == 'int':
-            x = struct.unpack('<i', bytes_io.read(4))[0]
-        elif type_ == '#':
-            x = struct.unpack('<I', bytes_io.read(4))[0]
-        elif type_ == 'long':
-            x = struct.unpack('<q', bytes_io.read(8))[0]
-        elif type_ == 'double':
-            x = struct.unpack('<d', bytes_io.read(8))[0]
-        elif type_ == 'int128':
-            t = struct.unpack('<16s', bytes_io.read(16))[0]
-            x = int.from_bytes(t, 'little')
-        elif type_ == 'int256':
-            t = struct.unpack('<32s', bytes_io.read(32))[0]
-            x = int.from_bytes(t, 'little')
-        elif type_ == 'bytes':
-            l = int.from_bytes(bytes_io.read(1), 'little')
+def deserialize(bytes_io, type_=None, subtype=None):
+    assert isinstance(bytes_io, io.BytesIO)
+
+    # Built-in bare types
+    if type_ == 'int':
+        x = struct.unpack('<i', bytes_io.read(4))[0]
+    elif type_ == '#':
+        x = struct.unpack('<I', bytes_io.read(4))[0]
+    elif type_ == 'long':
+        x = struct.unpack('<q', bytes_io.read(8))[0]
+    elif type_ == 'double':
+        x = struct.unpack('<d', bytes_io.read(8))[0]
+    elif type_ == 'int128':
+        t = struct.unpack('<16s', bytes_io.read(16))[0]
+        x = int.from_bytes(t, 'little')
+    elif type_ == 'int256':
+        t = struct.unpack('<32s', bytes_io.read(32))[0]
+        x = int.from_bytes(t, 'little')
+    elif type_ == 'bytes':
+        l = int.from_bytes(bytes_io.read(1), 'little')
+        x = bytes_io.read(l)
+        bytes_io.read(-(l+1) % 4)  # skip padding bytes
+    elif type_ == 'string':
+        l = int.from_bytes(bytes_io.read(1), 'little')
+        assert l <=254
+        if l == 254:
+            # We have a long string
+            long_len = int.from_bytes(bytes_io.read(3), 'little')
+            x = bytes_io.read(long_len)
+            bytes_io.read(-long_len % 4)  # skip padding bytes
+        else:
+            # We have a short string
             x = bytes_io.read(l)
             bytes_io.read(-(l+1) % 4)  # skip padding bytes
-        elif type_ == 'string':
-            l = int.from_bytes(bytes_io.read(1), 'little')
-            assert l <=254
-            if l == 254:
-                # We have a long string
-                long_len = int.from_bytes(bytes_io.read(3), 'little')
-                x = bytes_io.read(long_len)
-                bytes_io.read(-long_len % 4)  # skip padding bytes
-            else:
-                # We have a short string
-                x = bytes_io.read(l)
-                bytes_io.read(-(l+1) % 4)  # skip padding bytes
-            assert isinstance(x, bytes)
-        elif type_ == 'vector':
-            assert subtype is not None
-            count = int.from_bytes(bytes_io.read(4), 'little')
-            x = [self.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
-            try:
-                tl_elem = self.constructor_id[i]
-            except:
-                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 = self.deserialize(bytes_io, type_=tl_elem.predicate, subtype=subtype)
-            else:  # other types
-                x = {}
-                for arg in tl_elem.params:
-                    x[arg['name']] = self.deserialize(bytes_io, type_=arg['type'], subtype=arg['subtype'])
-        return x
-
+        assert isinstance(x, bytes)
+    elif type_ == 'vector':
+        assert subtype is not None
+        count = int.from_bytes(bytes_io.read(4), 'little')
+        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
+        try:
+            tl_elem = tl.constructor_id[i]
+        except:
+            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)
+        else:  # other types
+            x = {}
+            for arg in tl_elem.params:
+                x[arg['name']] = deserialize(bytes_io, type_=arg['type'], subtype=arg['subtype'])
+    return x
 
 class Session:
+    """ Manages TCP Transport. encryption and message frames """
     def __init__(self, ip, port):
         # creating socket
         self.sock = socket.socket()
@@ -211,3 +207,19 @@ class Session:
                 print('<<')
                 vis(data)  # Received message visualisation to console
                 return data
+
+    def method_call(self, method, **kwargs):
+        z = io.BytesIO()
+        tl_method = tl.method_name[method]
+        z.write(struct.pack('<i', tl_method.id))
+        for param in tl_method.params:
+            serialize(bytes_io=z, value=kwargs[param['name']], type_=param['type'])
+        self.send_message(z.getvalue())
+        server_answer = self.recv_message()
+        return deserialize(io.BytesIO(server_answer))
+
+class Telegram:
+    global tl
+    def connect(self, ip, port):
+        self.session = Session(ip, port)
+
diff --git a/testing.py b/testing.py
index b87adbb1..eabd20ec 100644
--- a/testing.py
+++ b/testing.py
@@ -2,7 +2,6 @@
 import mtproto
 import os
 import prime
-import io
 import configparser
 
 config = configparser.ConfigParser()
@@ -10,67 +9,10 @@ config.read('credentials')
 ip = config['App data']['ip_address']
 port = config['App data'].getint('port')
 
-TL=mtproto.TL()
-
 Session = mtproto.Session(ip, port)
-a = TL.method_call('req_pq', nonce=os.urandom(16))
-Session.send_message(a)
-b = Session.recv_message()
-x = TL.deserialize(io.BytesIO(b))
+x = Session.method_call('req_pq', nonce=os.urandom(16))
 
 PQ = int.from_bytes(x['pq'], 'big')
 [p, q] = prime.primefactors(PQ)
 
 print("PQ = %d\np = %d, q = %d" % (PQ, p, q))
-#sock = socket.socket()
-# sock.connect(("149.154.167.40", 443))
-# nonce = os.urandom(16)
-# tosend = b'\x78\x97\x46\x60' + nonce
-# good=(b'\x00\x00\x00\x00\x00\x00\x00\x00'
-#      #+b'\x00\x00\x00\x00\x0f\x35\xe8\x51
-#      +int(time.time()*2**32).to_bytes(8, 'little')
-#      +len(tosend).to_bytes(4, 'little')
-#      +tosend)
-#      #+b'\x14\x00\x00\x00'
-#      #+b'\x78\x97\x46\x60'
-#      #+b'\x15\x3e\xa9\xbe\xc5\x4f\xfc\x16'
-#      #+b'\x66\x23\x4a\xd0\x31\xc0\xf1\x3d')
-# mtproto.sendpacket(sock, good, 0)
-# (number, data) = mtproto.recvpacket(sock)
-# auth_key_id = data[0:8]
-# message_id =  data[8:16]
-# message_length = int.from_bytes(data[16:20], 'little')
-# resPQ = data[20:24]
-# nonce2 = data[24:40]
-# server_nonce = data[40:56]
-# pq = data[56:68]
-# fingerprint_count = int.from_bytes(data[72:76], 'little')
-# fingerprints = []
-# for i in range(fingerprint_count):
-#     fingerprints.append(data[76+i*8:76+8*(i+1)])
-# PQ = int.from_bytes(pq[1:9],'big')
-# server_public_key = fingerprints[0]
-# [p, q] = prime.primefactors(PQ)
-# #(p, q)=factorize(PQ)
-# def num_to_string(number, endianness):
-#     number_bytes = math.ceil(number.bit_length() / 8)
-#     return (number_bytes.to_bytes(3, endianness) +
-#             number.to_bytes(number_bytes, endianness) +
-#             b'\x00' * ((number_bytes + 3) %4) )
-# new_nonce = os.urandom(32)
-# data = (b'\x83\xc9\x5a\xec' +  #(p_q_inner_data)
-#         num_to_string(PQ, 'big') +
-#         num_to_string(p, 'big') +
-#         num_to_string(q, 'big') +
-#         nonce +
-#         server_nonce +
-#         new_nonce )
-# sock.close()
-# #34 00 00 00 - len
-# #00 00 00 00- num in the session
-# #00 00 00 00 00 00 00 00
-# #00 00 00 00 0f 35 e8 51 - message id (unixtime << 32)
-# #14 00 00 00 - message len
-# #78 97 46 60 - function
-# #15 3e a9 be c5 4f fc 16 66 23 4a d0 31 c0 f1 3d - nonce (random number 16 bytes)
-# #dd fe f0 07 - crc32