Xtao-Labs/MTPyroger
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add SOCKS5 proxy support

Browse Source
This commit is contained in:
Dan 2018-01-16 22:05:19 +01:00
parent e9f6bce579
commit ba21430512
8 changed files with 59 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
pyrogram/client/client.py
View File

@ -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
)

5
pyrogram/connection/connection.py
View File

@ -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...")

17
pyrogram/connection/transport/tcp/tcp.py
View File

@ -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:

6
pyrogram/connection/transport/tcp/tcp_abridged.py
View File

@ -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

6
pyrogram/connection/transport/tcp/tcp_full.py
View File

@ -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)

6
pyrogram/connection/transport/tcp/tcp_intermediate.py
View File

@ -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)

4
pyrogram/session/auth.py
View File

@ -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:

4
pyrogram/session/session.py
View File

@ -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