diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index ff2eda7a..ba7470a4 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,30 @@ 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 + ) + 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 +1308,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 +1373,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 +1390,7 @@ class Client: session = Session( dc_id, self.test_mode, + self.proxy, self.auth_key, self.config.api_id ) @@ -1433,7 +1446,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 429d5520..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...") @@ -56,7 +57,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 e7d635ae..2ec899a7 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -18,19 +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(socket.socket): - def __init__(self): + +class TCP(socks.socksocket): + def __init__(self, proxy: Proxy): super().__init__() + self.proxy_enabled = False - def send(self, *args): - pass + if proxy and proxy.enabled: + self.proxy_enabled = True - def recv(self, *args): - pass + 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 33b4a507..acb837af 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged.py @@ -24,16 +24,16 @@ 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 send(self, data: bytes): + def sendall(self, data: bytes, *args): length = len(data) // 4 data = ( @@ -48,20 +48,16 @@ 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) - - return packet + return super().recvall(int.from_bytes(length, "little") * 4) diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py index 8b932eee..5c1dc2c7 100644 --- a/pyrogram/connection/transport/tcp/tcp_full.py +++ b/pyrogram/connection/transport/tcp/tcp_full.py @@ -26,16 +26,16 @@ 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 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(". -try: - from pyaes import AES -except ImportError: - pass +from pyaes import AES BLOCK_SIZE = 16 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 diff --git a/setup.cfg b/setup.cfg index fd96b3b4..bc330f79 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,5 +26,6 @@ classifiers = [options] packages = find: zip_safe = False -install_requires = pyaes +setup_requires = pyaes; pysocks +install_requires = pyaes; pysocks include_package_data = True \ No newline at end of file