From 42a287884258476a4bd3679a24557f5e2e38d8f8 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Mon, 26 Mar 2018 03:39:30 +0200
Subject: [PATCH] Don't use sys.exit(), re-raise ImportError instead

---
 pyrogram/connection/transport/tcp/tcp.py | 7 ++++---
 pyrogram/crypto/aes.py                   | 9 ++++-----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
index 09df8c72..22b953c1 100644
--- a/pyrogram/connection/transport/tcp/tcp.py
+++ b/pyrogram/connection/transport/tcp/tcp.py
@@ -18,17 +18,18 @@
 
 import logging
 import socket
-import sys
 from collections import namedtuple
 
 try:
     import socks
-except ImportError:
-    sys.exit(
+except ImportError as e:
+    e.msg = (
         "PySocks is missing and Pyrogram can't run without. "
         "Please install it using \"pip3 install pysocks\"."
     )
 
+    raise e
+
 log = logging.getLogger(__name__)
 
 Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
diff --git a/pyrogram/crypto/aes.py b/pyrogram/crypto/aes.py
index 331dd5dd..8ca72535 100644
--- a/pyrogram/crypto/aes.py
+++ b/pyrogram/crypto/aes.py
@@ -16,19 +16,18 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
-import sys
-
 try:
     import tgcrypto
-except ImportError:
-    sys.exit(
+except ImportError as e:
+    e.msg = (
         "TgCrypto is missing and Pyrogram can't run without. "
         "Please install it using \"pip3 install tgcrypto\". "
         "More info: https://docs.pyrogram.ml/resources/TgCrypto"
     )
 
+    raise e
+
 
-# TODO: Ugly IFs
 class AES:
     @classmethod
     def ige_encrypt(cls, data: bytes, key: bytes, iv: bytes) -> bytes: