Avoid using BinaryWriter where possible

This commit is contained in:
Lonami Exo 2017-09-27 21:23:59 +02:00
parent 8a605f21e6
commit c667a00281
4 changed files with 26 additions and 40 deletions

View File

@ -1,7 +1,8 @@
import struct
from hashlib import sha1 from hashlib import sha1
from .. import helpers as utils from .. import helpers as utils
from ..extensions import BinaryReader, BinaryWriter from ..extensions import BinaryReader
class AuthKey: class AuthKey:
@ -17,10 +18,5 @@ class AuthKey:
"""Calculates the new nonce hash based on """Calculates the new nonce hash based on
the current class fields' values the current class fields' values
""" """
with BinaryWriter() as writer: data = new_nonce + struct.pack('<BQ', number, self.aux_hash)
writer.write(new_nonce) return utils.calc_msg_key(data)
writer.write_byte(number)
writer.write_long(self.aux_hash, signed=False)
new_nonce_hash = utils.calc_msg_key(writer.get_bytes())
return new_nonce_hash

View File

@ -7,7 +7,7 @@ except ImportError:
rsa = None rsa = None
raise ImportError('Missing module "rsa", please install via pip.') raise ImportError('Missing module "rsa", please install via pip.')
from ..extensions import BinaryWriter from ..tl import TLObject
# {fingerprint: Crypto.PublicKey.RSA._RSAobj} dictionary # {fingerprint: Crypto.PublicKey.RSA._RSAobj} dictionary
@ -34,11 +34,10 @@ def _compute_fingerprint(key):
"""For a given Crypto.RSA key, computes its 8-bytes-long fingerprint """For a given Crypto.RSA key, computes its 8-bytes-long fingerprint
in the same way that Telegram does. in the same way that Telegram does.
""" """
with BinaryWriter() as writer: n = TLObject.serialize_bytes(get_byte_array(key.n))
writer.tgwrite_bytes(get_byte_array(key.n)) e = TLObject.serialize_bytes(get_byte_array(key.e))
writer.tgwrite_bytes(get_byte_array(key.e)) # Telegram uses the last 8 bytes as the fingerprint
# Telegram uses the last 8 bytes as the fingerprint return sha1(n + e).digest()[-8:]
return sha1(writer.get_bytes()).digest()[-8:]
def add_key(pub): def add_key(pub):

View File

@ -1,4 +1,5 @@
import os import os
import struct
from datetime import timedelta from datetime import timedelta
from zlib import crc32 from zlib import crc32
from enum import Enum from enum import Enum
@ -6,7 +7,7 @@ from enum import Enum
import errno import errno
from ..crypto import AESModeCTR from ..crypto import AESModeCTR
from ..extensions import BinaryWriter, TcpClient from ..extensions import TcpClient
from ..errors import InvalidChecksumError from ..errors import InvalidChecksumError
@ -175,30 +176,22 @@ class Connection:
# https://core.telegram.org/mtproto#tcp-transport # https://core.telegram.org/mtproto#tcp-transport
# total length, sequence number, packet and checksum (CRC32) # total length, sequence number, packet and checksum (CRC32)
length = len(message) + 12 length = len(message) + 12
with BinaryWriter(known_length=length) as writer: data = struct.pack('<ii', length, self._send_counter) + message
writer.write_int(length) crc = struct.pack('<I', crc32(data))
writer.write_int(self._send_counter) self._send_counter += 1
writer.write(message) self.write(data + crc)
writer.write_int(crc32(writer.get_bytes()), signed=False)
self._send_counter += 1
self.write(writer.get_bytes())
def _send_intermediate(self, message): def _send_intermediate(self, message):
with BinaryWriter(known_length=len(message) + 4) as writer: self.write(struct.pack('<i', len(message)) + message)
writer.write_int(len(message))
writer.write(message)
self.write(writer.get_bytes())
def _send_abridged(self, message): def _send_abridged(self, message):
with BinaryWriter(known_length=len(message) + 4) as writer: length = len(message) >> 2
length = len(message) >> 2 if length < 127:
if length < 127: length = struct.pack('B', length)
writer.write_byte(length) else:
else: length = b'\x7f' + int.to_bytes(length, 3, 'little')
writer.write_byte(127)
writer.write(int.to_bytes(length, 3, 'little')) self.write(length + message)
writer.write(message)
self.write(writer.get_bytes())
# endregion # endregion

View File

@ -1,7 +1,7 @@
import gzip import gzip
import struct
from . import TLObject from . import TLObject
from ..extensions import BinaryWriter
class GzipPacked(TLObject): class GzipPacked(TLObject):
@ -29,10 +29,8 @@ class GzipPacked(TLObject):
def to_bytes(self): def to_bytes(self):
# TODO Maybe compress level could be an option # TODO Maybe compress level could be an option
with BinaryWriter() as writer: return struct.pack('<I', GzipPacked.constructor_id) + \
writer.write_int(GzipPacked.constructor_id, signed=False) TLObject.serialize_bytes(gzip.compress(self.data))
writer.tgwrite_bytes(gzip.compress(self.data))
return writer.get_bytes()
@staticmethod @staticmethod
def read(reader): def read(reader):