1616# You should have received a copy of the GNU Lesser General Public License
1717# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818
19+ import os
20+ import hashlib
1921import asyncio
2022import logging
2123from typing import Optional , Tuple , Union
2224
25+ from pyrogram .crypto import aes
2326from .tcp import TCP , ProxyDict
2427
2528log = logging .getLogger (__name__ )
2831class TCPAbridged (TCP ):
2932 def __init__ (
3033 self ,
34+ dc_id : int ,
3135 ipv6 : bool = False ,
3236 proxy : Optional [Union [str , ProxyDict ]] = None ,
3337 crypto_executor_workers : int = 1 ,
3438 loop : Optional [asyncio .AbstractEventLoop ] = None ,
3539 ) -> None :
3640 super ().__init__ (ipv6 , proxy , crypto_executor_workers , loop )
41+ self .dc_id = dc_id
42+ self .encrypt = None
43+ self .decrypt = None
44+ self .secret = None
45+ self .hostname = None
46+ self .port = None
3747
48+ if isinstance (proxy , dict ) and proxy .get ("scheme" , "" ).lower () == "mtproxy" :
49+ self .hostname = proxy .get ("hostname" )
50+ self .port = proxy .get ("port" )
51+ self .secret = proxy .get ("secret" )
52+
53+ self .is_mtproxy = bool (self .hostname and self .port and self .secret )
54+
55+ log .info (f"TCPAbridged initialized with server address: { self .hostname } :{ self .port } (MTProxy={ self .is_mtproxy } )." )
56+
57+ async def connect_via_mtproxy (self ):
58+ self .marker_event .clear ()
59+
60+ nonce = bytearray (os .urandom (64 ))
61+
62+ while (
63+ nonce [0 ] == 0xef
64+ or nonce [:4 ] in (b"HEAD" , b"POST" , b"GET " , b"OPTI" , b"\xee " * 4 )
65+ or nonce [4 :8 ] == b"\x00 " * 4
66+ ):
67+ nonce = bytearray (os .urandom (64 ))
68+
69+ temp = bytearray (nonce [55 :7 :- 1 ])
70+
71+ secret = self .secret
72+
73+ if isinstance (secret , str ):
74+ if secret .startswith ("dd" ):
75+ secret = secret [2 :]
76+ secret = bytes .fromhex (secret )
77+
78+ encrypt_key = hashlib .sha256 (
79+ bytes (nonce [8 :40 ]) + secret
80+ ).digest ()
81+
82+ decrypt_key = hashlib .sha256 (
83+ bytes (temp [:32 ]) + secret
84+ ).digest ()
85+
86+ self .encrypt = (
87+ encrypt_key ,
88+ bytes (nonce [40 :56 ]),
89+ bytearray (1 )
90+ )
91+
92+ self .decrypt = (
93+ decrypt_key ,
94+ bytes (temp [32 :48 ]),
95+ bytearray (1 )
96+ )
97+
98+ nonce [56 :60 ] = b'\xef \xef \xef \xef '
99+
100+ nonce [60 :62 ] = self .dc_id .to_bytes (
101+ 2 ,
102+ "little" ,
103+ signed = True
104+ )
105+
106+ encrypted = aes .ctr256_encrypt (
107+ nonce ,
108+ * self .encrypt
109+ )
110+
111+ nonce [56 :64 ] = encrypted [56 :64 ]
112+
113+ await super ().connect ((self .hostname , self .port ))
114+
115+ await super ().send (
116+ nonce ,
117+ wait_for_marker = False
118+ )
119+
120+ self .marker_event .set ()
121+
38122 async def connect (self , address : Tuple [str , int ]) -> None :
123+ if self .is_mtproxy :
124+ return await self .connect_via_mtproxy ()
125+
39126 self .marker_event .clear ()
40127 await super ().connect (address )
41128 await super ().send (b"\xef " , wait_for_marker = False )
42129 self .marker_event .set ()
43130
44- async def send (self , data : bytes , * args ) -> None :
131+ async def send_via_mtproxy (self , data : bytes , * args ):
45132 length = len (data ) // 4
46133
134+ payload = (
135+ bytes ([length ])
136+ if length <= 126
137+ else b"\x7f " + length .to_bytes (3 , "little" )
138+ ) + data
139+
140+ payload = await self .loop .run_in_executor (
141+ self .crypto_executor ,
142+ aes .ctr256_encrypt ,
143+ payload ,
144+ * self .encrypt
145+ )
146+ await super ().send (payload )
147+
148+ async def send (self , data : bytes , * args ) -> None :
149+ if self .is_mtproxy :
150+ return await self .send_via_mtproxy (data , * args )
151+
152+ length = len (data ) // 4
47153 await super ().send (
48154 (bytes ([length ])
49- if length <= 126
50- else b"\x7f " + length .to_bytes (3 , "little" ))
155+ if length <= 126
156+ else b"\x7f " + length .to_bytes (3 , "little" ))
51157 + data
52158 )
53159
160+ async def recv_via_mtproxy (self , length : int = 0 ):
161+ length = await super ().recv (1 )
162+
163+ if length is None :
164+ return None
165+
166+ length = aes .ctr256_decrypt (
167+ length ,
168+ * self .decrypt
169+ )
170+
171+ if length == b"\x7f " :
172+ length = await super ().recv (3 )
173+
174+ if length is None :
175+ return None
176+
177+ length = aes .ctr256_decrypt (
178+ length ,
179+ * self .decrypt
180+ )
181+
182+ packet_length = int .from_bytes (
183+ length ,
184+ "little"
185+ ) * 4
186+
187+ data = await super ().recv (packet_length )
188+
189+ if data is None :
190+ return None
191+
192+ return await self .loop .run_in_executor (
193+ self .crypto_executor ,
194+ aes .ctr256_decrypt ,
195+ data ,
196+ * self .decrypt
197+ )
198+
54199 async def recv (self , length : int = 0 ) -> Optional [bytes ]:
200+ if self .is_mtproxy :
201+ return await self .recv_via_mtproxy (length )
202+
55203 length = await super ().recv (1 )
56204
57205 if length is None :
@@ -63,4 +211,4 @@ async def recv(self, length: int = 0) -> Optional[bytes]:
63211 if length is None :
64212 return None
65213
66- return await super ().recv (int .from_bytes (length , "little" ) * 4 )
214+ return await super ().recv (int .from_bytes (length , "little" ) * 4 )
0 commit comments