Connection refactoring

This commit is contained in:
Dan 2018-05-24 21:19:57 +02:00
parent 9001ccd11f
commit 65c209000c
4 changed files with 17 additions and 10 deletions

View File

@ -843,8 +843,6 @@ class Client(Methods, BaseClient):
if self.proxy:
self.proxy["enabled"] = True
self.proxy["username"] = self.proxy.get("username", None)
self.proxy["password"] = self.proxy.get("password", None)
else:
self.proxy = {}

View File

@ -41,15 +41,15 @@ class TCP(socks.socksocket):
if proxy and self.proxy_enabled:
self.set_proxy(
proxy_type=socks.SOCKS5,
addr=proxy["hostname"],
port=proxy["port"],
username=proxy["username"],
password=proxy["password"]
addr=proxy.get("hostname", None),
port=proxy.get("port", None),
username=proxy.get("username", None),
password=proxy.get("password", None)
)
log.info("Using proxy {}:{}".format(
proxy["hostname"],
proxy["port"]
proxy.get("hostname", None),
proxy.get("port", None)
))
def close(self):

View File

@ -49,8 +49,9 @@ class Auth:
def __init__(self, dc_id: int, test_mode: bool, proxy: dict):
self.dc_id = dc_id
self.test_mode = test_mode
self.proxy = proxy
self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
self.connection = None
@staticmethod
def pack(data: Object) -> bytes:
@ -83,6 +84,8 @@ class Auth:
# The server may close the connection at any time, causing the auth key creation to fail.
# If that happens, just try again up to MAX_RETRIES times.
while True:
self.connection = Connection(DataCenter(self.dc_id, self.test_mode), self.proxy)
try:
log.info("Start creating a new auth key on DC{}".format(self.dc_id))

View File

@ -96,11 +96,15 @@ class Session:
print("Licensed under the terms of the " + __license__, end="\n\n")
Session.notice_displayed = True
self.connection = Connection(DataCenter(dc_id, test_mode), proxy)
self.dc_id = dc_id
self.test_mode = test_mode
self.proxy = proxy
self.api_id = api_id
self.is_cdn = is_cdn
self.client = client
self.connection = None
self.auth_key = auth_key
self.auth_key_id = sha1(auth_key).digest()[-8:]
@ -126,6 +130,8 @@ class Session:
def start(self):
while True:
self.connection = Connection(DataCenter(self.dc_id, self.test_mode), self.proxy)
try:
self.connection.connect()