From 0d79c88636f7e68bdd87916fa9b067b26ea8f336 Mon Sep 17 00:00:00 2001
From: Anton Grigoryev <mail@griganton.ru>
Date: Tue, 17 Mar 2015 01:04:23 +0300
Subject: [PATCH] Auth successful

Python 2 is not supported since to_bytes function in testing.py
---
 crypt.py   |  2 +-
 mtproto.py |  5 ++--
 testing.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++-------
 3 files changed, 72 insertions(+), 14 deletions(-)

diff --git a/crypt.py b/crypt.py
index f69ff43df..26200ccc5 100644
--- a/crypt.py
+++ b/crypt.py
@@ -40,12 +40,12 @@ class IGE:
          key must be 32 byte
          iv must be 32 byte (it's not internally used in AES 256 ECB, but it's
          needed for IGE)"""
+        blocksize = self.cipher.block_size
 
         if len(message) % blocksize != 0:
             raise ValueError("message must be a multiple of 16 bytes (try adding " +
                             str(16 - len(message) % 16) + " bytes of padding)")
 
-        blocksize = self.cipher.block_size
         ivp = self.iv[0:blocksize]
         ivp2 = self.iv[blocksize:]
 
diff --git a/mtproto.py b/mtproto.py
index 231032406..9774d6540 100644
--- a/mtproto.py
+++ b/mtproto.py
@@ -104,7 +104,6 @@ def serialize_method(bytes_io, type_, **kwargs):
 
 
 def serialize_param(bytes_io, type_, value):
-    print("type(value): " + str(type(value)))
     if type_ == "int":
         assert isinstance(value, int)
         bytes_io.write(struct.pack('<i', value))
@@ -176,11 +175,11 @@ def deserialize(bytes_io, type_=None, subtype=None):
 
 class Session:
     """ Manages TCP Transport. encryption and message frames """
-    def __init__(self, ip, port):
+    def __init__(self, ip, port, auth_key_id=None):
         # creating socket
         self.sock = socket.socket()
         self.sock.connect((ip, port))
-        self.auth_key_id = None
+        self.auth_key_id = auth_key_id
         self.number = 0
 
     def __del__(self):
diff --git a/testing.py b/testing.py
index 960df9341..4d1c9a2f9 100644
--- a/testing.py
+++ b/testing.py
@@ -9,6 +9,7 @@ except ImportError:
     import ConfigParser as configparser
 from Crypto.Hash import SHA
 from Crypto.PublicKey import RSA
+from Crypto.Util.strxor import strxor
 
 # local modules
 import crypt
@@ -30,6 +31,7 @@ x = Session.method_call('req_pq', nonce=client_nonce)
 
 server_nonce = x['server_nonce']
 public_key_fingerprint = x['server_public_key_fingerprints'][0]
+# TODO: selecting RSA public key based on this fingerprint
 PQ_bytes = x['pq']
 
 # doing len(PQ_bytes) I saw it was 8 bytes, so we unpack with Q
@@ -65,26 +67,21 @@ random_bytes = os.urandom(255-len(data)-len(sha_digest))
 to_encrypt = sha_digest + data + random_bytes
 encrypted_data = key.encrypt(to_encrypt, 0)[0]
 
-z = Session.method_call('req_DH_params',
+server_DH_params = Session.method_call('req_DH_params',
                         nonce=client_nonce, # 16 bytes
                         server_nonce=server_nonce,
                         p=P_bytes,
                         q=Q_bytes,
                         public_key_fingerprint=public_key_fingerprint,
                         encrypted_data=encrypted_data)
+assert client_nonce == server_DH_params['nonce']
+assert server_nonce == server_DH_params['server_nonce']
+
+encrypted_answer = server_DH_params['encrypted_answer']
 
-encrypted_answer = z['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]
 
-print("\ntmp_aes_key:")
-mtproto.vis(tmp_aes_key)
-print(tmp_aes_key.__repr__())
-
-print("\ntmp_aes_iv:")
-mtproto.vis(tmp_aes_iv)
-print(tmp_aes_iv.__repr__())
-
 crypter = crypt.IGE(tmp_aes_key, tmp_aes_iv)
 
 answer_with_hash = crypter.decrypt(encrypted_answer)
@@ -93,3 +90,65 @@ answer_hash = answer_with_hash[:20]
 answer = answer_with_hash[20:]
 mtproto.vis(answer) # To start off BA0D89 ...
 
+server_DH_inner_data = mtproto.deserialize(io.BytesIO(answer))
+assert client_nonce == server_DH_inner_data['nonce']
+assert server_nonce == server_DH_inner_data['server_nonce']
+dh_prime_str = server_DH_inner_data['dh_prime']
+g = server_DH_inner_data['g']
+g_a_str = server_DH_inner_data['g_a']
+server_time = server_DH_inner_data['server_time']
+
+
+dh_prime = int.from_bytes(dh_prime_str,'big')
+g_a = int.from_bytes(g_a_str,'big')
+print(dh_prime)
+print(g)
+print(g_a)
+print(prime.isprime(dh_prime))
+
+b_str = os.urandom(256)
+b = int.from_bytes(b_str,'big')
+g_b = pow(g,b,dh_prime)
+
+g_b_str = int.to_bytes(g_b, g_b.bit_length() // 8 + 1, 'big')
+
+
+retry_id = 0
+z = io.BytesIO()
+mtproto.serialize_obj(z, 'client_DH_inner_data',
+                      nonce=client_nonce,
+                      server_nonce=server_nonce,
+                      retry_id=retry_id,
+                      g_b=g_b_str)
+data = z.getvalue()
+data_with_sha = SHA.new(data).digest()+data
+data_with_sha_padded = data_with_sha + os.urandom(-len(data_with_sha) % 16)
+encrypted_data = crypter.encrypt(data_with_sha_padded)
+
+Set_client_DH_params_answer = Session.method_call('set_client_DH_params',
+                                                   nonce=client_nonce,
+                                                   server_nonce=server_nonce,
+                                                   encrypted_data=encrypted_data)
+print(encrypted_data)
+print(Set_client_DH_params_answer)
+
+auth_key = pow(g_a, b, dh_prime)
+auth_key_str = int.to_bytes(auth_key, auth_key.bit_length() // 8 + 1, 'big')
+print("auth_key = %d" % auth_key)
+
+auth_key_sha = SHA.new(auth_key_str).digest()
+auth_key_hash = auth_key_sha[-8:]
+auth_key_aux_hash = auth_key_sha[:8]
+
+new_nonce_hash1 = SHA.new(new_nonce+b'\x01'+auth_key_aux_hash).digest()[-16:]
+new_nonce_hash2 = SHA.new(new_nonce+b'\x02'+auth_key_aux_hash).digest()[-16:]
+new_nonce_hash3 = SHA.new(new_nonce+b'\x03'+auth_key_aux_hash).digest()[-16:]
+
+assert Set_client_DH_params_answer['nonce'] == client_nonce
+assert Set_client_DH_params_answer['server_nonce'] == server_nonce
+assert Set_client_DH_params_answer['new_nonce_hash1'] == new_nonce_hash1
+
+Session.__del__()
+
+server_salt = strxor(new_nonce[0:8], server_nonce[0:8])
+mtproto.vis(server_salt)
\ No newline at end of file