Add SOCKS5 proxy support

pyrogram/client/client.py

@@ -53,6 +53,7 @@ from pyrogram.session import Auth, Session
 log = logging.getLogger(__name__)
 
 Config = namedtuple("Config", ["api_id", "api_hash"])
+Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
 
 
 class Client:
@@ -90,6 +91,7 @@ class Client:
         self.markdown = Markdown(self.peers_by_id)
 
         self.config = None
+        self.proxy = None
         self.session = None
 
         self.update_handler = None
@@ -101,7 +103,7 @@ class Client:
         self.load_config()
         self.load_session(self.session_name)
 
-        self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
+        self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
 
         terms = self.session.start()
 
@@ -191,9 +193,9 @@ class Client:
                 self.session.stop()
 
                 self.dc_id = e.x
-                self.auth_key = Auth(self.dc_id, self.test_mode).create()
+                self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
 
-                self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
+                self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
                 self.session.start()
 
                 r = self.send(
@@ -290,21 +292,32 @@ class Client:
         return r.user.id
 
     def load_config(self):
-        config = ConfigParser()
+        parser = ConfigParser()
-        config.read("config.ini")
+        parser.read("config.ini")
 
         self.config = Config(
-            int(config["pyrogram"]["api_id"]),
+            api_id=parser.getint("pyrogram", "api_id"),
-            config["pyrogram"]["api_hash"]
+            api_hash=parser.get("pyrogram", "api_hash")
         )
 
+        if parser.has_section("proxy"):
+            self.proxy = Proxy(
+                enabled=parser.getboolean("proxy", "enabled"),
+                hostname=parser.get("proxy", "hostname"),
+                port=parser.getint("proxy", "port"),
+                username=parser.get("proxy", "username", fallback=None) or None,
+                password=parser.get("proxy", "password", fallback=None) or None
+            )
+
+        print(self.proxy)
+
     def load_session(self, session_name):
         try:
             with open("{}.session".format(session_name)) as f:
                 s = json.load(f)
         except FileNotFoundError:
             self.dc_id = 1
-            self.auth_key = Auth(self.dc_id, self.test_mode).create()
+            self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
         else:
             self.dc_id = s["dc_id"]
             self.test_mode = s["test_mode"]
@@ -1297,7 +1310,7 @@ class Client:
         file_id = file_id or self.rnd_id()
         md5_sum = md5() if not is_big and not is_missing_part else None
 
-        session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
+        session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
         session.start()
 
         try:
@@ -1362,7 +1375,8 @@ class Client:
             session = Session(
                 dc_id,
                 self.test_mode,
-                Auth(dc_id, self.test_mode).create(),
+                self.proxy,
+                Auth(dc_id, self.test_mode, self.proxy).create(),
                 self.config.api_id
             )
 
@@ -1378,6 +1392,7 @@ class Client:
             session = Session(
                 dc_id,
                 self.test_mode,
+                self.proxy,
                 self.auth_key,
                 self.config.api_id
             )
@@ -1433,7 +1448,8 @@ class Client:
                 cdn_session = Session(
                     r.dc_id,
                     self.test_mode,
-                    Auth(r.dc_id, self.test_mode).create(),
+                    self.proxy,
+                    Auth(r.dc_id, self.test_mode, self.proxy).create(),
                     self.config.api_id,
                     is_cdn=True
                 )

pyrogram/connection/connection.py

@@ -32,15 +32,16 @@ class Connection:
         2: TCPIntermediate
     }
 
-    def __init__(self, ipv4: str, mode: int = 1):
+    def __init__(self, ipv4: str, proxy: type, mode: int = 1):
         self.address = (ipv4, 80)
+        self.proxy = proxy
         self.mode = self.MODES.get(mode, TCPAbridged)
         self.lock = threading.Lock()
         self.connection = None
 
     def connect(self):
         while True:
-            self.connection = self.mode()
+            self.connection = self.mode(self.proxy)
 
             try:
                 log.info("Connecting...")

pyrogram/connection/transport/tcp/tcp.py

@@ -18,15 +18,30 @@
 
 import logging
 import socket
+from collections import namedtuple
 
 import socks
 
 log = logging.getLogger(__name__)
 
+Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
+
 
 class TCP(socks.socksocket):
-    def __init__(self):
+    def __init__(self, proxy: Proxy):
         super().__init__()
+        self.proxy_enabled = False
+
+        if proxy and proxy.enabled:
+            self.proxy_enabled = True
+
+            self.set_proxy(
+                proxy_type=socks.SOCKS5,
+                addr=proxy.hostname,
+                port=proxy.port,
+                username=proxy.username,
+                password=proxy.password
+            )
 
     def close(self):
         try:

pyrogram/connection/transport/tcp/tcp_abridged.py

@@ -24,14 +24,14 @@ log = logging.getLogger(__name__)
 
 
 class TCPAbridged(TCP):
-    def __init__(self):
+    def __init__(self, proxy: type):
-        super().__init__()
+        super().__init__(proxy)
         self.is_first_packet = None
 
     def connect(self, address: tuple):
         super().connect(address)
         self.is_first_packet = True
-        log.info("Connected!")
+        log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
 
     def sendall(self, data: bytes, *args):
         length = len(data) // 4

pyrogram/connection/transport/tcp/tcp_full.py

@@ -26,14 +26,14 @@ log = logging.getLogger(__name__)
 
 
 class TCPFull(TCP):
-    def __init__(self):
+    def __init__(self, proxy: type):
-        super().__init__()
+        super().__init__(proxy)
         self.seq_no = None
 
     def connect(self, address: tuple):
         super().connect(address)
         self.seq_no = 0
-        log.info("Connected!")
+        log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
 
     def sendall(self, data: bytes, *args):
         # 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)

pyrogram/connection/transport/tcp/tcp_intermediate.py

@@ -25,14 +25,14 @@ log = logging.getLogger(__name__)
 
 
 class TCPIntermediate(TCP):
-    def __init__(self):
+    def __init__(self, proxy: type):
-        super().__init__()
+        super().__init__(proxy)
         self.is_first_packet = None
 
     def connect(self, address: tuple):
         super().connect(address)
         self.is_first_packet = True
-        log.info("Connected!")
+        log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
 
     def sendall(self, data: bytes, *args):
         length = len(data)

pyrogram/session/auth.py

@@ -46,11 +46,11 @@ class Auth:
         16
     )
 
-    def __init__(self, dc_id: int, test_mode: bool):
+    def __init__(self, dc_id: int, test_mode: bool, proxy: type):
         self.dc_id = dc_id
         self.test_mode = test_mode
 
-        self.connection = Connection(DataCenter(dc_id, test_mode))
+        self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
         self.msg_id = MsgId()
 
     def pack(self, data: Object) -> bytes:

pyrogram/session/session.py

@@ -68,7 +68,7 @@ class Session:
 
     notice_displayed = False
 
-    def __init__(self, dc_id: int, test_mode: bool, auth_key: bytes, api_id: str, is_cdn: bool = False):
+    def __init__(self, dc_id: int, test_mode: bool, proxy: type, auth_key: bytes, api_id: str, is_cdn: bool = False):
         if not Session.notice_displayed:
             print("Pyrogram v{}, {}".format(__version__, __copyright__))
             print("Licensed under the terms of the " + __license__, end="\n\n")
@@ -76,7 +76,7 @@ class Session:
 
         self.is_cdn = is_cdn
 
-        self.connection = Connection(DataCenter(dc_id, test_mode))
+        self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
 
         self.api_id = api_id