From f993917c35d2207b72fb0f3e4240b476de84438a Mon Sep 17 00:00:00 2001
From: igerzog <igerzog@users.noreply.github.com>
Date: Tue, 2 Mar 2021 21:59:12 +0300
Subject: [PATCH] faster download and upload

---
 telethon/crypto/aes.py    | 28 +++++++++++++++++++---------
 telethon/crypto/aesctr.py | 30 +++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/telethon/crypto/aes.py b/telethon/crypto/aes.py
index 3cfcc1af..7bd2fe1a 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, tgcrypto 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.
 """
@@ -15,15 +16,20 @@ __log__ = logging.getLogger(__name__)
 
 
 try:
-    import cryptg
-    __log__.info('cryptg detected, it will be used for encryption')
+    import tgcrypto
+    __log__.debug('tgcrypto 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')
+    tgcrypto = None
+    try:
+        import cryptg
+        __log__.debug('cryptg detected, it will be used for encryption')
+    except ImportError:
+        cryptg = None
+        if libssl.encrypt_ige and libssl.decrypt_ige:
+            __log__.debug('libssl detected, it will be used for encryption')
+        else:
+            __log__.debug('tgcrypto 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 tgcrypto:
+            return tgcrypto.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 tgcrypto:
+            return tgcrypto.ige256_encrypt(plain_text, key, iv)
         if cryptg:
             return cryptg.encrypt_ige(plain_text, key, iv)
         if libssl.encrypt_ige:
diff --git a/telethon/crypto/aesctr.py b/telethon/crypto/aesctr.py
index 34422904..d1f5947f 100644
--- a/telethon/crypto/aesctr.py
+++ b/telethon/crypto/aesctr.py
@@ -2,6 +2,19 @@
 This module holds the AESModeCTR wrapper class.
 """
 import pyaes
+import logging
+
+
+__log__ = logging.getLogger(__name__)
+
+
+try:
+    import tgcrypto
+    __log__.debug('tgcrypto detected, it will be used for ctr encryption')
+except ImportError:
+    tgcrypto = None
+    __log__.debug('tgcrypto module not installed, '
+                'falling back to (slower) Python encryption')
 
 
 class AESModeCTR:
@@ -16,12 +29,15 @@ class AESModeCTR:
         :param iv: the bytes initialization vector. Must have a length of 16.
         """
         # TODO Use libssl if available
-        assert isinstance(key, bytes)
-        self._aes = pyaes.AESModeOfOperationCTR(key)
+        if tgcrypto:
+            self._aes = (key, iv, bytearray(1))
+        else:
+            assert isinstance(key, bytes)
+            self._aes = pyaes.AESModeOfOperationCTR(key)
 
-        assert isinstance(iv, bytes)
-        assert len(iv) == 16
-        self._aes._counter._counter = list(iv)
+            assert isinstance(iv, bytes)
+            assert len(iv) == 16
+            self._aes._counter._counter = list(iv)
 
     def encrypt(self, data):
         """
@@ -30,6 +46,8 @@ class AESModeCTR:
         :param data: the plain text to be encrypted.
         :return: the encrypted cipher text.
         """
+        if tgcrypto:
+            return tgcrypto.ctr256_encrypt(data, *self._aes)
         return self._aes.encrypt(data)
 
     def decrypt(self, data):
@@ -39,4 +57,6 @@ class AESModeCTR:
         :param data: the cipher text to be decrypted.
         :return: the decrypted plain text.
         """
+        if tgcrypto:
+            return tgcrypto.ctr256_decrypt(data, *self._aes)
         return self._aes.decrypt(data)