Skip to content

Commit ba21430

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

File tree

8 files changed

+59
-27
lines changed

8 files changed

+59
-27
lines changed

‎pyrogram/client/client.py‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
log = logging.getLogger(__name__)
5454

5555
Config = namedtuple("Config", ["api_id", "api_hash"])
56+
Proxy = namedtuple("Proxy", ["enabled", "hostname", "port", "username", "password"])
5657

5758

5859
class Client:
@@ -90,6 +91,7 @@ def __init__(self, session_name: str, test_mode: bool = False):
9091
self.markdown = Markdown(self.peers_by_id)
9192

9293
self.config = None
94+
self.proxy = None
9395
self.session = None
9496

9597
self.update_handler = None
@@ -101,7 +103,7 @@ def start(self):
101103
self.load_config()
102104
self.load_session(self.session_name)
103105

104-
self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
106+
self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
105107

106108
terms = self.session.start()
107109

@@ -191,9 +193,9 @@ def authorize(self):
191193
self.session.stop()
192194

193195
self.dc_id = e.x
194-
self.auth_key = Auth(self.dc_id, self.test_mode).create()
196+
self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
195197

196-
self.session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
198+
self.session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
197199
self.session.start()
198200

199201
r = self.send(
@@ -290,21 +292,32 @@ def authorize(self):
290292
return r.user.id
291293

292294
def load_config(self):
293-
config = ConfigParser()
294-
config.read("config.ini")
295+
parser = ConfigParser()
296+
parser.read("config.ini")
295297

296298
self.config = Config(
297-
int(config["pyrogram"]["api_id"]),
298-
config["pyrogram"]["api_hash"]
299+
api_id=parser.getint("pyrogram", "api_id"),
300+
api_hash=parser.get("pyrogram", "api_hash")
299301
)
300302

303+
if parser.has_section("proxy"):
304+
self.proxy = Proxy(
305+
enabled=parser.getboolean("proxy", "enabled"),
306+
hostname=parser.get("proxy", "hostname"),
307+
port=parser.getint("proxy", "port"),
308+
username=parser.get("proxy", "username", fallback=None) or None,
309+
password=parser.get("proxy", "password", fallback=None) or None
310+
)
311+
312+
print(self.proxy)
313+
301314
def load_session(self, session_name):
302315
try:
303316
with open("{}.session".format(session_name)) as f:
304317
s = json.load(f)
305318
except FileNotFoundError:
306319
self.dc_id = 1
307-
self.auth_key = Auth(self.dc_id, self.test_mode).create()
320+
self.auth_key = Auth(self.dc_id, self.test_mode, self.proxy).create()
308321
else:
309322
self.dc_id = s["dc_id"]
310323
self.test_mode = s["test_mode"]
@@ -1297,7 +1310,7 @@ def save_file(self, path: str, file_id: int = None, file_part: int = 0):
12971310
file_id = file_id or self.rnd_id()
12981311
md5_sum = md5() if not is_big and not is_missing_part else None
12991312

1300-
session = Session(self.dc_id, self.test_mode, self.auth_key, self.config.api_id)
1313+
session = Session(self.dc_id, self.test_mode, self.proxy, self.auth_key, self.config.api_id)
13011314
session.start()
13021315

13031316
try:
@@ -1362,7 +1375,8 @@ def get_file(self,
13621375
session = Session(
13631376
dc_id,
13641377
self.test_mode,
1365-
Auth(dc_id, self.test_mode).create(),
1378+
self.proxy,
1379+
Auth(dc_id, self.test_mode, self.proxy).create(),
13661380
self.config.api_id
13671381
)
13681382

@@ -1378,6 +1392,7 @@ def get_file(self,
13781392
session = Session(
13791393
dc_id,
13801394
self.test_mode,
1395+
self.proxy,
13811396
self.auth_key,
13821397
self.config.api_id
13831398
)
@@ -1433,7 +1448,8 @@ def get_file(self,
14331448
cdn_session = Session(
14341449
r.dc_id,
14351450
self.test_mode,
1436-
Auth(r.dc_id, self.test_mode).create(),
1451+
self.proxy,
1452+
Auth(r.dc_id, self.test_mode, self.proxy).create(),
14371453
self.config.api_id,
14381454
is_cdn=True
14391455
)

‎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
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,15 @@ class Session:
6868

6969
notice_displayed = False
7070

71-
def __init__(self, dc_id: int, test_mode: bool, auth_key: bytes, api_id: str, is_cdn: bool = False):
71+
def __init__(self, dc_id: int, test_mode: bool, proxy: type, auth_key: bytes, api_id: str, is_cdn: bool = False):
7272
if not Session.notice_displayed:
7373
print("Pyrogram v{}, {}".format(__version__, __copyright__))
7474
print("Licensed under the terms of the " + __license__, end="\n\n")
7575
Session.notice_displayed = True
7676

7777
self.is_cdn = is_cdn
7878

79-
self.connection = Connection(DataCenter(dc_id, test_mode))
79+
self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
8080

8181
self.api_id = api_id
8282

0 commit comments

Comments
 (0)