Naive TL message container implementation
This commit is contained in:
Anton Grigoryev
parent
79ea16f1dd
commit
77a192e0a8
20
TL.py
20
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,11 +140,17 @@ 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:
|
||||
# known types
|
||||
try:
|
||||
# Bare types
|
||||
tl_elem = tl.constructor_type[type_]
|
||||
except KeyError:
|
||||
# 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:
|
||||
|
||||
27
mtproto.py
27
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',
|
||||
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)
|
||||
data = z.getvalue()
|
||||
|
||||
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,13 +214,11 @@ class Session:
|
||||
|
||||
g_b_str = long_to_bytes(g_b)
|
||||
|
||||
z = io.BytesIO()
|
||||
TL.serialize_obj(z, 'client_DH_inner_data',
|
||||
data = TL.serialize_obj('client_DH_inner_data',
|
||||
nonce=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 = crypt.ige_encrypt(data_with_sha_padded, tmp_aes_key, tmp_aes_iv)
|
||||
|
||||
@ -23,3 +23,4 @@ Session = mtproto.Session(ip, port)
|
||||
Session.create_auth_key()
|
||||
|
||||
future_salts = Session.method_call('get_future_salts', num=3)
|
||||
print(future_salts)
|
||||
Loading…
Reference in New Issue
Block a user