From ba214305128abd142d2da1210138c30b2484bb51 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Tue, 16 Jan 2018 22:05:19 +0100 Subject: [PATCH] Add SOCKS5 proxy support --- pyrogram/client/client.py | 38 +++++++++++++------ pyrogram/connection/connection.py | 5 ++- pyrogram/connection/transport/tcp/tcp.py | 17 ++++++++- .../connection/transport/tcp/tcp_abridged.py | 6 +-- pyrogram/connection/transport/tcp/tcp_full.py | 6 +-- .../transport/tcp/tcp_intermediate.py | 6 +-- pyrogram/session/auth.py | 4 +- pyrogram/session/session.py | 4 +- 8 files changed, 59 insertions(+), 27 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index ff2eda7a..3ab86e0d 100644 --- a/pyrogram/client/client.py +++ b/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() - config.read("config.ini") + parser = ConfigParser() + parser.read("config.ini") self.config = Config( - int(config["pyrogram"]["api_id"]), - config["pyrogram"]["api_hash"] + api_id=parser.getint("pyrogram", "api_id"), + 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 ) diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py index 92f8389a..aa958a7a 100644 --- a/pyrogram/connection/connection.py +++ b/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...") diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index 4528a281..2ec899a7 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/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: diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py index 532c1875..f23afe4b 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged.py @@ -24,14 +24,14 @@ log = logging.getLogger(__name__) class TCPAbridged(TCP): - def __init__(self): - super().__init__() + def __init__(self, proxy: type): + 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 diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py index d25905b1..5c1dc2c7 100644 --- a/pyrogram/connection/transport/tcp/tcp_full.py +++ b/pyrogram/connection/transport/tcp/tcp_full.py @@ -26,14 +26,14 @@ log = logging.getLogger(__name__) class TCPFull(TCP): - def __init__(self): - super().__init__() + def __init__(self, proxy: type): + 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) diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate.py b/pyrogram/connection/transport/tcp/tcp_intermediate.py index 537d951e..0ba31c85 100644 --- a/pyrogram/connection/transport/tcp/tcp_intermediate.py +++ b/pyrogram/connection/transport/tcp/tcp_intermediate.py @@ -25,14 +25,14 @@ log = logging.getLogger(__name__) class TCPIntermediate(TCP): - def __init__(self): - super().__init__() + def __init__(self, proxy: type): + 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) diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py index cedf47fb..f3d7a3a3 100644 --- a/pyrogram/session/auth.py +++ b/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: diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index c78da3a7..c399d9ec 100644 --- a/pyrogram/session/session.py +++ b/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