test1

Author: Surendra9123

pyrogram/connection/connection.py

@@ -61,7 +61,7 @@ def __init__(
 
     async def connect(self) -> None:
         for i in range(Connection.MAX_CONNECTION_ATTEMPTS):
-            self.protocol = self.protocol_factory(ipv6=self.ipv6, proxy=self.proxy, crypto_executor_workers=self.crypto_executor_workers, loop=self.loop)
+            self.protocol = self.protocol_factory(dc_id=self.dc_id, ipv6=self.ipv6, proxy=self.proxy, crypto_executor_workers=self.crypto_executor_workers, loop=self.loop)
 
             try:
                 log.info("Connecting...")

pyrogram/connection/transport/tcp/tcp.py

@@ -36,6 +36,7 @@ class ProxyDict(TypedDict):
     scheme: str
     hostname: str
     port: int
+    secret: Optional[str]
     username: Optional[str]
     password: Optional[str]
 
@@ -154,7 +155,7 @@ async def _connect_via_direct(self, destination: Tuple[str, int]) -> None:
         log.info("Connection established")
 
     async def _connect(self, destination: Tuple[str, int]) -> None:
-        if self.proxy:
+        if self.proxy and isinstance(self.proxy, dict) and self.proxy.get("scheme", "").lower() != "mtproxy":
             await self._connect_via_proxy(destination)
         else:
             await self._connect_via_direct(destination)

pyrogram/connection/transport/tcp/tcp_abridged.py

@@ -16,10 +16,13 @@
 #  You should have received a copy of the GNU Lesser General Public License
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
+import hashlib
 import asyncio
 import logging
 from typing import Optional, Tuple, Union
 
+from pyrogram.crypto import aes
 from .tcp import TCP, ProxyDict
 
 log = logging.getLogger(__name__)
@@ -28,30 +31,170 @@
 class TCPAbridged(TCP):
     def __init__(
         self,
+        dc_id: int,
         ipv6: bool = False,
         proxy: Optional[Union[str, ProxyDict]] = None,
         crypto_executor_workers: int = 1,
         loop: Optional[asyncio.AbstractEventLoop] = None,
     ) -> None:
         super().__init__(ipv6, proxy, crypto_executor_workers, loop)
+        self.dc_id = dc_id
+        self.encrypt = None
+        self.decrypt = None
+        self.secret = None
+        self.hostname = None
+        self.port = None
+        if isinstance(proxy, dict) and proxy.get("scheme", "").lower() == "mtproxy":
+            self.hostname = proxy.get("hostname")
+            self.port = proxy.get("port")
+            self.secret = proxy.get("secret")
 
+    async def connect_via_mtproxy(self, address):
+        self.marker_event.clear()
+
+        nonce = bytearray(os.urandom(64))
+
+        while (
+            nonce[0] == 0xef
+            or nonce[:4] in (b"HEAD", b"POST", b"GET ", b"OPTI", b"\xee"*4)
+            or nonce[4:8] == b"\x00"*4
+        ):
+            nonce = bytearray(os.urandom(64))
+
+        temp = bytearray(nonce[55:7:-1])
+
+        secret = self.secret
+
+        if isinstance(secret, str):
+            if secret.startswith("dd"):
+                secret = secret[2:]
+            secret = bytes.fromhex(secret)
+
+        encrypt_key = hashlib.sha256(
+            bytes(nonce[8:40]) + secret
+        ).digest()
+
+        decrypt_key = hashlib.sha256(
+            bytes(temp[:32]) + secret
+        ).digest()
+
+        self.encrypt = (
+            encrypt_key,
+            bytes(nonce[40:56]),
+            bytearray(1)
+        )
+
+        self.decrypt = (
+            decrypt_key,
+            bytes(temp[32:48]),
+            bytearray(1)
+        )
+
+        nonce[56:60] =  b'\xef\xef\xef\xef'
+
+        nonce[60:62] = self.dc_id.to_bytes(
+            2,
+            "little",
+            signed=True
+        )
+
+        encrypted = aes.ctr256_encrypt(
+            nonce,
+            *self.encrypt
+        )
+
+        nonce[56:64] = encrypted[56:64]
+
+        await super().connect((self.hostname, self.port))
+
+        await super().send(
+            nonce,
+            wait_for_marker=False
+        )
+
+        self.marker_event.set()
+    
     async def connect(self, address: Tuple[str, int]) -> None:
+        if self.hostname and self.port and self.secret:
+            return await self.connect_via_mtproxy(address)
+
         self.marker_event.clear()
         await super().connect(address)
         await super().send(b"\xef", wait_for_marker=False)
         self.marker_event.set()
 
-    async def send(self, data: bytes, *args) -> None:
+    async def send_via_mtproxy(self, data: bytes, *args):
         length = len(data) // 4
 
+        payload = (
+            bytes([length])
+            if length <= 126
+            else b"\x7f" + length.to_bytes(3, "little")
+        ) + data
+
+        payload = await self.loop.run_in_executor(
+            self.crypto_executor,
+            aes.ctr256_encrypt,
+            payload,
+            *self.encrypt
+        )
+        await super().send(payload)
+
+    async def send(self, data: bytes, *args) -> None:
+        if self.hostname and self.port and self.secret:
+            return await self.send_via_mtproxy(data, *args)
+
+        length = len(data) // 4
         await super().send(
             (bytes([length])
-             if length <= 126
-             else b"\x7f" + length.to_bytes(3, "little"))
+            if length <= 126
+            else b"\x7f" + length.to_bytes(3, "little"))
             + data
         )
 
+    async def recv_via_mtproxy(self, length: int = 0):
+        length = await super().recv(1)
+
+        if length is None:
+            return None
+
+        length = aes.ctr256_decrypt(
+            length,
+            *self.decrypt
+        )
+
+        if length == b"\x7f":
+            length = await super().recv(3)
+
+            if length is None:
+                return None
+
+            length = aes.ctr256_decrypt(
+                length,
+                *self.decrypt
+            )
+
+        packet_length = int.from_bytes(
+            length,
+            "little"
+        ) * 4
+
+        data = await super().recv(packet_length)
+
+        if data is None:
+            return None
+
+        return await self.loop.run_in_executor(
+            self.crypto_executor,
+            aes.ctr256_decrypt,
+            data,
+            *self.decrypt
+        )
+
     async def recv(self, length: int = 0) -> Optional[bytes]:
+        if self.hostname and self.port and self.secret:
+            return await self.recv_via_mtproxy(length)
+
         length = await super().recv(1)
 
         if length is None:
@@ -63,4 +206,4 @@ async def recv(self, length: int = 0) -> Optional[bytes]:
             if length is None:
                 return None
 
-        return await super().recv(int.from_bytes(length, "little") * 4)
+        return await super().recv(int.from_bytes(length, "little") * 4)

pyrogram/session/session.py

@@ -181,6 +181,14 @@ async def start(self):
 
             init_connection_params = self.client.init_connection_params
 
+            proxy = None
+            
+            if isinstance(proxy, dict) and proxy.get("scheme", "").lower() == "mtproxy":
+                proxy = raw.types.InputClientProxy(
+                    address=proxy.get("hostname", ""),
+                    port=proxy.get("port", 0)
+                )
+
             if isinstance(init_connection_params, dict):
                 init_connection_params = utils.obj_to_jsonvalue(init_connection_params)
 
@@ -198,6 +206,7 @@ async def start(self):
                             lang_code=self.client.lang_code,
                             query=raw.functions.help.GetConfig(),
                             params=init_connection_params,
+                            proxy=proxy
                         )
                     ),
                     timeout=self.START_TIMEOUT