Use tgcrypto if available (#1715)

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.
 """
@@ -14,15 +15,20 @@ from . import libssl
 __log__ = logging.getLogger(__name__)
 
 
+try:
+    import tgcrypto
+    __log__.debug('tgcrypto detected, it will be used for encryption')
+except ImportError:
+    tgcrypto = None
     try:
         import cryptg
-    __log__.info('cryptg detected, it will be used for encryption')
+        __log__.debug('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')
+            __log__.debug('libssl detected, it will be used for encryption')
         else:
-        __log__.info('cryptg module not installed and libssl not found, '
+            __log__.debug('tgcrypto or cryptg modules not installed and libssl not found, '
                         'falling back to (slower) Python encryption')
 
 
@@ -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:

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,6 +29,9 @@ class AESModeCTR:
         :param iv: the bytes initialization vector. Must have a length of 16.
         """
         # TODO Use libssl if available
+        if tgcrypto:
+            self._aes = (key, iv, bytearray(1))
+        else:
             assert isinstance(key, bytes)
             self._aes = pyaes.AESModeOfOperationCTR(key)
 
@@ -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)