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: if self.proxy:
self.proxy["enabled"] = True self.proxy["enabled"] = True
self.proxy["username"] = self.proxy.get("username", None)
self.proxy["password"] = self.proxy.get("password", None)
else: else:
self.proxy = {} self.proxy = {}

View File

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

View File

@ -49,8 +49,9 @@ class Auth:
def __init__(self, dc_id: int, test_mode: bool, proxy: dict): def __init__(self, dc_id: int, test_mode: bool, proxy: dict):
self.dc_id = dc_id self.dc_id = dc_id
self.test_mode = test_mode self.test_mode = test_mode
self.proxy = proxy
self.connection = Connection(DataCenter(dc_id, test_mode), proxy) self.connection = None
@staticmethod @staticmethod
def pack(data: Object) -> bytes: 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. # 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. # If that happens, just try again up to MAX_RETRIES times.
while True: while True:
self.connection = Connection(DataCenter(self.dc_id, self.test_mode), self.proxy)
try: try:
log.info("Start creating a new auth key on DC{}".format(self.dc_id)) 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") print("Licensed under the terms of the " + __license__, end="\n\n")
Session.notice_displayed = True 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.api_id = api_id
self.is_cdn = is_cdn self.is_cdn = is_cdn
self.client = client self.client = client
self.connection = None
self.auth_key = auth_key self.auth_key = auth_key
self.auth_key_id = sha1(auth_key).digest()[-8:] self.auth_key_id = sha1(auth_key).digest()[-8:]
@ -126,6 +130,8 @@ class Session:
def start(self): def start(self):
while True: while True:
self.connection = Connection(DataCenter(self.dc_id, self.test_mode), self.proxy)
try: try:
self.connection.connect() self.connection.connect()