From b4737e52f3c832cc76cd2058188a2d6dfc3f1598 Mon Sep 17 00:00:00 2001
From: painor <pi.oussama@gmail.com>
Date: Thu, 1 Apr 2021 14:32:08 +0100
Subject: [PATCH] Add pyaesni support when found

---
 optional-requirements.txt |  1 +
 telethon/crypto/aes.py    | 32 +++++++++++++++++++++-----------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/optional-requirements.txt b/optional-requirements.txt
index 30326da9..aefb1ba7 100644
--- a/optional-requirements.txt
+++ b/optional-requirements.txt
@@ -1,3 +1,4 @@
+pyaesni
 cryptg
 pysocks
 python-socks[asyncio]
diff --git a/telethon/crypto/aes.py b/telethon/crypto/aes.py
index 3cfcc1af..30a86120 100644
--- a/telethon/crypto/aes.py
+++ b/telethon/crypto/aes.py
@@ -1,7 +1,8 @@
 """
 AES IGE implementation in Python.
 
-If available, cryptg will be used instead, otherwise
+If available, pyaesni will be used instead, otherwise
+if available, cryptg will be used instead, otherwise
 if available, libssl will be used instead, otherwise
 the Python implementation will be used.
 """
@@ -10,20 +11,25 @@ import pyaes
 import logging
 from . import libssl
 
-
 __log__ = logging.getLogger(__name__)
 
-
 try:
-    import cryptg
-    __log__.info('cryptg detected, it will be used for encryption')
+    import pyaesni
+
+    __log__.info('pyaesni detected, it will be used for encryption')
 except ImportError:
-    cryptg = None
-    if libssl.encrypt_ige and libssl.decrypt_ige:
-        __log__.info('libssl detected, it will be used for encryption')
-    else:
-        __log__.info('cryptg module not installed and libssl not found, '
-                     'falling back to (slower) Python encryption')
+    pyaesni = None
+    try:
+        import cryptg
+
+        __log__.info('cryptg detected, it will be used for encryption')
+    except ImportError:
+        cryptg = None
+        if libssl.encrypt_ige and libssl.decrypt_ige:
+            __log__.info('libssl detected, it will be used for encryption')
+        else:
+            __log__.info('pyaesni or cryptg modules not installed and libssl not found, '
+                          'falling back to (slower) Python encryption')
 
 
 class AES:
@@ -37,6 +43,8 @@ class AES:
         Decrypts the given text in 16-bytes blocks by using the
         given key and 32-bytes initialization vector.
         """
+        if pyaesni:
+            return pyaesni.ige256_decrypt(cipher_text, key, iv)
         if cryptg:
             return cryptg.decrypt_ige(cipher_text, key, iv)
         if libssl.decrypt_ige:
@@ -78,6 +86,8 @@ class AES:
         if padding:
             plain_text += os.urandom(16 - padding)
 
+        if pyaesni:
+            return pyaesni.ige256_encrypt(plain_text, key, iv)
         if cryptg:
             return cryptg.encrypt_ige(plain_text, key, iv)
         if libssl.encrypt_ige: