From f0c801be8c75158ea319410d36914494c2c6a215 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Fri, 18 May 2018 14:15:35 +0200 Subject: [PATCH] Rename AES class methods and add ctr256_encrypt --- pyrogram/crypto/aes.py | 19 ++++++++++--------- pyrogram/session/auth.py | 4 ++-- pyrogram/session/session.py | 4 ++-- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/pyrogram/crypto/aes.py b/pyrogram/crypto/aes.py index 05a01044..19eec9a1 100644 --- a/pyrogram/crypto/aes.py +++ b/pyrogram/crypto/aes.py @@ -38,31 +38,32 @@ else: # TODO: Ugly IFs class AES: @classmethod - def ige_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes: + def ige256_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes: if is_fast: - return tgcrypto.ige_encrypt(data, key, iv) + return tgcrypto.ige256_encrypt(data, key, iv) else: return cls.ige(data, key, iv, True) @classmethod - def ige_decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes: + def ige256_decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes: if is_fast: - return tgcrypto.ige_decrypt(data, key, iv) + return tgcrypto.ige256_decrypt(data, key, iv) else: return cls.ige(data, key, iv, False) @staticmethod - def ctr_decrypt(data: bytes, key: bytes, iv: bytes, offset: int) -> bytes: - replace = int.to_bytes(offset // 16, 4, "big") - iv = iv[:-4] + replace - + def ctr256_encrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes: if is_fast: - return tgcrypto.ctr_decrypt(data, key, iv) + return tgcrypto.ctr256_decrypt(data, key, iv, state) else: ctr = pyaes.AESModeOfOperationCTR(key) ctr._counter._counter = list(iv) return ctr.decrypt(data) + @staticmethod + def ctr256_decrypt(data: bytes, key: bytes, iv: bytes, state: bytes) -> bytes: + return AES.ctr256_encrypt(data, key, iv, state) + @staticmethod def xor(a: bytes, b: bytes) -> bytes: return int.to_bytes( diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py index 449524b3..888d3a4d 100644 --- a/pyrogram/session/auth.py +++ b/pyrogram/session/auth.py @@ -163,7 +163,7 @@ class Auth: server_nonce = int.from_bytes(server_nonce, "little", signed=True) - answer_with_hash = AES.ige_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv) + answer_with_hash = AES.ige256_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv) answer = answer_with_hash[20:] server_dh_inner_data = Object.read(BytesIO(answer)) @@ -192,7 +192,7 @@ class Auth: sha = sha1(data).digest() padding = urandom(- (len(data) + len(sha)) % 16) data_with_hash = sha + data + padding - encrypted_data = AES.ige_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv) + encrypted_data = AES.ige256_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv) log.debug("Send set_client_DH_params") set_client_dh_params_answer = self.send( diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index f8cdcc19..2564b7d1 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -222,14 +222,14 @@ class Session: msg_key = msg_key_large[8:24] aes_key, aes_iv = KDF(self.auth_key, msg_key, True) - return self.auth_key_id + msg_key + AES.ige_encrypt(data + padding, aes_key, aes_iv) + return self.auth_key_id + msg_key + AES.ige256_encrypt(data + padding, aes_key, aes_iv) def unpack(self, b: BytesIO) -> Message: assert b.read(8) == self.auth_key_id, b.getvalue() msg_key = b.read(16) aes_key, aes_iv = KDF(self.auth_key, msg_key, False) - data = BytesIO(AES.ige_decrypt(b.read(), aes_key, aes_iv)) + data = BytesIO(AES.ige256_decrypt(b.read(), aes_key, aes_iv)) data.read(8) # https://core.telegram.org/mtproto/security_guidelines#checking-session-id