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 .. import helpers as utils
from ..extensions import BinaryReader, BinaryWriter
from ..extensions import BinaryReader
class AuthKey:
@ -17,10 +18,5 @@ class AuthKey:
"""Calculates the new nonce hash based on
the current class fields' values
"""
with BinaryWriter() as writer:
writer.write(new_nonce)
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
data = new_nonce + struct.pack('<BQ', number, self.aux_hash)
return utils.calc_msg_key(data)

View File

@ -7,7 +7,7 @@ except ImportError:
rsa = None
raise ImportError('Missing module "rsa", please install via pip.')
from ..extensions import BinaryWriter
from ..tl import TLObject
# {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
in the same way that Telegram does.
"""
with BinaryWriter() as writer:
writer.tgwrite_bytes(get_byte_array(key.n))
writer.tgwrite_bytes(get_byte_array(key.e))
# Telegram uses the last 8 bytes as the fingerprint
return sha1(writer.get_bytes()).digest()[-8:]
n = TLObject.serialize_bytes(get_byte_array(key.n))
e = TLObject.serialize_bytes(get_byte_array(key.e))
# Telegram uses the last 8 bytes as the fingerprint
return sha1(n + e).digest()[-8:]
def add_key(pub):

View File

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

View File

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