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.
 """
@@ -15,15 +16,20 @@
 
 
 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 @@ def decrypt_ige(cipher_text, key, iv):
         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 @@ def encrypt_ige(plain_text, key, iv):
         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,12 +29,15 @@ def __init__(self, key, iv):
         :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 @@ def encrypt(self, data):
         :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 @@ def decrypt(self, data):
         :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)