From 532021ab8790285d72b25c98278d844c9199cba1 Mon Sep 17 00:00:00 2001
From: Alexander Bersenev <bay@hackerdom.ru>
Date: Thu, 28 Jun 2018 20:47:12 +0500
Subject: [PATCH] support for cryptography module and advise to use it

---
 mtprotoproxy.py | 61 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 57 insertions(+), 4 deletions(-)

diff --git a/mtprotoproxy.py b/mtprotoproxy.py
index 022eb22..128d7fd 100755
--- a/mtprotoproxy.py
+++ b/mtprotoproxy.py
@@ -20,7 +20,46 @@ try:
 except ImportError:
     pass
 
-try:
+
+def try_use_cryptography_module():
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+    from cryptography.hazmat.backends import default_backend
+
+    def create_aes_ctr(key, iv):
+        class EncryptorAdapter:
+            def __init__(self, cipher):
+                self.encryptor = cipher.encryptor()
+                self.decryptor = cipher.decryptor()
+
+            def encrypt(self, data):
+                return self.encryptor.update(data)
+
+            def decrypt(self, data):
+                return self.decryptor.update(data)
+
+        iv_bytes = int.to_bytes(iv, 16, "big")
+        cipher = Cipher(algorithms.AES(key), modes.CTR(iv_bytes), default_backend())
+        return EncryptorAdapter(cipher)
+
+    def create_aes_cbc(key, iv):
+        class EncryptorAdapter:
+            def __init__(self, cipher):
+                self.encryptor = cipher.encryptor()
+                self.decryptor = cipher.decryptor()
+
+            def encrypt(self, data):
+                return self.encryptor.update(data)
+
+            def decrypt(self, data):
+                return self.decryptor.update(data)
+
+        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), default_backend())
+        return EncryptorAdapter(cipher)
+
+    return create_aes_ctr, create_aes_cbc
+
+
+def try_use_pycrypto_or_pycryptodome_module():
     from Crypto.Cipher import AES
     from Crypto.Util import Counter
 
@@ -31,11 +70,16 @@ try:
     def create_aes_cbc(key, iv):
         return AES.new(key, AES.MODE_CBC, iv)
 
-except ImportError:
-    print("Failed to find pycryptodome or pycrypto, using slow AES implementation",
-          flush=True, file=sys.stderr)
+    return create_aes_ctr, create_aes_cbc
+
+
+def use_slow_bundled_cryptography_module():
     import pyaes
 
+    msg = "To make the program a *lot* faster, please install cryptography module: "
+    msg += "pip install cryptography\n"
+    print(msg, flush=True, file=sys.stderr)
+
     def create_aes_ctr(key, iv):
         ctr = pyaes.Counter(iv)
         return pyaes.AESModeOfOperationCTR(key, ctr)
@@ -55,8 +99,17 @@ except ImportError:
 
         mode = pyaes.AESModeOfOperationCBC(key, iv)
         return EncryptorAdapter(mode)
+    return create_aes_ctr, create_aes_cbc
 
 
+try:
+    create_aes_ctr, create_aes_cbc = try_use_cryptography_module()
+except ImportError:
+    try:
+        create_aes_ctr, create_aes_cbc = try_use_pycrypto_or_pycryptodome_module()
+    except ImportError:
+        create_aes_ctr, create_aes_cbc = use_slow_bundled_cryptography_module()
+
 try:
     import resource
     soft_fd_limit, hard_fd_limit = resource.getrlimit(resource.RLIMIT_NOFILE)