From cd027b8c1c252149cc7304b95b6ec8da53bd01fa Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 15 Dec 2021 13:18:13 +0100
Subject: [PATCH 1/8] Implement missing MTProto checks

---
 pyrogram/crypto/mtproto.py        | 44 ++++++++++++++++++++++++++++---
 pyrogram/raw/core/future_salt.py  |  9 +++++++
 pyrogram/raw/core/future_salts.py | 14 ++++++++++
 pyrogram/session/session.py       |  8 +++++-
 4 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index 803db297..1eec7b7a 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -19,9 +19,11 @@
 from hashlib import sha256
 from io import BytesIO
 from os import urandom
+from typing import Optional, List
 
 from pyrogram.raw.core import Message, Long
 from . import aes
+from ..session.internals import MsgId
 
 
 def kdf(auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple:
@@ -49,13 +51,19 @@ def pack(message: Message, salt: int, session_id: bytes, auth_key: bytes, auth_k
     return auth_key_id + msg_key + aes.ige256_encrypt(data + padding, aes_key, aes_iv)
 
 
-def unpack(b: BytesIO, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -> Message:
+def unpack(
+    b: BytesIO,
+    session_id: bytes,
+    auth_key: bytes,
+    auth_key_id: bytes,
+    stored_msg_ids: List[int]
+) -> Optional[Message]:
     assert b.read(8) == auth_key_id, b.getvalue()
 
     msg_key = b.read(16)
     aes_key, aes_iv = kdf(auth_key, msg_key, False)
     data = BytesIO(aes.ige256_decrypt(b.read(), aes_key, aes_iv))
-    data.read(8)
+    data.read(8)  # Salt
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
     assert data.read(8) == session_id
@@ -75,11 +83,41 @@ def unpack(b: BytesIO, session_id: bytes, auth_key: bytes, auth_key_id: bytes) -
         raise ValueError(f"The server sent an unknown constructor: {hex(e.args[0])}\n{left}")
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key
-    # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
     # 96 = 88 + 8 (incoming message)
     assert msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24]
 
+    # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
+    data.seek(32)  # Get to the payload, skip salt (8) + session_id (8) + msg_id (8) + seq_no (4) + length (4)
+    payload = data.read()
+    padding = payload[message.length:]
+    assert 12 <= len(padding) <= 1024
+    assert len(payload) % 4 == 0
+
     # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
     assert message.msg_id % 2 != 0
 
+    if len(stored_msg_ids) > 200:
+        stored_msg_ids = stored_msg_ids[50:]
+
+    if stored_msg_ids:
+        # Ignored message: msg_id is lower than all of the stored values
+        if message.msg_id < stored_msg_ids[0]:
+            return None
+
+        # Ignored message: msg_id is equal to any of the stored values
+        if message.msg_id in stored_msg_ids:
+            return None
+
+        time_diff = (message.msg_id - MsgId()) / 2 ** 32
+
+        # Ignored message: msg_id belongs over 30 seconds in the future
+        if time_diff > 30:
+            return None
+
+        # Ignored message: msg_id belongs over 300 seconds in the past
+        if time_diff < -300:
+            return None
+
+    stored_msg_ids.append(message.msg_id)
+
     return message
diff --git a/pyrogram/raw/core/future_salt.py b/pyrogram/raw/core/future_salt.py
index 85303d12..54a12963 100644
--- a/pyrogram/raw/core/future_salt.py
+++ b/pyrogram/raw/core/future_salt.py
@@ -42,3 +42,12 @@ class FutureSalt(TLObject):
         salt = Long.read(data)
 
         return FutureSalt(valid_since, valid_until, salt)
+
+    def write(self, *args: Any) -> bytes:
+        b = BytesIO()
+
+        b.write(Int(self.valid_since))
+        b.write(Int(self.valid_until))
+        b.write(Long(self.salt))
+
+        return b.getvalue()
diff --git a/pyrogram/raw/core/future_salts.py b/pyrogram/raw/core/future_salts.py
index faa4b741..9fa2f8e9 100644
--- a/pyrogram/raw/core/future_salts.py
+++ b/pyrogram/raw/core/future_salts.py
@@ -45,3 +45,17 @@ class FutureSalts(TLObject):
         salts = [FutureSalt.read(data) for _ in range(count)]
 
         return FutureSalts(req_msg_id, now, salts)
+
+    def write(self, *args: Any) -> bytes:
+        b = BytesIO()
+
+        b.write(Long(self.req_msg_id))
+        b.write(Int(self.now))
+
+        count = len(self.salts)
+        b.write(Int(count))
+
+        for salt in self.salts:
+            b.write(salt.write())
+
+        return b.getvalue()
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 721586a0..39fe605c 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -102,6 +102,8 @@ class Session:
 
         self.results = {}
 
+        self.stored_msg_ids = []
+
         self.ping_task = None
         self.ping_task_event = asyncio.Event()
 
@@ -224,9 +226,13 @@ class Session:
             BytesIO(packet),
             self.session_id,
             self.auth_key,
-            self.auth_key_id
+            self.auth_key_id,
+            self.stored_msg_ids
         )
 
+        if data is None:
+            return
+
         messages = (
             data.body.messages
             if isinstance(data.body, MsgContainer)

From bc420da0e2ec34506263d0a8abef7019e1770194 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 15 Dec 2021 15:04:44 +0100
Subject: [PATCH 2/8] Maintain a sorted list of stored_msg_ids

---
 pyrogram/crypto/mtproto.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index 1eec7b7a..81cf56f4 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+import bisect
 from hashlib import sha256
 from io import BytesIO
 from os import urandom
@@ -118,6 +119,6 @@ def unpack(
         if time_diff < -300:
             return None
 
-    stored_msg_ids.append(message.msg_id)
+    bisect.insort(stored_msg_ids, message.msg_id)
 
     return message

From 2a1af2b8e9afb54cc8f96bffb95317e52fbafcd3 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 15 Dec 2021 16:02:39 +0100
Subject: [PATCH 3/8] Close and reestablish the TCP connection in case of
 mismatch

---
 pyrogram/crypto/mtproto.py  | 14 +++++++-------
 pyrogram/session/session.py | 25 +++++++++++++++----------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index 81cf56f4..a22693fe 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -20,7 +20,7 @@ import bisect
 from hashlib import sha256
 from io import BytesIO
 from os import urandom
-from typing import Optional, List
+from typing import List, Tuple
 
 from pyrogram.raw.core import Message, Long
 from . import aes
@@ -58,7 +58,7 @@ def unpack(
     auth_key: bytes,
     auth_key_id: bytes,
     stored_msg_ids: List[int]
-) -> Optional[Message]:
+) -> Tuple[Message, bool]:
     assert b.read(8) == auth_key_id, b.getvalue()
 
     msg_key = b.read(16)
@@ -103,22 +103,22 @@ def unpack(
     if stored_msg_ids:
         # Ignored message: msg_id is lower than all of the stored values
         if message.msg_id < stored_msg_ids[0]:
-            return None
+            return message, False
 
         # Ignored message: msg_id is equal to any of the stored values
         if message.msg_id in stored_msg_ids:
-            return None
+            return message, False
 
         time_diff = (message.msg_id - MsgId()) / 2 ** 32
 
         # Ignored message: msg_id belongs over 30 seconds in the future
         if time_diff > 30:
-            return None
+            return message, False
 
         # Ignored message: msg_id belongs over 300 seconds in the past
         if time_diff < -300:
-            return None
+            return message, False
 
     bisect.insort(stored_msg_ids, message.msg_id)
 
-    return message
+    return message, True
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 39fe605c..b3875f3c 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -220,17 +220,22 @@ class Session:
         await self.start()
 
     async def handle_packet(self, packet):
-        data = await self.loop.run_in_executor(
-            pyrogram.crypto_executor,
-            mtproto.unpack,
-            BytesIO(packet),
-            self.session_id,
-            self.auth_key,
-            self.auth_key_id,
-            self.stored_msg_ids
-        )
+        try:
+            data, ok = await self.loop.run_in_executor(
+                pyrogram.crypto_executor,
+                mtproto.unpack,
+                BytesIO(packet),
+                self.session_id,
+                self.auth_key,
+                self.auth_key_id,
+                self.stored_msg_ids
+            )
+        except AssertionError:
+            self.connection.close()
+            return
 
-        if data is None:
+        if not ok:
+            self.connection.close()
             return
 
         messages = (

From c2a29c8c302feff24e84e7974b7618f4437db105 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 15 Dec 2021 16:53:17 +0100
Subject: [PATCH 4/8] Tune stored_msg_ids max size

---
 pyrogram/crypto/mtproto.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index a22693fe..2b4f71af 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -26,6 +26,8 @@ from pyrogram.raw.core import Message, Long
 from . import aes
 from ..session.internals import MsgId
 
+STORED_MSG_IDS_MAX_SIZE = 1000 * 2
+
 
 def kdf(auth_key: bytes, msg_key: bytes, outgoing: bool) -> tuple:
     # https://core.telegram.org/mtproto/description#defining-aes-key-and-initialization-vector
@@ -97,8 +99,8 @@ def unpack(
     # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
     assert message.msg_id % 2 != 0
 
-    if len(stored_msg_ids) > 200:
-        stored_msg_ids = stored_msg_ids[50:]
+    if len(stored_msg_ids) > STORED_MSG_IDS_MAX_SIZE:
+        del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2]
 
     if stored_msg_ids:
         # Ignored message: msg_id is lower than all of the stored values

From ed9c7e4694f770f2ef90adc87152889c8530fc01 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 15 Dec 2021 19:26:54 +0100
Subject: [PATCH 5/8] Simplify the error handling a bit

---
 pyrogram/crypto/mtproto.py  | 16 ++++++++--------
 pyrogram/session/session.py |  6 +-----
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index 2b4f71af..f9b885f6 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -20,7 +20,7 @@ import bisect
 from hashlib import sha256
 from io import BytesIO
 from os import urandom
-from typing import List, Tuple
+from typing import List
 
 from pyrogram.raw.core import Message, Long
 from . import aes
@@ -60,8 +60,8 @@ def unpack(
     auth_key: bytes,
     auth_key_id: bytes,
     stored_msg_ids: List[int]
-) -> Tuple[Message, bool]:
-    assert b.read(8) == auth_key_id, b.getvalue()
+) -> Message:
+    assert b.read(8) == auth_key_id
 
     msg_key = b.read(16)
     aes_key, aes_iv = kdf(auth_key, msg_key, False)
@@ -105,22 +105,22 @@ def unpack(
     if stored_msg_ids:
         # Ignored message: msg_id is lower than all of the stored values
         if message.msg_id < stored_msg_ids[0]:
-            return message, False
+            assert False
 
         # Ignored message: msg_id is equal to any of the stored values
         if message.msg_id in stored_msg_ids:
-            return message, False
+            assert False
 
         time_diff = (message.msg_id - MsgId()) / 2 ** 32
 
         # Ignored message: msg_id belongs over 30 seconds in the future
         if time_diff > 30:
-            return message, False
+            assert False
 
         # Ignored message: msg_id belongs over 300 seconds in the past
         if time_diff < -300:
-            return message, False
+            assert False
 
     bisect.insort(stored_msg_ids, message.msg_id)
 
-    return message, True
+    return message
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index b3875f3c..3504c3ad 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -221,7 +221,7 @@ class Session:
 
     async def handle_packet(self, packet):
         try:
-            data, ok = await self.loop.run_in_executor(
+            data = await self.loop.run_in_executor(
                 pyrogram.crypto_executor,
                 mtproto.unpack,
                 BytesIO(packet),
@@ -234,10 +234,6 @@ class Session:
             self.connection.close()
             return
 
-        if not ok:
-            self.connection.close()
-            return
-
         messages = (
             data.body.messages
             if isinstance(data.body, MsgContainer)

From a720726479bd1f9b6804b837a7679bc4b9f682b1 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 16 Dec 2021 21:05:01 +0100
Subject: [PATCH 6/8] Remove unneeded assertion

---
 pyrogram/methods/auth/accept_terms_of_service.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pyrogram/methods/auth/accept_terms_of_service.py b/pyrogram/methods/auth/accept_terms_of_service.py
index b5abab86..c8cfd36d 100644
--- a/pyrogram/methods/auth/accept_terms_of_service.py
+++ b/pyrogram/methods/auth/accept_terms_of_service.py
@@ -36,6 +36,4 @@ class AcceptTermsOfService(Scaffold):
             )
         )
 
-        assert r
-
-        return True
+        return bool(r)

From 8aa358129c3484a2cd35454dc82228e34663be13 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 16 Dec 2021 21:38:24 +0100
Subject: [PATCH 7/8] Use specialized exceptions for handling security checks

---
 pyrogram/client.py          |  3 ++-
 pyrogram/crypto/mtproto.py  | 21 +++++++++++----------
 pyrogram/errors/__init__.py | 24 ++++++++++++++++++++++++
 pyrogram/session/auth.py    | 25 +++++++++++++------------
 pyrogram/session/session.py |  7 +++++--
 5 files changed, 55 insertions(+), 25 deletions(-)

diff --git a/pyrogram/client.py b/pyrogram/client.py
index b54794fa..8867d2f9 100644
--- a/pyrogram/client.py
+++ b/pyrogram/client.py
@@ -35,6 +35,7 @@ import pyrogram
 from pyrogram import raw
 from pyrogram import utils
 from pyrogram.crypto import aes
+from pyrogram.errors import CDNFileHashMismatch
 from pyrogram.errors import (
     SessionPasswordNeeded,
     VolumeLocNotFound, ChannelPrivate,
@@ -1009,7 +1010,7 @@ class Client(Methods, Scaffold):
                             # https://core.telegram.org/cdn#verifying-files
                             for i, h in enumerate(hashes):
                                 cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)]
-                                assert h.hash == sha256(cdn_chunk).digest(), f"Invalid CDN hash part {i}"
+                                CDNFileHashMismatch.check(h.hash == sha256(cdn_chunk).digest())
 
                             f.write(decrypted_chunk)
 
diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index f9b885f6..ccea119c 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -22,6 +22,7 @@ from io import BytesIO
 from os import urandom
 from typing import List
 
+from pyrogram.errors import SecurityCheckMismatch
 from pyrogram.raw.core import Message, Long
 from . import aes
 from ..session.internals import MsgId
@@ -61,7 +62,7 @@ def unpack(
     auth_key_id: bytes,
     stored_msg_ids: List[int]
 ) -> Message:
-    assert b.read(8) == auth_key_id
+    SecurityCheckMismatch.check(b.read(8) == auth_key_id)
 
     msg_key = b.read(16)
     aes_key, aes_iv = kdf(auth_key, msg_key, False)
@@ -69,7 +70,7 @@ def unpack(
     data.read(8)  # Salt
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-session-id
-    assert data.read(8) == session_id
+    SecurityCheckMismatch.check(data.read(8) == session_id)
 
     try:
         message = Message.read(data)
@@ -87,17 +88,17 @@ def unpack(
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-sha256-hash-value-of-msg-key
     # 96 = 88 + 8 (incoming message)
-    assert msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24]
+    SecurityCheckMismatch.check(msg_key == sha256(auth_key[96:96 + 32] + data.getvalue()).digest()[8:24])
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-message-length
     data.seek(32)  # Get to the payload, skip salt (8) + session_id (8) + msg_id (8) + seq_no (4) + length (4)
     payload = data.read()
     padding = payload[message.length:]
-    assert 12 <= len(padding) <= 1024
-    assert len(payload) % 4 == 0
+    SecurityCheckMismatch.check(12 <= len(padding) <= 1024)
+    SecurityCheckMismatch.check(len(payload) % 4 == 0)
 
     # https://core.telegram.org/mtproto/security_guidelines#checking-msg-id
-    assert message.msg_id % 2 != 0
+    SecurityCheckMismatch.check(message.msg_id % 2 != 0)
 
     if len(stored_msg_ids) > STORED_MSG_IDS_MAX_SIZE:
         del stored_msg_ids[:STORED_MSG_IDS_MAX_SIZE // 2]
@@ -105,21 +106,21 @@ def unpack(
     if stored_msg_ids:
         # Ignored message: msg_id is lower than all of the stored values
         if message.msg_id < stored_msg_ids[0]:
-            assert False
+            SecurityCheckMismatch.check(False)
 
         # Ignored message: msg_id is equal to any of the stored values
         if message.msg_id in stored_msg_ids:
-            assert False
+            SecurityCheckMismatch.check(False)
 
         time_diff = (message.msg_id - MsgId()) / 2 ** 32
 
         # Ignored message: msg_id belongs over 30 seconds in the future
         if time_diff > 30:
-            assert False
+            SecurityCheckMismatch.check(False)
 
         # Ignored message: msg_id belongs over 300 seconds in the past
         if time_diff < -300:
-            assert False
+            SecurityCheckMismatch.check(False)
 
     bisect.insort(stored_msg_ids, message.msg_id)
 
diff --git a/pyrogram/errors/__init__.py b/pyrogram/errors/__init__.py
index 1b94700f..5011b080 100644
--- a/pyrogram/errors/__init__.py
+++ b/pyrogram/errors/__init__.py
@@ -18,3 +18,27 @@
 
 from .exceptions import *
 from .rpc_error import UnknownError
+
+
+class SecurityError(Exception):
+    """Generic security error."""
+
+    @classmethod
+    def check(cls, cond: bool):
+        """Raises this exception if the condition is false"""
+        if not cond:
+            raise cls
+
+
+class SecurityCheckMismatch(SecurityError):
+    """Raised when a security check mismatch occurs."""
+
+    def __init__(self):
+        super().__init__("A security check mismatch has occurred.")
+
+
+class CDNFileHashMismatch(SecurityError):
+    """Raised when a CDN file hash mismatch occurs."""
+
+    def __init__(self):
+        super().__init__("A CDN file hash mismatch has occurred.")
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
index a3e87ff4..6b1ad953 100644
--- a/pyrogram/session/auth.py
+++ b/pyrogram/session/auth.py
@@ -27,6 +27,7 @@ import pyrogram
 from pyrogram import raw
 from pyrogram.connection import Connection
 from pyrogram.crypto import aes, rsa, prime
+from pyrogram.errors import SecurityCheckMismatch
 from pyrogram.raw.core import TLObject, Long, Int
 from .internals import MsgId
 
@@ -210,33 +211,33 @@ class Auth:
                 # Security checks
                 #######################
 
-                assert dh_prime == prime.CURRENT_DH_PRIME
+                SecurityCheckMismatch.check(dh_prime == prime.CURRENT_DH_PRIME)
                 log.debug("DH parameters check: OK")
 
                 # https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
                 g_b = int.from_bytes(g_b, "big")
-                assert 1 < g < dh_prime - 1
-                assert 1 < g_a < dh_prime - 1
-                assert 1 < g_b < dh_prime - 1
-                assert 2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)
-                assert 2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)
+                SecurityCheckMismatch.check(1 < g < dh_prime - 1)
+                SecurityCheckMismatch.check(1 < g_a < dh_prime - 1)
+                SecurityCheckMismatch.check(1 < g_b < dh_prime - 1)
+                SecurityCheckMismatch.check(2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64))
+                SecurityCheckMismatch.check(2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64))
                 log.debug("g_a and g_b validation: OK")
 
                 # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
                 answer = server_dh_inner_data.write()  # Call .write() to remove padding
-                assert answer_with_hash[:20] == sha1(answer).digest()
+                SecurityCheckMismatch.check(answer_with_hash[:20] == sha1(answer).digest())
                 log.debug("SHA1 hash values check: OK")
 
                 # https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
                 # 1st message
-                assert nonce == res_pq.nonce
+                SecurityCheckMismatch.check(nonce == res_pq.nonce)
                 # 2nd message
                 server_nonce = int.from_bytes(server_nonce, "little", signed=True)
-                assert nonce == server_dh_params.nonce
-                assert server_nonce == server_dh_params.server_nonce
+                SecurityCheckMismatch.check(nonce == server_dh_params.nonce)
+                SecurityCheckMismatch.check(server_nonce == server_dh_params.server_nonce)
                 # 3rd message
-                assert nonce == set_client_dh_params_answer.nonce
-                assert server_nonce == set_client_dh_params_answer.server_nonce
+                SecurityCheckMismatch.check(nonce == set_client_dh_params_answer.nonce)
+                SecurityCheckMismatch.check(server_nonce == set_client_dh_params_answer.server_nonce)
                 server_nonce = server_nonce.to_bytes(16, "little", signed=True)
                 log.debug("Nonce fields check: OK")
 
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 3504c3ad..72f27621 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -29,7 +29,10 @@ from pyrogram import __copyright__, __license__, __version__
 from pyrogram import raw
 from pyrogram.connection import Connection
 from pyrogram.crypto import mtproto
-from pyrogram.errors import RPCError, InternalServerError, AuthKeyDuplicated, FloodWait, ServiceUnavailable
+from pyrogram.errors import (
+    RPCError, InternalServerError, AuthKeyDuplicated, FloodWait,
+    ServiceUnavailable, SecurityCheckMismatch
+)
 from pyrogram.raw.all import layer
 from pyrogram.raw.core import TLObject, MsgContainer, Int, FutureSalt, FutureSalts
 from .internals import MsgId, MsgFactory
@@ -230,7 +233,7 @@ class Session:
                 self.auth_key_id,
                 self.stored_msg_ids
             )
-        except AssertionError:
+        except SecurityCheckMismatch:
             self.connection.close()
             return
 

From ea3281b5f62caa916ea6f1ebb2f406bb66a2b36f Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Thu, 16 Dec 2021 21:39:52 +0100
Subject: [PATCH 8/8] Raise directly when not checking a boolean expression

---
 pyrogram/crypto/mtproto.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pyrogram/crypto/mtproto.py b/pyrogram/crypto/mtproto.py
index ccea119c..2fc3b9f8 100644
--- a/pyrogram/crypto/mtproto.py
+++ b/pyrogram/crypto/mtproto.py
@@ -106,21 +106,21 @@ def unpack(
     if stored_msg_ids:
         # Ignored message: msg_id is lower than all of the stored values
         if message.msg_id < stored_msg_ids[0]:
-            SecurityCheckMismatch.check(False)
+            raise SecurityCheckMismatch
 
         # Ignored message: msg_id is equal to any of the stored values
         if message.msg_id in stored_msg_ids:
-            SecurityCheckMismatch.check(False)
+            raise SecurityCheckMismatch
 
         time_diff = (message.msg_id - MsgId()) / 2 ** 32
 
         # Ignored message: msg_id belongs over 30 seconds in the future
         if time_diff > 30:
-            SecurityCheckMismatch.check(False)
+            raise SecurityCheckMismatch
 
         # Ignored message: msg_id belongs over 300 seconds in the past
         if time_diff < -300:
-            SecurityCheckMismatch.check(False)
+            raise SecurityCheckMismatch
 
     bisect.insort(stored_msg_ids, message.msg_id)