Use server address from session instead of hardcoded ips

Author: KurimuzonAkuma

pyrogram/client.py

@@ -381,10 +381,8 @@ def __init__(
         self.business_connections = {}
 
         self.sessions = {}
-        self.sessions_lock = asyncio.Lock()
-
         self.media_sessions = {}
-        self.media_sessions_lock = asyncio.Lock()
+        self.sessions_lock = asyncio.Lock()
 
         self.save_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions)
         self.get_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions)
@@ -416,6 +414,8 @@ def __init__(
         else:
             self.loop = asyncio.get_event_loop()
 
+        self.__config = None
+
     def __enter__(self):
         return self.start()
 
@@ -898,12 +898,17 @@ async def load_session(self):
             await self.storage.api_id(self.api_id)
 
             await self.storage.dc_id(2)
+            await self.storage.server_address("149.154.167.51")
+            await self.storage.port(443)
             await self.storage.date(0)
 
             await self.storage.test_mode(self.test_mode)
             await self.storage.auth_key(
                 await Auth(
-                    self, await self.storage.dc_id(),
+                    self,
+                    await self.storage.dc_id(),
+                    await self.storage.server_address(),
+                    await self.storage.port(),
                     await self.storage.test_mode()
                 ).create()
             )
@@ -1136,10 +1141,20 @@ async def get_file(
             try:
                 session = self.media_sessions.get(dc_id)
                 if not session:
+                    dc_option = await self.get_dc_option(dc_id, is_media=True, ipv6=self.ipv6)
+
                     session = self.media_sessions[dc_id] = Session(
-                        self, dc_id,
-                        await Auth(self, dc_id, await self.storage.test_mode()).create()
-                        if dc_id != await self.storage.dc_id()
+                        self,
+                        dc_id,
+                        dc_option.ip_address,
+                        dc_option.port,
+                        await Auth(
+                            self,
+                            dc_id,
+                            dc_option.ip_address,
+                            dc_option.port,
+                            await self.storage.test_mode()
+                        ).create() if dc_id != await self.storage.dc_id()
                         else await self.storage.auth_key(),
                         await self.storage.test_mode(),
                         is_media=True
@@ -1214,9 +1229,23 @@ async def get_file(
                         )
 
                 elif isinstance(r, raw.types.upload.FileCdnRedirect):
+                    dc_option = await self.get_dc_option(dc_id, is_cdn=True, ipv6=self.ipv6)
+
                     cdn_session = Session(
-                        self, r.dc_id, await Auth(self, r.dc_id, await self.storage.test_mode()).create(),
-                        await self.storage.test_mode(), is_media=True, is_cdn=True
+                        self,
+                        r.dc_id,
+                        dc_option.ip_address,
+                        dc_option.port,
+                        await Auth(
+                            self,
+                            r.dc_id,
+                            dc_option.ip_address,
+                            dc_option.port,
+                            await self.storage.test_mode()
+                        ).create(),
+                        await self.storage.test_mode(),
+                        is_media=True,
+                        is_cdn=True
                     )
 
                     try:
@@ -1302,6 +1331,186 @@ async def get_file(
             except Exception as e:
                 log.exception(e)
 
+    async def get_session(
+        self,
+        dc_id: int,
+        is_media: Optional[bool] = False,
+        export_authorization: Optional[bool] = True,
+        server_address: Optional[str] = None,
+        port: Optional[int] = None
+    ) -> "Session":
+        """Get existing session or create a new one.
+
+        Parameters:
+            dc_id (``int``):
+                Datacenter identifier.
+
+            is_media (``bool``, *optional*):
+                Pass True to get or create a media session.
+
+            export_authorization (``bool``, *optional*):
+                Pass True to export authorization after creating the session.
+                Used only when creating a new session.
+
+            server_address (``str``, *optional*):
+                Custom server address to connect to.
+                Used only when creating a new session.
+
+            port (``int``, *optional*):
+                Custom port to connect to.
+                Used only when creating a new session.
+        """
+        if dc_id == await self.storage.dc_id():
+            return self.session
+
+        sessions = self.media_sessions if is_media else self.sessions
+
+        async with self.sessions_lock:
+            if sessions.get(dc_id):
+                return sessions[dc_id]
+
+            dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6)
+
+            session = self.media_sessions[dc_id] = Session(
+                self,
+                dc_id,
+                server_address or dc_option.ip_address,
+                port or dc_option.port,
+                await Auth(
+                    self,
+                    dc_id,
+                    server_address or dc_option.ip_address,
+                    port or dc_option.port,
+                    await self.storage.test_mode()
+                ).create(),
+                await self.storage.test_mode(), is_media=is_media
+            )
+
+            await session.start()
+
+            if export_authorization:
+                for _ in range(3):
+                    exported_auth = await self.invoke(
+                        raw.functions.auth.ExportAuthorization(
+                            dc_id=dc_id
+                        )
+                    )
+
+                    try:
+                        await session.invoke(
+                            raw.functions.auth.ImportAuthorization(
+                                id=exported_auth.id,
+                                bytes=exported_auth.bytes
+                            )
+                        )
+                    except AuthBytesInvalid:
+                        continue
+                    else:
+                        break
+                else:
+                    await session.stop()
+                    raise AuthBytesInvalid
+
+            return session
+
+    async def get_dc_option(
+        self,
+        dc_id: int,
+        is_media: bool = False,
+        is_cdn: bool = False,
+        ipv6: bool = False
+    ) -> "raw.types.DcOption":
+        if not self.__config:
+            self.__config = await self.invoke(raw.functions.help.GetConfig())
+
+        options = [dc for dc in self.__config.dc_options if dc.id == dc_id and dc.ipv6 == ipv6] # type: List[raw.types.DcOption]
+
+        if not options:
+            raise ValueError(f"DC{dc_id} not found")
+
+        if is_cdn:
+            cdn_options = [dc for dc in options if dc.cdn]
+
+            if cdn_options:
+                return cdn_options[0]
+
+            log.debug(
+                "No CDN datacenter found for DC%s, falling back to prod DC",
+                dc_id
+            )
+
+            is_media = True
+
+        if is_media:
+            media_options = [dc for dc in options if dc.media_only]
+
+            if media_options:
+                return media_options[0]
+
+            log.debug(
+                "No media datacenter found for DC%s, falling back to prod DC",
+                dc_id
+            )
+
+        prod_options = [dc for dc in options if not dc.media_only]
+
+        if prod_options:
+            return prod_options[0]
+
+        raise ValueError("No suitable DC found")
+
+    async def set_dc(
+        self,
+        dc_id: Optional[int] = None,
+        *,
+        server_address: Optional[str] = None,
+        port: Optional[int] = None
+    ):
+        """Be careful with this method, you can easily break your session."""
+        if not self.__config:
+            self.__config = await self.invoke(raw.functions.help.GetConfig())
+
+        dc_id = dc_id or self.__config.this_dc
+
+        dc_option = await self.get_dc_option(dc_id, ipv6=self.ipv6)
+
+        server_address = server_address or dc_option.ip_address
+        port = port or dc_option.port
+
+        if dc_id == self.__config.this_dc and (self.session.server_address != server_address or self.session.port != port):
+            await self.storage.server_address(server_address)
+            await self.storage.port(port)
+
+            self.session.server_address = await self.storage.server_address()
+            self.session.port = await self.storage.port()
+
+            await self.session.restart()
+            log.info("Changed the current session DC%s address to %s:%s", dc_id, server_address, port)
+        else:
+            prod_session = self.sessions.get(dc_id)
+
+            if prod_session and (prod_session.server_address != server_address or prod_session.port != port):
+                prod_session.server_address = server_address
+                prod_session.port = port
+
+                await prod_session.restart()
+                log.info("Changed session DC%s address to %s:%s", dc_id, server_address, port)
+            else:
+                await self.get_session(dc_id, server_address=server_address, port=port)
+                log.info("Created new session DC%s with address %s:%s", dc_id, server_address, port)
+
+            media_session = self.media_sessions.get(dc_id)
+
+            if media_session and (media_session.server_address != server_address or media_session.port != port):
+                media_session.server_address = server_address
+                media_session.port = port
+
+                await media_session.restart()
+                log.info("Changed session DC%s (media) address to %s:%s", dc_id, server_address, port)
+            else:
+                await self.get_session(dc_id, is_media=True, server_address=server_address, port=port)
+                log.info("Created new session DC%s (media) with address %s:%s", dc_id, server_address, port)
+
     def guess_mime_type(self, filename: Union[str, BytesIO]) -> Optional[str]:
         if isinstance(filename, BytesIO):
             return self.mimetypes.guess_type(filename.name)[0]

pyrogram/connection/connection.py

@@ -21,7 +21,6 @@
 from typing import Optional, Type
 
 from .transport import TCP, TCPAbridged
-from ..session.internals import DataCenter
 
 log = logging.getLogger(__name__)
 
@@ -32,6 +31,8 @@ class Connection:
     def __init__(
         self,
         dc_id: int,
+        server_address: str,
+        port: int,
         test_mode: bool,
         ipv6: bool,
         proxy: dict,
@@ -41,14 +42,15 @@ def __init__(
         loop: Optional[asyncio.AbstractEventLoop] = None
     ) -> None:
         self.dc_id = dc_id
+        self.server_address = server_address
+        self.port = port
         self.test_mode = test_mode
         self.ipv6 = ipv6
         self.proxy = proxy
         self.media = media
         self.protocol_factory = protocol_factory
         self.crypto_executor_workers = crypto_executor_workers
 
-        self.address = DataCenter(dc_id, test_mode, ipv6, media)
         self.protocol: Optional[TCP] = None
 
         if isinstance(loop, asyncio.AbstractEventLoop):
@@ -62,7 +64,7 @@ async def connect(self) -> None:
 
             try:
                 log.info("Connecting...")
-                await self.protocol.connect(self.address)
+                await self.protocol.connect((self.server_address, self.port))
             except OSError as e:
                 log.warning("Unable to connect due to network issues: %s", e)
                 await self.protocol.close()

pyrogram/methods/advanced/save_file.py

@@ -145,9 +145,16 @@ async def worker(session):
             dc_id = await self.storage.dc_id()
 
             session = self.media_sessions.get(dc_id)
+
             if not session:
+                dc_option = await self.get_dc_option(dc_id, is_media=True, ipv6=self.ipv6)
+
                 session = self.media_sessions[dc_id] = Session(
-                    self, dc_id, await self.storage.auth_key(),
+                    self,
+                    dc_id,
+                    dc_option.ip_address,
+                    dc_option.port,
+                    await self.storage.auth_key(),
                     await self.storage.test_mode(), is_media=True
                 )
                 await session.start()

pyrogram/methods/auth/connect.py

@@ -40,8 +40,12 @@ async def connect(
         await self.load_session()
 
         self.session = Session(
-            self, await self.storage.dc_id(),
-            await self.storage.auth_key(), await self.storage.test_mode()
+            self,
+            await self.storage.dc_id(),
+            await self.storage.server_address(),
+            await self.storage.port(),
+            await self.storage.auth_key(),
+            await self.storage.test_mode()
         )
 
         await self.session.start()

pyrogram/methods/auth/send_code.py

@@ -23,6 +23,7 @@
 from pyrogram import raw
 from pyrogram import types
 from pyrogram.errors import PhoneMigrate, NetworkMigrate
+from pyrogram.raw.base import dc_option
 from pyrogram.session import Session, Auth
 
 log = logging.getLogger(__name__)
@@ -102,18 +103,30 @@ async def send_code(
                     )
                 )
             except (PhoneMigrate, NetworkMigrate) as e:
+                dc_option = await self.get_dc_option(e.value, ipv6=self.ipv6)
                 await self.session.stop()
 
                 await self.storage.dc_id(e.value)
+                await self.storage.server_address(dc_option.ip_address)
+                await self.storage.port(dc_option.port)
+
                 await self.storage.auth_key(
                     await Auth(
-                        self, await self.storage.dc_id(),
+                        self,
+                        await self.storage.dc_id(),
+                        await self.storage.server_address(),
+                        await self.storage.port(),
                         await self.storage.test_mode()
                     ).create()
                 )
+
                 self.session = Session(
-                    self, await self.storage.dc_id(),
-                    await self.storage.auth_key(), await self.storage.test_mode()
+                    self,
+                    await self.storage.dc_id(),
+                    await self.storage.server_address(),
+                    await self.storage.port(),
+                    await self.storage.auth_key(),
+                    await self.storage.test_mode()
                 )
 
                 await self.session.start()