Many code-style improvements

This commit is contained in:
Fadi Hadzh
2016-11-30 00:29:42 +03:00
parent ef264ae83f
commit d087941bd0
25 changed files with 698 additions and 499 deletions

View File

@@ -1,9 +1,10 @@
import os
import time
import telethon.helpers as utils
from telethon.utils import BinaryWriter, BinaryReader
from telethon.crypto import AES, AuthKey, Factorizator, RSA
from telethon.crypto import AES, RSA, AuthKey, Factorizator
from telethon.network import MtProtoPlainSender
from telethon.utils import BinaryReader, BinaryWriter
def do_authentication(transport):
@@ -23,7 +24,8 @@ def do_authentication(transport):
with BinaryReader(sender.receive()) as reader:
response_code = reader.read_int(signed=False)
if response_code != 0x05162463:
raise AssertionError('Invalid response code: {}'.format(hex(response_code)))
raise AssertionError('Invalid response code: {}'.format(
hex(response_code)))
nonce_from_server = reader.read(16)
if nonce_from_server != nonce:
@@ -36,7 +38,8 @@ def do_authentication(transport):
vector_id = reader.read_int()
if vector_id != 0x1cb5c415:
raise AssertionError('Invalid vector constructor ID: {}'.format(hex(response_code)))
raise AssertionError('Invalid vector constructor ID: {}'.format(
hex(response_code)))
fingerprints = []
fingerprint_count = reader.read_int()
@@ -47,32 +50,46 @@ def do_authentication(transport):
new_nonce = os.urandom(32)
p, q = Factorizator.factorize(pq)
with BinaryWriter() as pq_inner_data_writer:
pq_inner_data_writer.write_int(0x83c95aec, signed=False) # PQ Inner Data
pq_inner_data_writer.write_int(
0x83c95aec, signed=False) # PQ Inner Data
pq_inner_data_writer.tgwrite_bytes(get_byte_array(pq, signed=False))
pq_inner_data_writer.tgwrite_bytes(get_byte_array(min(p, q), signed=False))
pq_inner_data_writer.tgwrite_bytes(get_byte_array(max(p, q), signed=False))
pq_inner_data_writer.tgwrite_bytes(
get_byte_array(
min(p, q), signed=False))
pq_inner_data_writer.tgwrite_bytes(
get_byte_array(
max(p, q), signed=False))
pq_inner_data_writer.write(nonce)
pq_inner_data_writer.write(server_nonce)
pq_inner_data_writer.write(new_nonce)
cipher_text, target_fingerprint = None, None
for fingerprint in fingerprints:
cipher_text = RSA.encrypt(get_fingerprint_text(fingerprint), pq_inner_data_writer.get_bytes())
cipher_text = RSA.encrypt(
get_fingerprint_text(fingerprint),
pq_inner_data_writer.get_bytes())
if cipher_text is not None:
target_fingerprint = fingerprint
break
if cipher_text is None:
raise AssertionError('Could not find a valid key for fingerprints: {}'
.format(', '.join([get_fingerprint_text(f) for f in fingerprints])))
raise AssertionError(
'Could not find a valid key for fingerprints: {}'
.format(', '.join([get_fingerprint_text(f)
for f in fingerprints])))
with BinaryWriter() as req_dh_params_writer:
req_dh_params_writer.write_int(0xd712e4be, signed=False) # Req DH Params
req_dh_params_writer.write_int(
0xd712e4be, signed=False) # Req DH Params
req_dh_params_writer.write(nonce)
req_dh_params_writer.write(server_nonce)
req_dh_params_writer.tgwrite_bytes(get_byte_array(min(p, q), signed=False))
req_dh_params_writer.tgwrite_bytes(get_byte_array(max(p, q), signed=False))
req_dh_params_writer.tgwrite_bytes(
get_byte_array(
min(p, q), signed=False))
req_dh_params_writer.tgwrite_bytes(
get_byte_array(
max(p, q), signed=False))
req_dh_params_writer.write(target_fingerprint)
req_dh_params_writer.tgwrite_bytes(cipher_text)
@@ -88,7 +105,8 @@ def do_authentication(transport):
raise AssertionError('Server DH params fail: TODO')
if response_code != 0xd0e8075c:
raise AssertionError('Invalid response code: {}'.format(hex(response_code)))
raise AssertionError('Invalid response code: {}'.format(
hex(response_code)))
nonce_from_server = reader.read(16)
if nonce_from_server != nonce:
@@ -106,7 +124,6 @@ def do_authentication(transport):
g, dh_prime, ga, time_offset = None, None, None, None
with BinaryReader(plain_text_answer) as dh_inner_data_reader:
hashsum = dh_inner_data_reader.read(20)
code = dh_inner_data_reader.read_int(signed=False)
if code != 0xb5890dba:
raise AssertionError('Invalid DH Inner Data code: {}'.format(code))
@@ -132,26 +149,34 @@ def do_authentication(transport):
# Prepare client DH Inner Data
with BinaryWriter() as client_dh_inner_data_writer:
client_dh_inner_data_writer.write_int(0x6643b654, signed=False) # Client DH Inner Data
client_dh_inner_data_writer.write_int(
0x6643b654, signed=False) # Client DH Inner Data
client_dh_inner_data_writer.write(nonce)
client_dh_inner_data_writer.write(server_nonce)
client_dh_inner_data_writer.write_long(0) # TODO retry_id
client_dh_inner_data_writer.tgwrite_bytes(get_byte_array(gb, signed=False))
client_dh_inner_data_writer.tgwrite_bytes(
get_byte_array(
gb, signed=False))
with BinaryWriter() as client_dh_inner_data_with_hash_writer:
client_dh_inner_data_with_hash_writer.write(utils.sha1(client_dh_inner_data_writer.get_bytes()))
client_dh_inner_data_with_hash_writer.write(client_dh_inner_data_writer.get_bytes())
client_dh_inner_data_bytes = client_dh_inner_data_with_hash_writer.get_bytes()
client_dh_inner_data_with_hash_writer.write(
utils.sha1(client_dh_inner_data_writer.get_bytes()))
client_dh_inner_data_with_hash_writer.write(
client_dh_inner_data_writer.get_bytes())
client_dh_inner_data_bytes = client_dh_inner_data_with_hash_writer.get_bytes(
)
# Encryption
client_dh_inner_data_encrypted_bytes = AES.encrypt_ige(client_dh_inner_data_bytes, key, iv)
client_dh_inner_data_encrypted_bytes = AES.encrypt_ige(
client_dh_inner_data_bytes, key, iv)
# Prepare Set client DH params
with BinaryWriter() as set_client_dh_params_writer:
set_client_dh_params_writer.write_int(0xf5045f1f, signed=False)
set_client_dh_params_writer.write(nonce)
set_client_dh_params_writer.write(server_nonce)
set_client_dh_params_writer.tgwrite_bytes(client_dh_inner_data_encrypted_bytes)
set_client_dh_params_writer.tgwrite_bytes(
client_dh_inner_data_encrypted_bytes)
set_client_dh_params_bytes = set_client_dh_params_writer.get_bytes()
sender.send(set_client_dh_params_bytes)
@@ -171,7 +196,8 @@ def do_authentication(transport):
new_nonce_hash1 = reader.read(16)
auth_key = AuthKey(get_byte_array(gab, signed=False))
new_nonce_hash_calculated = auth_key.calc_new_nonce_hash(new_nonce, 1)
new_nonce_hash_calculated = auth_key.calc_new_nonce_hash(new_nonce,
1)
if new_nonce_hash1 != new_nonce_hash_calculated:
raise AssertionError('Invalid new nonce hash')
@@ -200,7 +226,8 @@ def get_byte_array(integer, signed):
"""Gets the arbitrary-length byte array corresponding to the given integer"""
bits = integer.bit_length()
byte_length = (bits + 8 - 1) // 8 # 8 bits per byte
return int.to_bytes(integer, length=byte_length, byteorder='big', signed=signed)
return int.to_bytes(
integer, length=byte_length, byteorder='big', signed=signed)
def get_int(byte_array, signed=True):

View File

@@ -1,10 +1,12 @@
import time
import random
from telethon.utils import BinaryWriter, BinaryReader
import time
from telethon.utils import BinaryReader, BinaryWriter
class MtProtoPlainSender:
"""MTProto Mobile Protocol plain sender (https://core.telegram.org/mtproto/description#unencrypted-messages)"""
def __init__(self, transport):
self._sequence = 0
self._time_offset = 0
@@ -37,9 +39,12 @@ class MtProtoPlainSender:
"""Generates a new message ID based on the current time (in ms) since epoch"""
# 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"
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:

View File

@@ -1,18 +1,19 @@
import gzip
from telethon.errors import *
from time import sleep
from datetime import timedelta
from threading import Thread, RLock
from threading import RLock, Thread
from time import sleep
import telethon.helpers as utils
from telethon.crypto import AES
from telethon.utils import BinaryWriter, BinaryReader
from telethon.tl.types import MsgsAck
from telethon.errors import *
from telethon.tl.all_tlobjects import tlobjects
from telethon.tl.types import MsgsAck
from telethon.utils import BinaryReader, BinaryWriter
class MtProtoSender:
"""MTProto Mobile Protocol sender (https://core.telegram.org/mtproto/description)"""
def __init__(self, transport, session):
self.transport = transport
self.session = session
@@ -27,7 +28,8 @@ class MtProtoSender:
# We need this to avoid using the updates thread if we're waiting to read
self.waiting_receive = False
self.updates_thread = Thread(target=self.updates_thread_method, name='Updates thread')
self.updates_thread = Thread(
target=self.updates_thread_method, name='Updates thread')
self.updates_thread_running = False
self.updates_thread_receiving = False
@@ -118,7 +120,8 @@ class MtProtoSender:
message, remote_msg_id, remote_sequence = self.decode_msg(body)
with BinaryReader(message) as reader:
self.process_msg(remote_msg_id, remote_sequence, reader, request)
self.process_msg(remote_msg_id, remote_sequence, reader,
request)
# We can now set the flag to False thus resuming the updates thread
self.waiting_receive = False
@@ -148,7 +151,8 @@ class MtProtoSender:
# And then finally send the encrypted packet
with BinaryWriter() as cipher_writer:
cipher_writer.write_long(self.session.auth_key.key_id, signed=False)
cipher_writer.write_long(
self.session.auth_key.key_id, signed=False)
cipher_writer.write(msg_key)
cipher_writer.write(cipher_text)
self.transport.send(cipher_writer.get_bytes())
@@ -168,7 +172,8 @@ class MtProtoSender:
msg_key = reader.read(16)
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)
plain_text = AES.decrypt_ige(
reader.read(len(body) - reader.tell_position()), key, iv)
with BinaryReader(plain_text) as plain_text_reader:
remote_salt = plain_text_reader.read_long()
@@ -198,7 +203,8 @@ class MtProtoSender:
if code == 0x3072cfa1: # gzip_packed
return self.handle_gzip_packed(msg_id, sequence, reader, request)
if code == 0xedab447b: # bad_server_salt
return self.handle_bad_server_salt(msg_id, sequence, reader, request)
return self.handle_bad_server_salt(msg_id, sequence, reader,
request)
if code == 0xa7eff811: # bad_msg_notification
return self.handle_bad_msg_notification(msg_id, sequence, reader)
@@ -253,7 +259,8 @@ class MtProtoSender:
self.session.salt = new_salt
if request is None:
raise ValueError('Tried to handle a bad server salt with no request specified')
raise ValueError(
'Tried to handle a bad server salt with no request specified')
# Resend
self.send(request)
@@ -277,15 +284,18 @@ class MtProtoSender:
request.confirm_received = True
if inner_code == 0x2144ca19: # RPC Error
error = RPCError(code=reader.read_int(), message=reader.tgread_string())
error = RPCError(
code=reader.read_int(), message=reader.tgread_string())
if error.must_resend:
if not request:
raise ValueError('The previously sent request must be resent. '
'However, no request was previously sent (called from updates thread).')
raise ValueError(
'The previously sent request must be resent. '
'However, no request was previously sent (called from updates thread).')
request.confirm_received = False
if error.message.startswith('FLOOD_WAIT_'):
print('Should wait {}s. Sleeping until then.'.format(error.additional_data))
print('Should wait {}s. Sleeping until then.'.format(
error.additional_data))
sleep(error.additional_data)
elif error.message.startswith('PHONE_MIGRATE_'):
@@ -295,7 +305,8 @@ class MtProtoSender:
raise error
else:
if not request:
raise ValueError('Cannot receive a request from inside an RPC result from the updates thread.')
raise ValueError(
'Cannot receive a request from inside an RPC result from the updates thread.')
if inner_code == 0x3072cfa1: # GZip packed
unpacked_data = gzip.decompress(reader.tgread_bytes())
@@ -311,7 +322,8 @@ class MtProtoSender:
unpacked_data = gzip.decompress(packed_data)
with BinaryReader(unpacked_data) as compressed_reader:
return self.process_msg(msg_id, sequence, compressed_reader, request)
return self.process_msg(msg_id, sequence, compressed_reader,
request)
# endregion
@@ -340,10 +352,12 @@ class MtProtoSender:
try:
self.updates_thread_receiving = True
seq, body = self.transport.receive(timeout)
message, remote_msg_id, remote_sequence = self.decode_msg(body)
message, remote_msg_id, remote_sequence = self.decode_msg(
body)
with BinaryReader(message) as reader:
self.process_msg(remote_msg_id, remote_sequence, reader)
self.process_msg(remote_msg_id, remote_sequence,
reader)
except (ReadCancelledError, TimeoutError):
pass

View File

@@ -78,7 +78,8 @@ class TcpClient:
if timeout:
time_passed = datetime.now() - start_time
if time_passed > timeout:
raise TimeoutError('The read operation exceeded the timeout.')
raise TimeoutError(
'The read operation exceeded the timeout.')
# If everything went fine, return the read bytes
return writer.get_bytes()

View File

@@ -1,8 +1,8 @@
from binascii import crc32
from datetime import timedelta
from telethon.network import TcpClient
from telethon.errors import *
from telethon.network import TcpClient
from telethon.utils import BinaryWriter
@@ -27,7 +27,7 @@ class TcpTransport:
crc = crc32(writer.get_bytes())
writer.write_int(crc, signed=False)
self.send_counter += 1
self.tcp_client.write(writer.get_bytes())
@@ -45,9 +45,8 @@ class TcpTransport:
body = self.tcp_client.read(packet_length - 12, timeout)
checksum = int.from_bytes(self.tcp_client.read(4, timeout),
byteorder='little',
signed=False)
checksum = int.from_bytes(
self.tcp_client.read(4, timeout), byteorder='little', signed=False)
# Then perform the checks
rv = packet_length_bytes + seq_bytes + body