Skip to content

Commit cd6bfbb

Browse files
committed
t2
1 parent 9936b98 commit cd6bfbb

4 files changed

Lines changed: 164 additions & 6 deletions

File tree

pyrogram/connection/connection.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def __init__(
6161

6262
async def connect(self) -> None:
6363
for i in range(Connection.MAX_CONNECTION_ATTEMPTS):
64-
self.protocol = self.protocol_factory(ipv6=self.ipv6, proxy=self.proxy, crypto_executor_workers=self.crypto_executor_workers, loop=self.loop)
64+
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)
6565

6666
try:
6767
log.info("Connecting...")

pyrogram/connection/transport/tcp/tcp.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ class ProxyDict(TypedDict):
3636
scheme: str
3737
hostname: str
3838
port: int
39+
secret: Optional[str]
3940
username: Optional[str]
4041
password: Optional[str]
4142

@@ -154,7 +155,7 @@ async def _connect_via_direct(self, destination: Tuple[str, int]) -> None:
154155
log.info("Connection established")
155156

156157
async def _connect(self, destination: Tuple[str, int]) -> None:
157-
if self.proxy:
158+
if self.proxy and isinstance(self.proxy, dict) and self.proxy.get("scheme", "").lower() != "mtproxy":
158159
await self._connect_via_proxy(destination)
159160
else:
160161
await self._connect_via_direct(destination)

pyrogram/connection/transport/tcp/tcp_abridged.py

Lines changed: 152 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
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
1921
import asyncio
2022
import logging
2123
from typing import Optional, Tuple, Union
2224

25+
from pyrogram.crypto import aes
2326
from .tcp import TCP, ProxyDict
2427

2528
log = logging.getLogger(__name__)
@@ -28,30 +31,175 @@
2831
class 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)

pyrogram/session/session.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,14 @@ async def start(self):
181181

182182
init_connection_params = self.client.init_connection_params
183183

184+
proxy = None
185+
186+
if isinstance(proxy, dict) and proxy.get("scheme", "").lower() == "mtproxy":
187+
proxy = raw.types.InputClientProxy(
188+
address=proxy.get("hostname", ""),
189+
port=proxy.get("port", 0)
190+
)
191+
184192
if isinstance(init_connection_params, dict):
185193
init_connection_params = utils.obj_to_jsonvalue(init_connection_params)
186194

@@ -198,6 +206,7 @@ async def start(self):
198206
lang_code=self.client.lang_code,
199207
query=raw.functions.help.GetConfig(),
200208
params=init_connection_params,
209+
proxy=proxy
201210
)
202211
),
203212
timeout=self.START_TIMEOUT

0 commit comments

Comments
 (0)