Completely refactored unit tests, removed unused code

This commit is contained in:
Lonami
2016-09-08 16:11:37 +02:00
parent a4f68dd29a
commit b2425eeea9
21 changed files with 398 additions and 460 deletions

View File

@@ -1,6 +1,5 @@
from .mtproto_plain_sender import MtProtoPlainSender
from .tcp_client import TcpClient
from .tcp_message import TcpMessage
from .authenticator import do_authentication
from .mtproto_sender import MtProtoSender
from .tcp_transport import TcpTransport

View File

@@ -37,9 +37,7 @@ def do_authentication(transport):
server_nonce = reader.read(16)
pq_bytes = reader.tgread_bytes()
# "string pq is a representation of a natural number (in binary big endian format)"
# See https://core.telegram.org/mtproto/auth_key#dh-exchange-initiation
pq = int.from_bytes(pq_bytes, byteorder='big')
pq = get_int(pq_bytes)
vector_id = reader.read_int()
if vector_id != 0x1cb5c415:
@@ -55,9 +53,9 @@ def do_authentication(transport):
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.tgwrite_bytes(utils.get_byte_array(pq, signed=False))
pq_inner_data_writer.tgwrite_bytes(utils.get_byte_array(min(p, q), signed=False))
pq_inner_data_writer.tgwrite_bytes(utils.get_byte_array(max(p, q), signed=False))
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.write(nonce)
pq_inner_data_writer.write(server_nonce)
pq_inner_data_writer.write(new_nonce)
@@ -78,8 +76,8 @@ def do_authentication(transport):
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(utils.get_byte_array(min(p, q), signed=False))
req_dh_params_writer.tgwrite_bytes(utils.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)
@@ -127,15 +125,13 @@ def do_authentication(transport):
raise AssertionError('Invalid server nonce in encrypted answer')
g = dh_inner_data_reader.read_int()
# "current value of dh_prime equals (in big-endian byte order)"
# See https://core.telegram.org/mtproto/auth_key#presenting-proof-of-work-server-authentication
dh_prime = int.from_bytes(dh_inner_data_reader.tgread_bytes(), byteorder='big', signed=False)
ga = int.from_bytes(dh_inner_data_reader.tgread_bytes(), byteorder='big', signed=False)
dh_prime = get_int(dh_inner_data_reader.tgread_bytes(), signed=False)
ga = get_int(dh_inner_data_reader.tgread_bytes(), signed=False)
server_time = dh_inner_data_reader.read_int()
time_offset = server_time - int(time.time())
b = int.from_bytes(utils.generate_random_bytes(2048), byteorder='big', signed=False)
b = get_int(utils.generate_random_bytes(2048), signed=False)
gb = pow(g, b, dh_prime)
gab = pow(ga, b, dh_prime)
@@ -145,7 +141,7 @@ def do_authentication(transport):
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(utils.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()))
@@ -178,7 +174,7 @@ def do_authentication(transport):
raise NotImplementedError('Invalid server nonce from server')
new_nonce_hash1 = reader.read(16)
auth_key = AuthKey(gab)
auth_key = AuthKey(get_byte_array(gab, signed=False))
new_nonce_hash_calculated = auth_key.calc_new_nonce_hash(new_nonce, 1)
if new_nonce_hash1 != new_nonce_hash_calculated:
@@ -199,3 +195,20 @@ def do_authentication(transport):
def get_fingerprint_text(fingerprint):
"""Gets a fingerprint text in 01-23-45-67-89-AB-CD-EF format (no hyphens)"""
return ''.join(hex(b)[2:].rjust(2, '0').upper() for b in fingerprint)
# The following methods operate in big endian (unlike most of Telegram API) because:
# > "...pq is a representation of a natural number (in binary *big endian* format)..."
# > "...current value of dh_prime equals (in *big-endian* byte order)..."
# Reference: https://core.telegram.org/mtproto/auth_key
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)
def get_int(byte_array, signed=True):
"""Gets the specified integer from its byte array. This should be used by the authenticator,
who requires the data to be in big endian"""
return int.from_bytes(byte_array, byteorder='big', signed=signed)

View File

@@ -26,8 +26,8 @@ class MtProtoPlainSender:
def receive(self):
"""Receives a plain packet, returning the body of the response"""
result = self._transport.receive()
with BinaryReader(result.body) as reader:
seq, body = self._transport.receive()
with BinaryReader(body) as reader:
auth_key_id = reader.read_long()
msg_id = reader.read_long()
message_length = reader.read_int()

View File

@@ -13,10 +13,9 @@ from tl.all_tlobjects import tlobjects
class MtProtoSender:
"""MTProto Mobile Protocol sender (https://core.telegram.org/mtproto/description)"""
def __init__(self, transport, session, swallow_errors=True):
def __init__(self, transport, session):
self.transport = transport
self.session = session
self.swallow_errors = swallow_errors
self.need_confirmation = [] # Message IDs that need confirmation
self.on_update_handlers = []
@@ -59,18 +58,11 @@ class MtProtoSender:
def receive(self, request):
"""Receives the specified MTProtoRequest ("fills in it" the received data)"""
while not request.confirm_received:
try:
message, remote_msg_id, remote_sequence = self.decode_msg(self.transport.receive().body)
with BinaryReader(message) as reader:
self.process_msg(remote_msg_id, remote_sequence, reader, request)
except RPCError as error:
if self.swallow_errors:
print('A RPC error occurred when decoding a message: {}'.format(error))
else:
raise error
seq, body = self.transport.receive()
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)
# endregion

View File

@@ -1,62 +0,0 @@
# This file is based on TLSharp
# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpMessage.cs
from utils import BinaryWriter, BinaryReader
from binascii import crc32
from errors import *
class TcpMessage:
def __init__(self, seq_number, body):
"""
:param seq_number: Sequence number
:param body: Message body byte array
"""
if body is None:
raise InvalidParameterError('body cannot be None')
self.sequence_number = seq_number
self.body = body
def encode(self):
"""Returns the bytes of the this message encoded, following Telegram's guidelines"""
with BinaryWriter() as writer:
''' https://core.telegram.org/mtproto#tcp-transport
4 length bytes are added at the front
(to include the length, the sequence number, and CRC32; always divisible by 4)
and 4 bytes with the packet sequence number within this TCP connection
(the first packet sent is numbered 0, the next one 1, etc.),
and 4 CRC32 bytes at the end (length, sequence number, and payload together).
'''
writer.write_int(len(self.body) + 12)
writer.write_int(self.sequence_number)
writer.write(self.body)
crc = crc32(writer.get_bytes())
writer.write_int(crc, signed=False)
return writer.get_bytes()
@staticmethod
def decode(body):
"""Returns a TcpMessage from the given encoded bytes, decoding them previously"""
if body is None:
raise InvalidParameterError('body cannot be None')
if len(body) < 12:
raise InvalidParameterError('Wrong size of input packet')
with BinaryReader(body) as reader:
packet_len = int.from_bytes(reader.read(4), byteorder='little')
if packet_len < 12:
raise InvalidParameterError('Invalid packet length in body: {}'.format(packet_len))
seq = reader.read_int()
packet = reader.read(packet_len - 12)
checksum = reader.read_int(signed=False)
valid_checksum = crc32(body[:packet_len - 4])
if checksum != valid_checksum:
raise InvalidChecksumError(checksum, valid_checksum)
return TcpMessage(seq, packet)

View File

@@ -1,41 +1,49 @@
# This file is based on TLSharp
# https://github.com/sochix/TLSharp/blob/master/TLSharp.Core/Network/TcpTransport.cs
from network import TcpMessage, TcpClient
from network import TcpClient
from binascii import crc32
from errors import *
from utils import BinaryWriter
class TcpTransport:
def __init__(self, ip_address, port):
self._tcp_client = TcpClient()
self._send_counter = 0
self.tcp_client = TcpClient()
self.send_counter = 0
self._tcp_client.connect(ip_address, port)
self.tcp_client.connect(ip_address, port)
# Original reference: https://core.telegram.org/mtproto#tcp-transport
# The packets are encoded as: total length, sequence number, packet and checksum (CRC32)
def send(self, packet):
"""Sends the given packet (bytes array) to the connected peer"""
if not self._tcp_client.connected:
if not self.tcp_client.connected:
raise ConnectionError('Client not connected to server.')
# Get a TcpMessage which contains the given packet
tcp_message = TcpMessage(self._send_counter, packet)
with BinaryWriter() as writer:
writer.write_int(len(packet) + 12) # 12 = size_of (integer) * 3
writer.write_int(self.send_counter)
writer.write(packet)
self._tcp_client.write(tcp_message.encode())
self._send_counter += 1
crc = crc32(writer.get_bytes())
writer.write_int(crc, signed=False)
self.tcp_client.write(writer.get_bytes())
self.send_counter += 1
def receive(self):
"""Receives a TcpMessage from the connected peer"""
"""Receives a TCP message (tuple(sequence number, body)) from the connected peer"""
# First read everything
packet_length_bytes = self._tcp_client.read(4)
# First read everything we need
packet_length_bytes = self.tcp_client.read(4)
packet_length = int.from_bytes(packet_length_bytes, byteorder='little')
seq_bytes = self._tcp_client.read(4)
seq_bytes = self.tcp_client.read(4)
seq = int.from_bytes(seq_bytes, byteorder='little')
body = self._tcp_client.read(packet_length - 12)
body = self.tcp_client.read(packet_length - 12)
checksum = int.from_bytes(self._tcp_client.read(4), byteorder='little', signed=False)
checksum = int.from_bytes(self.tcp_client.read(4), byteorder='little', signed=False)
# Then perform the checks
rv = packet_length_bytes + seq_bytes + body
@@ -44,9 +52,9 @@ class TcpTransport:
if checksum != valid_checksum:
raise InvalidChecksumError(checksum, valid_checksum)
# If we passed the tests, we can then return a valid TcpMessage
return TcpMessage(seq, body)
# If we passed the tests, we can then return a valid TCP message
return seq, body
def close(self):
if self._tcp_client.connected:
self._tcp_client.close()
if self.tcp_client.connected:
self.tcp_client.close()