From 268e6643d07e4fa429726cc350ea4a6ce91ed991 Mon Sep 17 00:00:00 2001
From: BlackCatDevel0per <bcdev@mail.ru>
Date: Wed, 23 Mar 2022 20:48:52 +0500
Subject: [PATCH] Added tgcrypto lib from pyrogram project

---
 telethon/_crypto/aes.py | 43 ++++++++++++++++++++++++++++++++---------
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/telethon/_crypto/aes.py b/telethon/_crypto/aes.py
index 3cfcc1af..61e482b1 100644
--- a/telethon/_crypto/aes.py
+++ b/telethon/_crypto/aes.py
@@ -2,6 +2,7 @@
 AES IGE implementation in Python.
 
 If available, cryptg will be used instead, otherwise
+If available, pyaesni(intel aes-ni for python) will be used instead, otherwise
 if available, libssl will be used instead, otherwise
 the Python implementation will be used.
 """
@@ -12,19 +13,35 @@ from . import libssl
 
 
 __log__ = logging.getLogger(__name__)
-
+# tgcrypto
+try:
+    import tgcrypto
+    __log__.info('pyaesni detected, it will be used for encryption')
+except ImportError:
+    tgcrypto = None
+    __log__.info('tgcrypto module not installed and libssl not found, '
+                 'falling back to (slower) Python encryption')
+# pyaesni (intel aes-ni)
+try:
+    if not tgcrypto:
+        import pyaesni
+        __log__.info('pyaesni detected, it will be used for encryption')
+except ImportError:
+    pyaesni = None
+    __log__.info('pyaesni module not installed and libssl not found, '
+                 'falling back to (slower) Python encryption')
 
 try:
-    import cryptg
-    __log__.info('cryptg detected, it will be used for encryption')
+    if not tgcrypto and not pyaesni:
+      import cryptg
+      __log__.info('cryptg detected, it will be used for encryption')
 except ImportError:
     cryptg = None
+    __log__.info('cryptg module not installed and libssl not found, '
+                 'falling back to (slower) Python encryption')
+if not tgcrypto and not pyaesni and not cryptg:
     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')
-
 
 class AES:
     """
@@ -37,8 +54,12 @@ class AES:
         Decrypts the given text in 16-bytes blocks by using the
         given key and 32-bytes initialization vector.
         """
-        if cryptg:
+        if tgcrypto:
+            return tgcrypto.ige256_decrypt(cipher_text, key, iv)
+        elif cryptg:
             return cryptg.decrypt_ige(cipher_text, key, iv)
+        elif pyaesni:
+            return pyaesni.ige256_decrypt(cipher_text, key, iv)
         if libssl.decrypt_ige:
             return libssl.decrypt_ige(cipher_text, key, iv)
 
@@ -78,8 +99,12 @@ class AES:
         if padding:
             plain_text += os.urandom(16 - padding)
 
-        if cryptg:
+        if tgcrypto:
+            return tgcrypto.ige256_encrypt(plain_text, key, iv)
+        elif cryptg:
             return cryptg.encrypt_ige(plain_text, key, iv)
+        elif pyaesni:
+            return pyaesni.ige256_encrypt(plain_text, key, iv)
         if libssl.encrypt_ige:
             return libssl.encrypt_ige(plain_text, key, iv)