Skip to content

Commit fe7662e

Browse files
Surendra9123Surendra9123
committed
New Fix
1 parent 7cab2ad commit fe7662e

2 files changed

Lines changed: 104 additions & 140 deletions

File tree

‎pyrogram/client.py‎

Lines changed: 103 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -1153,49 +1153,7 @@ async def get_file(
11531153
dc_id = file_id.dc_id
11541154

11551155
try:
1156-
session = self.media_sessions.get(dc_id)
1157-
if not session:
1158-
dc_option = await self.get_dc_option(dc_id, is_media=True, ipv6=self.ipv6)
1159-
1160-
session = self.media_sessions[dc_id] = Session(
1161-
self,
1162-
dc_id,
1163-
dc_option.ip_address,
1164-
dc_option.port,
1165-
await Auth(
1166-
self,
1167-
dc_id,
1168-
dc_option.ip_address,
1169-
dc_option.port,
1170-
await self.storage.test_mode()
1171-
).create() if dc_id != await self.storage.dc_id()
1172-
else await self.storage.auth_key(),
1173-
await self.storage.test_mode(),
1174-
is_media=True
1175-
)
1176-
await session.start()
1177-
1178-
if dc_id != await self.storage.dc_id():
1179-
for _ in range(3):
1180-
exported_auth = await self.invoke(
1181-
raw.functions.auth.ExportAuthorization(
1182-
dc_id=dc_id
1183-
)
1184-
)
1185-
1186-
try:
1187-
await session.invoke(
1188-
raw.functions.auth.ImportAuthorization(
1189-
id=exported_auth.id,
1190-
bytes=exported_auth.bytes
1191-
)
1192-
)
1193-
except AuthBytesInvalid:
1194-
continue
1195-
else:
1196-
break
1197-
else:
1198-
raise AuthBytesInvalid
1156+
session = await self.get_session(dc_id, is_media=True)
11991157

12001158
r = await session.invoke(
12011159
raw.functions.upload.GetFile(
@@ -1243,28 +1201,10 @@ async def get_file(
12431201
)
12441202

12451203
elif isinstance(r, raw.types.upload.FileCdnRedirect):
1246-
dc_option = await self.get_dc_option(dc_id, is_cdn=True, ipv6=self.ipv6)
1247-
1248-
cdn_session = Session(
1249-
self,
1250-
r.dc_id,
1251-
dc_option.ip_address,
1252-
dc_option.port,
1253-
await Auth(
1254-
self,
1255-
r.dc_id,
1256-
dc_option.ip_address,
1257-
dc_option.port,
1258-
await self.storage.test_mode()
1259-
).create(),
1260-
await self.storage.test_mode(),
1261-
is_media=True,
1262-
is_cdn=True
1263-
)
12641204

1265-
try:
1266-
await cdn_session.start()
1205+
cdn_session = await self.get_session(dc_id, is_cdn=True, temporary=True)
12671206

1207+
try:
12681208
while True:
12691209
r2 = await cdn_session.invoke(
12701210
raw.functions.upload.GetCdnFile(
@@ -1290,13 +1230,12 @@ async def get_file(
12901230
chunk = r2.bytes
12911231

12921232
# https://core.telegram.org/cdn#decrypting-files
1293-
decrypted_chunk = aes.ctr256_decrypt(
1233+
decrypted_chunk = await self.loop.run_in_executor(
1234+
self.executor,
1235+
aes.ctr256_decrypt,
12941236
chunk,
12951237
r.encryption_key,
1296-
bytearray(
1297-
r.encryption_iv[:-4]
1298-
+ (offset_bytes // 16).to_bytes(4, "big")
1299-
)
1238+
bytearray(r.encryption_iv[:-4] + (offset_bytes // 16).to_bytes(4, "big"))
13001239
)
13011240

13021241
hashes = await session.invoke(
@@ -1307,12 +1246,14 @@ async def get_file(
13071246
)
13081247

13091248
# https://core.telegram.org/cdn#verifying-files
1310-
for i, h in enumerate(hashes):
1311-
cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)]
1312-
CDNFileHashMismatch.check(
1313-
h.hash == sha256(cdn_chunk).digest(),
1314-
"h.hash == sha256(cdn_chunk).digest()"
1315-
)
1249+
def _check_all():
1250+
for i, h in enumerate(hashes):
1251+
cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)]
1252+
CDNFileHashMismatch.check(
1253+
h.hash == sha256(cdn_chunk).digest(),
1254+
"h.hash == sha256(cdn_chunk).digest()"
1255+
)
1256+
await self.loop.run_in_executor(self.executor, _check_all)
13161257

13171258
yield decrypted_chunk
13181259

@@ -1345,13 +1286,89 @@ async def get_file(
13451286
except Exception as e:
13461287
log.exception(e)
13471288

1289+
async def _get_session(
1290+
self,
1291+
dc_id: int,
1292+
is_media: Optional[bool] = False,
1293+
is_cdn: Optional[bool] = False,
1294+
export_authorization: Optional[bool] = True,
1295+
server_address: Optional[str] = None,
1296+
port: Optional[int] = None,
1297+
temporary: Optional[bool] = False
1298+
) -> "Session":
1299+
if dc_id == await self.storage.dc_id() and (not is_media) and (not temporary):
1300+
return self.session
1301+
1302+
sessions = self.media_sessions if is_media else self.sessions
1303+
1304+
if sessions.get(dc_id) and (not temporary):
1305+
return sessions[dc_id]
1306+
1307+
dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6, is_cdn=is_cdn)
1308+
1309+
if is_media:
1310+
auth_key = (await self._get_session(dc_id)).auth_key
1311+
else:
1312+
if dc_id != await self.storage.dc_id():
1313+
auth_key = await Auth(
1314+
self,
1315+
dc_id,
1316+
server_address or dc_option.ip_address,
1317+
port or dc_option.port,
1318+
await self.storage.test_mode()
1319+
).create()
1320+
else:
1321+
auth_key = await self.storage.auth_key()
1322+
1323+
session = Session(
1324+
self,
1325+
dc_id,
1326+
server_address or dc_option.ip_address,
1327+
port or dc_option.port,
1328+
auth_key,
1329+
await self.storage.test_mode(),
1330+
is_media=is_media
1331+
)
1332+
1333+
if not temporary:
1334+
sessions[dc_id] = session
1335+
1336+
await session.start()
1337+
1338+
if dc_id != await self.storage.dc_id() and export_authorization:
1339+
for _ in range(3):
1340+
exported_auth = await self.invoke(
1341+
raw.functions.auth.ExportAuthorization(
1342+
dc_id=dc_id
1343+
)
1344+
)
1345+
1346+
try:
1347+
await session.invoke(
1348+
raw.functions.auth.ImportAuthorization(
1349+
id=exported_auth.id,
1350+
bytes=exported_auth.bytes
1351+
)
1352+
)
1353+
except AuthBytesInvalid:
1354+
continue
1355+
else:
1356+
break
1357+
else:
1358+
await session.stop()
1359+
raise AuthBytesInvalid
1360+
1361+
return session
1362+
13481363
async def get_session(
13491364
self,
13501365
dc_id: int,
13511366
is_media: Optional[bool] = False,
1367+
is_cdn: Optional[bool] = False,
13521368
export_authorization: Optional[bool] = True,
13531369
server_address: Optional[str] = None,
1354-
port: Optional[int] = None
1370+
port: Optional[int] = None,
1371+
temporary: Optional[bool] = False
13551372
) -> "Session":
13561373
"""Get existing session or create a new one.
13571374
@@ -1362,6 +1379,9 @@ async def get_session(
13621379
is_media (``bool``, *optional*):
13631380
Pass True to get or create a media session.
13641381
1382+
is_cdn (``bool``, *optional*):
1383+
Pass True to get or create a cdn session.
1384+
13651385
export_authorization (``bool``, *optional*):
13661386
Pass True to export authorization after creating the session.
13671387
Used only when creating a new session.
@@ -1373,59 +1393,16 @@ async def get_session(
13731393
port (``int``, *optional*):
13741394
Custom port to connect to.
13751395
Used only when creating a new session.
1376-
"""
1377-
if dc_id == await self.storage.dc_id() and not is_media:
1378-
return self.session
1379-
1380-
sessions = self.media_sessions if is_media else self.sessions
13811396
1397+
temporary (``bool``, *optional*):
1398+
Temporary sessions for perticular invoke.
1399+
Used only when uploading/downloading and don't forget to stop it.
1400+
"""
13821401
async with self.sessions_lock:
1383-
if sessions.get(dc_id):
1384-
return sessions[dc_id]
1385-
1386-
dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6)
1387-
1388-
session = self.media_sessions[dc_id] = Session(
1389-
self,
1390-
dc_id,
1391-
server_address or dc_option.ip_address,
1392-
port or dc_option.port,
1393-
await Auth(
1394-
self,
1395-
dc_id,
1396-
server_address or dc_option.ip_address,
1397-
port or dc_option.port,
1398-
await self.storage.test_mode()
1399-
).create(),
1400-
await self.storage.test_mode(), is_media=is_media
1401-
)
1402-
1403-
await session.start()
1404-
1405-
if export_authorization:
1406-
for _ in range(3):
1407-
exported_auth = await self.invoke(
1408-
raw.functions.auth.ExportAuthorization(
1409-
dc_id=dc_id
1410-
)
1411-
)
1412-
1413-
try:
1414-
await session.invoke(
1415-
raw.functions.auth.ImportAuthorization(
1416-
id=exported_auth.id,
1417-
bytes=exported_auth.bytes
1418-
)
1419-
)
1420-
except AuthBytesInvalid:
1421-
continue
1422-
else:
1423-
break
1424-
else:
1425-
await session.stop()
1426-
raise AuthBytesInvalid
1427-
1428-
return session
1402+
return await self._get_session(
1403+
dc_id, is_media, is_cdn, export_authorization,
1404+
server_address, port, temporary
1405+
)
14291406

14301407
async def get_dc_option(
14311408
self,

‎pyrogram/methods/advanced/save_file.py‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -144,20 +144,7 @@ async def worker(session):
144144
md5_sum = md5() if not is_big and not is_missing_part else None
145145
dc_id = await self.storage.dc_id()
146146

147-
session = self.media_sessions.get(dc_id)
148-
149-
if not session:
150-
dc_option = await self.get_dc_option(dc_id, is_media=True, ipv6=self.ipv6)
151-
152-
session = self.media_sessions[dc_id] = Session(
153-
self,
154-
dc_id,
155-
dc_option.ip_address,
156-
dc_option.port,
157-
await self.storage.auth_key(),
158-
await self.storage.test_mode(), is_media=True
159-
)
160-
await session.start()
147+
session = await self.get_session(dc_id, is_media=True)
161148

162149
workers = [self.loop.create_task(worker(session)) for _ in range(workers_count)]
163150
queue = asyncio.Queue(1)

0 commit comments

Comments
 (0)