From 0aed7bf24ae2b2c9801392138dcfa812b67d3df6 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Tue, 16 Jan 2018 16:26:20 +0100 Subject: [PATCH 1/8] Use socks socket to allow enabling a proxy --- pyrogram/connection/transport/tcp/tcp.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index e7d635ae..11990202 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -19,10 +19,12 @@ import logging import socket +import socks + log = logging.getLogger(__name__) -class TCP(socket.socket): +class TCP(socks.socksocket): def __init__(self): super().__init__() From e9f6bce5795e2ee881c1e23ad1856e16f295f2ad Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Tue, 16 Jan 2018 16:28:50 +0100 Subject: [PATCH 2/8] Reformat connection classes to accommodate proxy settings --- pyrogram/connection/connection.py | 4 ++-- pyrogram/connection/transport/tcp/tcp.py | 6 ------ pyrogram/connection/transport/tcp/tcp_abridged.py | 10 +++++----- pyrogram/connection/transport/tcp/tcp_full.py | 8 ++++---- pyrogram/connection/transport/tcp/tcp_intermediate.py | 8 ++++---- 5 files changed, 15 insertions(+), 21 deletions(-) diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py index 429d5520..92f8389a 100644 --- a/pyrogram/connection/connection.py +++ b/pyrogram/connection/connection.py @@ -56,7 +56,7 @@ class Connection: def send(self, data: bytes): with self.lock: - self.connection.send(data) + self.connection.sendall(data) def recv(self) -> bytes or None: - return self.connection.recv() + return self.connection.recvall() diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index 11990202..4528a281 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -28,12 +28,6 @@ class TCP(socks.socksocket): def __init__(self): super().__init__() - def send(self, *args): - pass - - def recv(self, *args): - pass - def close(self): try: self.shutdown(socket.SHUT_RDWR) diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py index 33b4a507..532c1875 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged.py @@ -33,7 +33,7 @@ class TCPAbridged(TCP): self.is_first_packet = True log.info("Connected!") - def send(self, data: bytes): + def sendall(self, data: bytes, *args): length = len(data) // 4 data = ( @@ -48,20 +48,20 @@ class TCPAbridged(TCP): super().sendall(data) - def recv(self) -> bytes or None: - length = self.recvall(1) + def recvall(self, length: int = 0) -> bytes or None: + length = super().recvall(1) if length is None: return None if length == b"\x7f": - length = self.recvall(3) + length = super().recvall(3) if length is None: return None length = int.from_bytes(length, "little") * 4 - packet = self.recvall(length) + packet = super().recvall(length) return packet diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py index 8b932eee..d25905b1 100644 --- a/pyrogram/connection/transport/tcp/tcp_full.py +++ b/pyrogram/connection/transport/tcp/tcp_full.py @@ -35,7 +35,7 @@ class TCPFull(TCP): self.seq_no = 0 log.info("Connected!") - def send(self, data: bytes): + def sendall(self, data: bytes, *args): # 12 = packet_length (4), seq_no (4), crc32 (4) (at the end) data = pack(" bytes or None: - length = self.recvall(4) + def recvall(self, length: int = 0) -> bytes or None: + length = super().recvall(4) if length is None: return None - packet = self.recvall(unpack(" bytes or None: - length = self.recvall(4) + def recvall(self, length: int = 0) -> bytes or None: + length = super().recvall(4) if length is None: return None - packet = self.recvall(unpack(" Date: Tue, 16 Jan 2018 22:05:19 +0100 Subject: [PATCH 3/8] 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 From 789cffa4f7d1ce06a153c2827df85597e3ff818c Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Tue, 16 Jan 2018 22:06:04 +0100 Subject: [PATCH 4/8] Remove print call --- pyrogram/client/client.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 3ab86e0d..ba7470a4 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -309,8 +309,6 @@ class Client: 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: From ed1e86dce15c602f60d283e1f4b311e98d26c701 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 17 Jan 2018 00:10:57 +0100 Subject: [PATCH 5/8] Add pysocks to required packages --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index fd96b3b4..7dfb86cf 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,5 +26,5 @@ classifiers = [options] packages = find: zip_safe = False -install_requires = pyaes +install_requires = pyaes; pysocks include_package_data = True \ No newline at end of file From bbad1f774ebbc9ae49251f9b221757ae931f3363 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 17 Jan 2018 00:11:51 +0100 Subject: [PATCH 6/8] Add setup_requires entry --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 7dfb86cf..bc330f79 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,5 +26,6 @@ classifiers = [options] packages = find: zip_safe = False +setup_requires = pyaes; pysocks install_requires = pyaes; pysocks include_package_data = True \ No newline at end of file From 3686f6921de22b3972c5a21c294e3eb5ef56fd61 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 17 Jan 2018 00:16:23 +0100 Subject: [PATCH 7/8] Remove try..except block, now pyaes is made available during setup --- pyrogram/crypto/ige.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pyrogram/crypto/ige.py b/pyrogram/crypto/ige.py index 18d7a09a..da0a2398 100644 --- a/pyrogram/crypto/ige.py +++ b/pyrogram/crypto/ige.py @@ -16,10 +16,7 @@ # You should have received a copy of the GNU Lesser General Public License # along with Pyrogram. If not, see . -try: - from pyaes import AES -except ImportError: - pass +from pyaes import AES BLOCK_SIZE = 16 From 4784ba53e1b226940be4bd20afc3925c82a3692b Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Thu, 18 Jan 2018 12:22:16 +0100 Subject: [PATCH 8/8] Remove redundant code --- pyrogram/connection/transport/tcp/tcp_abridged.py | 6 +----- pyrogram/connection/transport/tcp/tcp_intermediate.py | 4 +--- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py index f23afe4b..acb837af 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged.py @@ -60,8 +60,4 @@ class TCPAbridged(TCP): if length is None: return None - length = int.from_bytes(length, "little") * 4 - - packet = super().recvall(length) - - return packet + return super().recvall(int.from_bytes(length, "little") * 4) diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate.py b/pyrogram/connection/transport/tcp/tcp_intermediate.py index 0ba31c85..301a88f6 100644 --- a/pyrogram/connection/transport/tcp/tcp_intermediate.py +++ b/pyrogram/connection/transport/tcp/tcp_intermediate.py @@ -50,6 +50,4 @@ class TCPIntermediate(TCP): if length is None: return None - packet = super().recvall(unpack("