Add SOCKS5 proxy support · Problem-2/pyrogram@ba21430 · GitHub
Skip to content

Commit ba21430

Browse files
committed
Add SOCKS5 proxy support
1 parent e9f6bce commit ba21430

8 files changed

Lines changed: 59 additions & 27 deletions

File tree

pyrogram/client/client.py

Lines changed: 27 additions & 11 deletions

pyrogram/connection/connection.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,16 @@ class Connection:
3232
2: TCPIntermediate
3333
}
3434

35-
def __init__(self, ipv4: str, mode: int = 1):
35+
def __init__(self, ipv4: str, proxy: type, mode: int = 1):
3636
self.address = (ipv4, 80)
37+
self.proxy = proxy
3738
self.mode = self.MODES.get(mode, TCPAbridged)
3839
self.lock = threading.Lock()
3940
self.connection = None
4041

4142
def connect(self):
4243
while True:
43-
self.connection = self.mode()
44+
self.connection = self.mode(self.proxy)
4445

4546
try:
4647
log.info("Connecting...")

pyrogram/connection/transport/tcp/tcp.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,30 @@
1818

1919
import logging
2020
import socket
21+
from collections import namedtuple
2122

2223
import socks
2324

2425
log = logging.getLogger(__name__)
2526

27+
Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
28+
2629

2730
class TCP(socks.socksocket):
28-
def __init__(self):
31+
def __init__(self, proxy: Proxy):
2932
super().__init__()
33+
self.proxy_enabled = False
34+
35+
if proxy and proxy.enabled:
36+
self.proxy_enabled = True
37+
38+
self.set_proxy(
39+
proxy_type=socks.SOCKS5,
40+
addr=proxy.hostname,
41+
port=proxy.port,
42+
username=proxy.username,
43+
password=proxy.password
44+
)
3045

3146
def close(self):
3247
try:

pyrogram/connection/transport/tcp/tcp_abridged.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,14 @@
2424

2525

2626
class TCPAbridged(TCP):
27-
def __init__(self):
28-
super().__init__()
27+
def __init__(self, proxy: type):
28+
super().__init__(proxy)
2929
self.is_first_packet = None
3030

3131
def connect(self, address: tuple):
3232
super().connect(address)
3333
self.is_first_packet = True
34-
log.info("Connected!")
34+
log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
3535

3636
def sendall(self, data: bytes, *args):
3737
length = len(data) // 4

pyrogram/connection/transport/tcp/tcp_full.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,14 @@
2626

2727

2828
class TCPFull(TCP):
29-
def __init__(self):
30-
super().__init__()
29+
def __init__(self, proxy: type):
30+
super().__init__(proxy)
3131
self.seq_no = None
3232

3333
def connect(self, address: tuple):
3434
super().connect(address)
3535
self.seq_no = 0
36-
log.info("Connected!")
36+
log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
3737

3838
def sendall(self, data: bytes, *args):
3939
# 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)

pyrogram/connection/transport/tcp/tcp_intermediate.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,14 @@
2525

2626

2727
class TCPIntermediate(TCP):
28-
def __init__(self):
29-
super().__init__()
28+
def __init__(self, proxy: type):
29+
super().__init__(proxy)
3030
self.is_first_packet = None
3131

3232
def connect(self, address: tuple):
3333
super().connect(address)
3434
self.is_first_packet = True
35-
log.info("Connected!")
35+
log.info("Connected{}!".format(" with proxy" if self.proxy_enabled else ""))
3636

3737
def sendall(self, data: bytes, *args):
3838
length = len(data)

pyrogram/session/auth.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,11 @@ class Auth:
4646
16
4747
)
4848

49-
def __init__(self, dc_id: int, test_mode: bool):
49+
def __init__(self, dc_id: int, test_mode: bool, proxy: type):
5050
self.dc_id = dc_id
5151
self.test_mode = test_mode
5252

53-
self.connection = Connection(DataCenter(dc_id, test_mode))
53+
self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
5454
self.msg_id = MsgId()
5555

5656
def pack(self, data: Object) -> bytes:

pyrogram/session/session.py

Lines changed: 2 additions & 2 deletions

0 commit comments

Comments
 (0)