From 2b7425019bc410802c7f36ea3beb0b26a16e9191 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sun, 28 Jan 2018 01:44:38 +0100
Subject: [PATCH] Merge IGE and CTR into a single class (AES)
---
pyrogram/client/client.py | 6 +--
pyrogram/crypto/__init__.py | 3 +-
pyrogram/crypto/aes.py | 88 +++++++++++++++++++++++++++++++++++++
pyrogram/crypto/ctr.py | 35 ---------------
pyrogram/crypto/ige.py | 64 ---------------------------
pyrogram/session/auth.py | 8 ++--
pyrogram/session/session.py | 6 +--
7 files changed, 98 insertions(+), 112 deletions(-)
create mode 100644 pyrogram/crypto/aes.py
delete mode 100644 pyrogram/crypto/ctr.py
delete mode 100644 pyrogram/crypto/ige.py
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index 5fd75cc0..d58656c9 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -46,7 +46,7 @@ from pyrogram.api.types import (
InputPeerEmpty, InputPeerSelf,
InputPeerUser, InputPeerChat, InputPeerChannel
)
-from pyrogram.crypto import CTR
+from pyrogram.crypto import AES
from pyrogram.session import Auth, Session
from .style import Markdown, HTML
@@ -1633,8 +1633,6 @@ class Client:
)
)
if isinstance(r, types.upload.FileCdnRedirect):
- ctr = CTR(r.encryption_key, r.encryption_iv)
-
cdn_session = Session(
r.dc_id,
self.test_mode,
@@ -1673,7 +1671,7 @@ class Client:
break
# https://core.telegram.org/cdn#decrypting-files
- decrypted_chunk = ctr.decrypt(chunk, offset)
+ decrypted_chunk = AES.ctr_decrypt(chunk, r.encryption_key, r.encryption_iv, offset)
# TODO: https://core.telegram.org/cdn#verifying-files
# TODO: Save to temp file, flush each chunk, rename to full if everything is ok
diff --git a/pyrogram/crypto/__init__.py b/pyrogram/crypto/__init__.py
index fa9b528d..08ed44f0 100644
--- a/pyrogram/crypto/__init__.py
+++ b/pyrogram/crypto/__init__.py
@@ -16,8 +16,7 @@
# You should have received a copy of the GNU Lesser General Public License
# along with Pyrogram. If not, see .
-from .ctr import CTR
-from .ige import IGE
+from .aes import AES
from .kdf import KDF
from .prime import Prime
from .rsa import RSA
diff --git a/pyrogram/crypto/aes.py b/pyrogram/crypto/aes.py
new file mode 100644
index 00000000..8d971370
--- /dev/null
+++ b/pyrogram/crypto/aes.py
@@ -0,0 +1,88 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2018 Dan Tès
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram. If not, see .
+
+import logging
+
+log = logging.getLogger(__name__)
+
+try:
+ import tgcrypto
+except ImportError:
+ logging.warning("Warning: TgCrypto is missing")
+ is_fast = False
+ import pyaes
+else:
+ log.info("Using TgCrypto")
+ is_fast = True
+
+
+# TODO: Ugly IFs
+class AES:
+ @classmethod
+ def ige_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
+ if is_fast:
+ return tgcrypto.ige_encrypt(data, key, iv)
+ else:
+ return cls.ige(data, key, iv, True)
+
+ @classmethod
+ def ige_decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
+ if is_fast:
+ return tgcrypto.ige_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, byteorder="big", length=4)
+ iv = iv[:-4] + replace
+
+ if is_fast:
+ return tgcrypto.ctr_decrypt(data, key, iv)
+ else:
+ ctr = pyaes.AESModeOfOperationCTR(key)
+ ctr._counter._counter = list(iv)
+ return ctr.decrypt(data)
+
+ @staticmethod
+ def xor(a: bytes, b: bytes) -> bytes:
+ return int.to_bytes(
+ int.from_bytes(a, "big") ^ int.from_bytes(b, "big"),
+ len(a),
+ "big",
+ )
+
+ @classmethod
+ def ige(cls, data: bytes, key: bytes, iv: bytes, encrypt: bool) -> bytes:
+ cipher = pyaes.AES(key)
+
+ iv_1 = iv[:16]
+ iv_2 = iv[16:]
+
+ data = [data[i: i + 16] for i in range(0, len(data), 16)]
+
+ if encrypt:
+ for i, chunk in enumerate(data):
+ iv_1 = data[i] = cls.xor(cipher.encrypt(cls.xor(chunk, iv_1)), iv_2)
+ iv_2 = chunk
+ else:
+ for i, chunk in enumerate(data):
+ iv_2 = data[i] = cls.xor(cipher.decrypt(cls.xor(chunk, iv_2)), iv_1)
+ iv_1 = chunk
+
+ return b"".join(data)
diff --git a/pyrogram/crypto/ctr.py b/pyrogram/crypto/ctr.py
deleted file mode 100644
index 25cd4181..00000000
--- a/pyrogram/crypto/ctr.py
+++ /dev/null
@@ -1,35 +0,0 @@
-# Pyrogram - Telegram MTProto API Client Library for Python
-# Copyright (C) 2017-2018 Dan Tès
-#
-# This file is part of Pyrogram.
-#
-# Pyrogram is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Lesser General Public License as published
-# by the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# Pyrogram is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Lesser General Public License for more details.
-#
-# You should have received a copy of the GNU Lesser General Public License
-# along with Pyrogram. If not, see .
-
-try:
- from pyaes import AESModeOfOperationCTR
-except ImportError:
- pass
-
-
-class CTR:
- def __init__(self, key: bytes, iv: bytes):
- self.ctr = AESModeOfOperationCTR(key)
- self.iv = iv
-
- def decrypt(self, data: bytes, offset: int) -> bytes:
- replace = int.to_bytes(offset // 16, byteorder="big", length=4)
- iv = self.iv[:-4] + replace
- self.ctr._counter._counter = list(iv)
-
- return self.ctr.decrypt(data)
diff --git a/pyrogram/crypto/ige.py b/pyrogram/crypto/ige.py
deleted file mode 100644
index 03b4c399..00000000
--- a/pyrogram/crypto/ige.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# Pyrogram - Telegram MTProto API Client Library for Python
-# Copyright (C) 2017-2018 Dan Tès
-#
-# This file is part of Pyrogram.
-#
-# Pyrogram is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Lesser General Public License as published
-# by the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# Pyrogram is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Lesser General Public License for more details.
-#
-# You should have received a copy of the GNU Lesser General Public License
-# along with Pyrogram. If not, see .
-
-# from pyaes import AES
-import tgcrypto
-
-BLOCK_SIZE = 16
-
-
-# TODO: Performance optimization
-
-class IGE:
- @classmethod
- def encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
- return tgcrypto.ige_encrypt(data, key, iv)
- # return cls.ige(data, key, iv, True)
-
- @classmethod
- def decrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes:
- return tgcrypto.ige_decrypt(data, key, iv)
- # return cls.ige(data, key, iv, False)
-
- @staticmethod
- def xor(a: bytes, b: bytes) -> bytes:
- return int.to_bytes(
- int.from_bytes(a, "big") ^ int.from_bytes(b, "big"),
- len(a),
- "big",
- )
-
- # @classmethod
- # def ige(cls, data: bytes, key: bytes, iv: bytes, encrypt: bool) -> bytes:
- # cipher = AES(key)
- #
- # iv_1 = iv[:BLOCK_SIZE]
- # iv_2 = iv[BLOCK_SIZE:]
- #
- # data = [data[i: i + BLOCK_SIZE] for i in range(0, len(data), BLOCK_SIZE)]
- #
- # if encrypt:
- # for i, chunk in enumerate(data):
- # iv_1 = data[i] = cls.xor(cipher.encrypt(cls.xor(chunk, iv_1)), iv_2)
- # iv_2 = chunk
- # else:
- # for i, chunk in enumerate(data):
- # iv_2 = data[i] = cls.xor(cipher.decrypt(cls.xor(chunk, iv_2)), iv_1)
- # iv_1 = chunk
- #
- # return b"".join(data)
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
index f3d7a3a3..741e9a44 100644
--- a/pyrogram/session/auth.py
+++ b/pyrogram/session/auth.py
@@ -25,7 +25,7 @@ from os import urandom
from pyrogram.api import functions, types
from pyrogram.api.core import Object, Long, Int
from pyrogram.connection import Connection
-from pyrogram.crypto import IGE, RSA, Prime
+from pyrogram.crypto import AES, RSA, Prime
from .internals import MsgId, DataCenter
log = logging.getLogger(__name__)
@@ -152,7 +152,7 @@ class Auth:
server_nonce = int.from_bytes(server_nonce, "little", signed=True)
- answer_with_hash = IGE.decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
+ answer_with_hash = AES.ige_decrypt(encrypted_answer, tmp_aes_key, tmp_aes_iv)
answer = answer_with_hash[20:]
server_dh_inner_data = Object.read(BytesIO(answer))
@@ -181,7 +181,7 @@ class Auth:
sha = sha1(data).digest()
padding = urandom(- (len(data) + len(sha)) % 16)
data_with_hash = sha + data + padding
- encrypted_data = IGE.encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
+ encrypted_data = AES.ige_encrypt(data_with_hash, tmp_aes_key, tmp_aes_iv)
log.debug("Send set_client_DH_params")
set_client_dh_params_answer = self.send(
@@ -236,7 +236,7 @@ class Auth:
log.debug("Nonce fields check: OK")
# Step 9
- server_salt = IGE.xor(new_nonce[:8], server_nonce[:8])
+ server_salt = AES.xor(new_nonce[:8], server_nonce[:8])
log.debug("Server salt: {}".format(int.from_bytes(server_salt, "little")))
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 23d686a4..89d905d1 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -32,7 +32,7 @@ from pyrogram.api.all import layer
from pyrogram.api.core import Message, Object, MsgContainer, Long, FutureSalt, Int
from pyrogram.api.errors import Error
from pyrogram.connection import Connection
-from pyrogram.crypto import IGE, KDF
+from pyrogram.crypto import AES, KDF
from .internals import MsgId, MsgFactory, DataCenter
log = logging.getLogger(__name__)
@@ -204,14 +204,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 + IGE.encrypt(data + padding, aes_key, aes_iv)
+ return self.auth_key_id + msg_key + AES.ige_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(IGE.decrypt(b.read(), aes_key, aes_iv))
+ data = BytesIO(AES.ige_decrypt(b.read(), aes_key, aes_iv))
data.read(8)
# https://core.telegram.org/mtproto/security_guidelines#checking-session-id