feat(crypto): enhance mtproto padding and rsa key validation

Author: SychO3

pyrogram/crypto/aes.py

@@ -20,6 +20,15 @@
 
 log = logging.getLogger(__name__)
 
+
+def xor(a: bytes, b: bytes) -> bytes:
+    return int.to_bytes(
+        int.from_bytes(a, "big") ^ int.from_bytes(b, "big"),
+        len(a),
+        "big",
+    )
+
+
 try:
     import tgcrypto
 
@@ -41,13 +50,6 @@ def ctr256_encrypt(data: bytes, key: bytes, iv: bytearray, state: bytearray = No
     def ctr256_decrypt(data: bytes, key: bytes, iv: bytearray, state: bytearray = None) -> bytes:
         return tgcrypto.ctr256_decrypt(data, key, iv, state or bytearray(1))
 
-
-    def xor(a: bytes, b: bytes) -> bytes:
-        return int.to_bytes(
-            int.from_bytes(a, "big") ^ int.from_bytes(b, "big"),
-            len(a),
-            "big",
-        )
 except ImportError:
     import pyaes
 
@@ -74,14 +76,6 @@ def ctr256_decrypt(data: bytes, key: bytes, iv: bytearray, state: bytearray = No
         return ctr(data, key, iv, state or bytearray(1))
 
 
-    def xor(a: bytes, b: bytes) -> bytes:
-        return int.to_bytes(
-            int.from_bytes(a, "big") ^ int.from_bytes(b, "big"),
-            len(a),
-            "big",
-        )
-
-
     def ige(data: bytes, key: bytes, iv: bytes, encrypt: bool) -> bytes:
         cipher = pyaes.AES(key)
 

pyrogram/crypto/mtproto.py

@@ -19,6 +19,7 @@
 from hashlib import sha256
 from io import BytesIO
 from os import urandom
+from random import randint
 
 from pyrogram.errors import SecurityCheckMismatch
 from pyrogram.raw.core import Message, Long
@@ -40,7 +41,9 @@ def kdf(auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple:
 
 def pack(message: Message, salt: int, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -> bytes:
     data = Long(salt) + session_id + message.write()
-    padding = urandom(-(len(data) + 12) % 16 + 12)
+    # Spec allows 12-1024 bytes; add extra random multiples of 16 for traffic analysis resistance
+    min_padding = -(len(data) + 12) % 16 + 12
+    padding = urandom(min_padding + 16 * randint(0, 4))
 
     # 88 = 88 + 0 (outgoing message)
     msg_key_large = sha256(auth_key[88: 88 + 32] + data + padding).digest()

pyrogram/crypto/rsa.py

@@ -252,8 +252,14 @@
 
 
 def encrypt(data: bytes, fingerprint: int) -> bytes:
+    key = server_public_keys.get(fingerprint)
+    if key is None:
+        raise ValueError(
+            f"Unknown server public key fingerprint: {fingerprint:#018x}. "
+            f"Known fingerprints: {', '.join(f'{fp:#018x}' for fp in server_public_keys)}"
+        )
     return pow(
         int.from_bytes(data, "big"),
-        server_public_keys[fingerprint].e,
-        server_public_keys[fingerprint].m
+        key.e,
+        key.m
     ).to_bytes(256, "big")