Replace .to_bytes() with the special .__bytes__ function

This commit is contained in:
Lonami Exo 2017-10-17 19:54:24 +02:00
parent 63dfb1e3ea
commit adb79b21cf
7 changed files with 22 additions and 22 deletions

View File

@ -42,7 +42,7 @@ def _do_authentication(connection):
req_pq_request = ReqPqRequest( req_pq_request = ReqPqRequest(
nonce=int.from_bytes(os.urandom(16), 'big', signed=True) nonce=int.from_bytes(os.urandom(16), 'big', signed=True)
) )
sender.send(req_pq_request.to_bytes()) sender.send(bytes(req_pq_request))
with BinaryReader(sender.receive()) as reader: with BinaryReader(sender.receive()) as reader:
req_pq_request.on_response(reader) req_pq_request.on_response(reader)
@ -60,12 +60,12 @@ def _do_authentication(connection):
p, q = rsa.get_byte_array(min(p, q)), rsa.get_byte_array(max(p, q)) p, q = rsa.get_byte_array(min(p, q)), rsa.get_byte_array(max(p, q))
new_nonce = int.from_bytes(os.urandom(32), 'little', signed=True) new_nonce = int.from_bytes(os.urandom(32), 'little', signed=True)
pq_inner_data = PQInnerData( pq_inner_data = bytes(PQInnerData(
pq=rsa.get_byte_array(pq), p=p, q=q, pq=rsa.get_byte_array(pq), p=p, q=q,
nonce=res_pq.nonce, nonce=res_pq.nonce,
server_nonce=res_pq.server_nonce, server_nonce=res_pq.server_nonce,
new_nonce=new_nonce new_nonce=new_nonce
).to_bytes() ))
# sha_digest + data + random_bytes # sha_digest + data + random_bytes
cipher_text, target_fingerprint = None, None cipher_text, target_fingerprint = None, None
@ -90,7 +90,7 @@ def _do_authentication(connection):
public_key_fingerprint=target_fingerprint, public_key_fingerprint=target_fingerprint,
encrypted_data=cipher_text encrypted_data=cipher_text
) )
sender.send(req_dh_params.to_bytes()) sender.send(bytes(req_dh_params))
# Step 2 response: DH Exchange # Step 2 response: DH Exchange
with BinaryReader(sender.receive()) as reader: with BinaryReader(sender.receive()) as reader:
@ -138,12 +138,12 @@ def _do_authentication(connection):
gab = pow(g_a, b, dh_prime) gab = pow(g_a, b, dh_prime)
# Prepare client DH Inner Data # Prepare client DH Inner Data
client_dh_inner = ClientDHInnerData( client_dh_inner = bytes(ClientDHInnerData(
nonce=res_pq.nonce, nonce=res_pq.nonce,
server_nonce=res_pq.server_nonce, server_nonce=res_pq.server_nonce,
retry_id=0, # TODO Actual retry ID retry_id=0, # TODO Actual retry ID
g_b=rsa.get_byte_array(gb) g_b=rsa.get_byte_array(gb)
).to_bytes() ))
client_dh_inner_hashed = sha1(client_dh_inner).digest() + client_dh_inner client_dh_inner_hashed = sha1(client_dh_inner).digest() + client_dh_inner
@ -156,7 +156,7 @@ def _do_authentication(connection):
server_nonce=res_pq.server_nonce, server_nonce=res_pq.server_nonce,
encrypted_data=client_dh_encrypted, encrypted_data=client_dh_encrypted,
) )
sender.send(set_client_dh.to_bytes()) sender.send(bytes(set_client_dh))
# Step 3 response: Complete DH Exchange # Step 3 response: Complete DH Exchange
with BinaryReader(sender.receive()) as reader: with BinaryReader(sender.receive()) as reader:

View File

@ -125,7 +125,7 @@ class MtProtoSender:
plain_text = \ plain_text = \
struct.pack('<QQ', self.session.salt, self.session.id) \ struct.pack('<QQ', self.session.salt, self.session.id) \
+ message.to_bytes() + bytes(message)
msg_key = utils.calc_msg_key(plain_text) msg_key = utils.calc_msg_key(plain_text)
key_id = struct.pack('<Q', self.session.auth_key.key_id) key_id = struct.pack('<Q', self.session.auth_key.key_id)

View File

@ -13,21 +13,21 @@ class GzipPacked(TLObject):
@staticmethod @staticmethod
def gzip_if_smaller(request): def gzip_if_smaller(request):
"""Calls request.to_bytes(), and based on a certain threshold, """Calls bytes(request), and based on a certain threshold,
optionally gzips the resulting data. If the gzipped data is optionally gzips the resulting data. If the gzipped data is
smaller than the original byte array, this is returned instead. smaller than the original byte array, this is returned instead.
Note that this only applies to content related requests. Note that this only applies to content related requests.
""" """
data = request.to_bytes() data = bytes(request)
# TODO This threshold could be configurable # TODO This threshold could be configurable
if request.content_related and len(data) > 512: if request.content_related and len(data) > 512:
gzipped = GzipPacked(data).to_bytes() gzipped = bytes(GzipPacked(data))
return gzipped if len(gzipped) < len(data) else data return gzipped if len(gzipped) < len(data) else data
else: else:
return data return data
def to_bytes(self): def __bytes__(self):
# TODO Maybe compress level could be an option # TODO Maybe compress level could be an option
return struct.pack('<I', GzipPacked.CONSTRUCTOR_ID) + \ return struct.pack('<I', GzipPacked.CONSTRUCTOR_ID) + \
TLObject.serialize_bytes(gzip.compress(self.data)) TLObject.serialize_bytes(gzip.compress(self.data))

View File

@ -11,10 +11,10 @@ class MessageContainer(TLObject):
self.content_related = False self.content_related = False
self.messages = messages self.messages = messages
def to_bytes(self): def __bytes__(self):
return struct.pack( return struct.pack(
'<Ii', MessageContainer.CONSTRUCTOR_ID, len(self.messages) '<Ii', MessageContainer.CONSTRUCTOR_ID, len(self.messages)
) + b''.join(m.to_bytes() for m in self.messages) ) + b''.join(bytes(m) for m in self.messages)
@staticmethod @staticmethod
def iter_read(reader): def iter_read(reader):

View File

@ -12,6 +12,6 @@ class TLMessage(TLObject):
self.seq_no = session.generate_sequence(request.content_related) self.seq_no = session.generate_sequence(request.content_related)
self.request = request self.request = request
def to_bytes(self): def __bytes__(self):
body = GzipPacked.gzip_if_smaller(self.request) body = GzipPacked.gzip_if_smaller(self.request)
return struct.pack('<qii', self.msg_id, self.seq_no, len(body)) + body return struct.pack('<qii', self.msg_id, self.seq_no, len(body)) + body

View File

@ -125,7 +125,7 @@ class TLObject:
def to_dict(self, recursive=True): def to_dict(self, recursive=True):
return {} return {}
def to_bytes(self): def __bytes__(self):
return b'' return b''
@staticmethod @staticmethod

View File

@ -154,7 +154,7 @@ class TLGenerator:
# for all those TLObjects with arg.can_be_inferred. # for all those TLObjects with arg.can_be_inferred.
builder.writeln('import os') builder.writeln('import os')
# Import struct for the .to_bytes(self) serialization # Import struct for the .__bytes__(self) serialization
builder.writeln('import struct') builder.writeln('import struct')
# Generate the class for every TLObject # Generate the class for every TLObject
@ -302,8 +302,8 @@ class TLGenerator:
builder.end_block() builder.end_block()
# Write the .to_bytes() function # Write the .__bytes__() function
builder.writeln('def to_bytes(self):') builder.writeln('def __bytes__(self):')
# Some objects require more than one flag parameter to be set # Some objects require more than one flag parameter to be set
# at the same time. In this case, add an assertion. # at the same time. In this case, add an assertion.
@ -441,10 +441,10 @@ class TLGenerator:
@staticmethod @staticmethod
def write_to_bytes(builder, arg, args, name=None): def write_to_bytes(builder, arg, args, name=None):
""" """
Writes the .to_bytes() code for the given argument Writes the .__bytes__() code for the given argument
:param builder: The source code builder :param builder: The source code builder
:param arg: The argument to write :param arg: The argument to write
:param args: All the other arguments in TLObject same to_bytes. :param args: All the other arguments in TLObject same __bytes__.
This is required to determine the flags value This is required to determine the flags value
:param name: The name of the argument. Defaults to "self.argname" :param name: The name of the argument. Defaults to "self.argname"
This argument is an option because it's required when This argument is an option because it's required when
@ -540,7 +540,7 @@ class TLGenerator:
else: else:
# Else it may be a custom type # Else it may be a custom type
builder.write('{}.to_bytes()'.format(name)) builder.write('bytes({})'.format(name))
if arg.is_flag: if arg.is_flag:
builder.write(')') builder.write(')')