diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index dba204fd..7d195e01 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -146,6 +146,7 @@ class Client(Methods, BaseClient): device_model: str = None, system_version: str = None, lang_code: str = None, + ipv6: bool = False, proxy: dict = None, test_mode: bool = False, phone_number: str = None, @@ -166,6 +167,7 @@ class Client(Methods, BaseClient): self.device_model = device_model self.system_version = system_version self.lang_code = lang_code + self.ipv6 = ipv6 # TODO: Make code consistent, use underscore for private/protected fields self._proxy = proxy self.test_mode = test_mode @@ -388,7 +390,7 @@ class Client(Methods, BaseClient): self.session.stop() self.dc_id = e.x - self.auth_key = Auth(self.dc_id, self.test_mode, self._proxy).create() + self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() self.session = Session( self, @@ -433,7 +435,7 @@ class Client(Methods, BaseClient): self.session.stop() self.dc_id = e.x - self.auth_key = Auth(self.dc_id, self.test_mode, self._proxy).create() + self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() self.session = Session( self, @@ -936,7 +938,7 @@ class Client(Methods, BaseClient): except FileNotFoundError: self.dc_id = 1 self.date = 0 - self.auth_key = Auth(self.dc_id, self.test_mode, self._proxy).create() + self.auth_key = Auth(self.dc_id, self.test_mode, self.ipv6, self._proxy).create() else: self.dc_id = s["dc_id"] self.test_mode = s["test_mode"] @@ -1072,7 +1074,7 @@ class Client(Methods, BaseClient): 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, self.dc_id, self.auth_key, is_media=True) + session = Session(self, self.ipv6, self.dc_id, self.auth_key, is_media=True) session.start() try: diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py index a53295ce..ca9bd96c 100644 --- a/pyrogram/connection/connection.py +++ b/pyrogram/connection/connection.py @@ -21,6 +21,7 @@ import threading import time from .transport import * +from ..session.internals import DataCenter log = logging.getLogger(__name__) @@ -36,21 +37,24 @@ class Connection: 4: TCPIntermediateO } - def __init__(self, address: tuple, proxy: dict, mode: int = 1): - self.address = address + def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict, mode: int = 1): + self.ipv6 = ipv6 self.proxy = proxy + self.address = DataCenter(dc_id, test_mode, ipv6) self.mode = self.MODES.get(mode, TCPAbridged) + self.lock = threading.Lock() self.connection = None def connect(self): for i in range(Connection.MAX_RETRIES): - self.connection = self.mode(self.proxy) + self.connection = self.mode(self.ipv6, self.proxy) try: log.info("Connecting...") self.connection.connect(self.address) - except OSError: + except OSError as e: + log.warning(e) # TODO: Remove self.connection.close() time.sleep(1) else: diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index 5df8aacb..8560604f 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -33,8 +33,9 @@ log = logging.getLogger(__name__) class TCP(socks.socksocket): - def __init__(self, proxy: dict): - super().__init__() + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(family=socket.AF_INET6 if ipv6 else socket.AF_INET) + self.settimeout(10) self.proxy_enabled = proxy.get("enabled", False) diff --git a/pyrogram/connection/transport/tcp/tcp_abridged.py b/pyrogram/connection/transport/tcp/tcp_abridged.py index 472f4799..d89421f5 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged.py @@ -24,8 +24,8 @@ log = logging.getLogger(__name__) class TCPAbridged(TCP): - def __init__(self, proxy: dict): - super().__init__(proxy) + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(ipv6, proxy) def connect(self, address: tuple): super().connect(address) diff --git a/pyrogram/connection/transport/tcp/tcp_abridged_o.py b/pyrogram/connection/transport/tcp/tcp_abridged_o.py index bba88e34..57cc0336 100644 --- a/pyrogram/connection/transport/tcp/tcp_abridged_o.py +++ b/pyrogram/connection/transport/tcp/tcp_abridged_o.py @@ -28,8 +28,9 @@ log = logging.getLogger(__name__) class TCPAbridgedO(TCP): RESERVED = (b"HEAD", b"POST", b"GET ", b"OPTI", b"\xee" * 4) - def __init__(self, proxy: dict): - super().__init__(proxy) + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(ipv6, proxy) + self.encrypt = None self.decrypt = None diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py index 1b131678..0aac1a14 100644 --- a/pyrogram/connection/transport/tcp/tcp_full.py +++ b/pyrogram/connection/transport/tcp/tcp_full.py @@ -26,8 +26,9 @@ log = logging.getLogger(__name__) class TCPFull(TCP): - def __init__(self, proxy: dict): - super().__init__(proxy) + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(ipv6, proxy) + self.seq_no = None def connect(self, address: tuple): diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate.py b/pyrogram/connection/transport/tcp/tcp_intermediate.py index 4b2e2596..b33fe466 100644 --- a/pyrogram/connection/transport/tcp/tcp_intermediate.py +++ b/pyrogram/connection/transport/tcp/tcp_intermediate.py @@ -25,8 +25,8 @@ log = logging.getLogger(__name__) class TCPIntermediate(TCP): - def __init__(self, proxy: dict): - super().__init__(proxy) + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(ipv6, proxy) def connect(self, address: tuple): super().connect(address) diff --git a/pyrogram/connection/transport/tcp/tcp_intermediate_o.py b/pyrogram/connection/transport/tcp/tcp_intermediate_o.py index df79352a..e8c96ebc 100644 --- a/pyrogram/connection/transport/tcp/tcp_intermediate_o.py +++ b/pyrogram/connection/transport/tcp/tcp_intermediate_o.py @@ -29,8 +29,9 @@ log = logging.getLogger(__name__) class TCPIntermediateO(TCP): RESERVED = (b"HEAD", b"POST", b"GET ", b"OPTI", b"\xee" * 4) - def __init__(self, proxy: dict): - super().__init__(proxy) + def __init__(self, ipv6: bool, proxy: dict): + super().__init__(ipv6, proxy) + self.encrypt = None self.decrypt = None diff --git a/pyrogram/session/auth.py b/pyrogram/session/auth.py index 80956187..2142be59 100644 --- a/pyrogram/session/auth.py +++ b/pyrogram/session/auth.py @@ -26,7 +26,7 @@ from pyrogram.api import functions, types from pyrogram.api.core import Object, Long, Int from pyrogram.connection import Connection from pyrogram.crypto import AES, RSA, Prime -from .internals import MsgId, DataCenter +from .internals import MsgId log = logging.getLogger(__name__) @@ -46,9 +46,10 @@ class Auth: 16 ) - def __init__(self, dc_id: int, test_mode: bool, proxy: dict): + def __init__(self, dc_id: int, test_mode: bool, ipv6: bool, proxy: dict): self.dc_id = dc_id self.test_mode = test_mode + self.ipv6 = ipv6 self.proxy = proxy self.connection = None @@ -84,7 +85,7 @@ 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) + self.connection = Connection(self.dc_id, self.test_mode, self.ipv6, self.proxy) try: log.info("Start creating a new auth key on DC{}".format(self.dc_id)) diff --git a/pyrogram/session/internals/data_center.py b/pyrogram/session/internals/data_center.py index 232ca13b..154cd9e0 100644 --- a/pyrogram/session/internals/data_center.py +++ b/pyrogram/session/internals/data_center.py @@ -34,5 +34,24 @@ class DataCenter: 121: "95.213.217.195" } - def __new__(cls, dc_id: int, test_mode: bool): - return (cls.TEST[dc_id], 80) if test_mode else (cls.PROD[dc_id], 443) + TEST_IPV6 = { + 1: "2001:0b28:f23d:f001:0000:0000:0000:000e", + 2: "2001:067c:04e8:f002:0000:0000:0000:000e", + 3: "2001:0b28:f23d:f003:0000:0000:0000:000e", + 121: "2a03:b0c0:0003:00d0:0000:0000:0114:d001" + } + + PROD_IPV6 = { + 1: "2001:0b28:f23d:f001:0000:0000:0000:000a", + 2: "2001:067c:04e8:f002:0000:0000:0000:000a", + 3: "2001:0b28:f23d:f003:0000:0000:0000:000a", + 4: "2001:067c:04e8:f004:0000:0000:0000:000a", + 5: "2001:0b28:f23f:f005:0000:0000:0000:000a", + 121: "2a03:b0c0:0003:00d0:0000:0000:0114:d001" + } + + def __new__(cls, dc_id: int, test_mode: bool, ipv6: bool): + if ipv6: + return (cls.TEST_IPV6[dc_id], 80) if test_mode else (cls.PROD_IPV6[dc_id], 443) + else: + return (cls.TEST[dc_id], 80) if test_mode else (cls.PROD[dc_id], 443) diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index b7645b11..c87f0698 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -34,7 +34,7 @@ from pyrogram.api.core import Message, Object, MsgContainer, Long, FutureSalt, I from pyrogram.api.errors import Error, InternalServerError, AuthKeyDuplicated from pyrogram.connection import Connection from pyrogram.crypto import AES, KDF -from .internals import MsgId, MsgFactory, DataCenter +from .internals import MsgId, MsgFactory log = logging.getLogger(__name__) @@ -112,7 +112,7 @@ class Session: def start(self): while True: - self.connection = Connection(DataCenter(self.dc_id, self.client.test_mode), self.client.proxy) + self.connection = Connection(self.dc_id, self.client.test_mode, self.client.ipv6, self.client.proxy) try: self.connection.connect()