Connection refactoring

delivrance · delivrance · commit 65c209000cb2 · 2018-05-24T21:19:57.000+02:00
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
@@ -843,8 +843,6 @@ def load_config(self):
 
         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 = {}
 
diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
@@ -41,15 +41,15 @@ def __init__(self, proxy: dict):
         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):
diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py
@@ -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 @@ def create(self):
         # 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))
 
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
@@ -96,11 +96,15 @@ def __init__(self,
             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 @@ def __init__(self,
 
     def start(self):
         while True:
+            self.connection = Connection(DataCenter(self.dc_id, self.test_mode), self.proxy)
+
             try:
                 self.connection.connect()
 
