async def recv(self, length: int = 0) -> Optional[bytes]:
length = await super().recv(4)
if length is None:
return None
length = aes.ctr256_decrypt(length, *self.decrypt)
data = await super().recv(unpack("<i", length)[0])
if data is None:
return None
return aes.ctr256_decrypt(data, *self.decrypt)
async def recv(self, length: int = 0) -> bytes or None:
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)
data = await super().recv(int.from_bytes(length, "little") * 4)
if data is None:
return None
return aes.ctr256_decrypt(data, *self.decrypt)
async def recv(self, length: int = 0) -> Optional[bytes]:
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)
data = await super().recv(int.from_bytes(length, "little") * 4)
if data is None:
return None
return await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_decrypt, data, *self.decrypt)
async def recv(self, length: int = 0) -> Optional[bytes]:
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)
data = await super().recv(int.from_bytes(length, "little") * 4)
if data is None:
return None
return await utils.maybe_run_in_executor(aes.ctr256_decrypt, data,
len(data), self.loop,
*self.decrypt)
async def get_file(self,
media_type: int,
dc_id: int,
document_id: int,
access_hash: int,
thumb_size: str,
peer_id: int,
peer_type: str,
peer_access_hash: int,
volume_id: int,
local_id: int,
file_ref: str,
file_size: int,
is_big: bool,
progress: callable,
progress_args: tuple = ()) -> str:
async with self.media_sessions_lock:
session = self.media_sessions.get(dc_id, None)
if session is None:
if dc_id != await self.storage.dc_id():
session = Session(self,
dc_id,
await
Auth(self, dc_id, await
self.storage.test_mode()).create(),
await self.storage.test_mode(),
is_media=True)
await session.start()
for _ in range(3):
exported_auth = await self.send(
raw.functions.auth.ExportAuthorization(dc_id=dc_id)
)
try:
await session.send(
raw.functions.auth.ImportAuthorization(
id=exported_auth.id,
bytes=exported_auth.bytes))
except AuthBytesInvalid:
continue
else:
break
else:
await session.stop()
raise AuthBytesInvalid
else:
session = Session(self,
dc_id,
await self.storage.auth_key(),
await self.storage.test_mode(),
is_media=True)
await session.start()
self.media_sessions[dc_id] = session
file_ref = utils.decode_file_ref(file_ref)
if media_type == 1:
if peer_type == "user":
peer = raw.types.InputPeerUser(user_id=peer_id,
access_hash=peer_access_hash)
elif peer_type == "chat":
peer = raw.types.InputPeerChat(chat_id=peer_id)
else:
peer = raw.types.InputPeerChannel(channel_id=peer_id,
access_hash=peer_access_hash)
location = raw.types.InputPeerPhotoFileLocation(
peer=peer,
volume_id=volume_id,
local_id=local_id,
big=is_big or None)
elif media_type in (0, 2):
location = raw.types.InputPhotoFileLocation(
id=document_id,
access_hash=access_hash,
file_reference=file_ref,
thumb_size=thumb_size)
elif media_type == 14:
location = raw.types.InputDocumentFileLocation(
id=document_id,
access_hash=access_hash,
file_reference=file_ref,
thumb_size=thumb_size)
else:
location = raw.types.InputDocumentFileLocation(
id=document_id,
access_hash=access_hash,
file_reference=file_ref,
thumb_size="")
limit = 1024 * 1024
offset = 0
file_name = ""
try:
r = await session.send(raw.functions.upload.GetFile(
location=location, offset=offset, limit=limit),
sleep_threshold=30)
if isinstance(r, raw.types.upload.File):
with tempfile.NamedTemporaryFile("wb", delete=False) as f:
file_name = f.name
while True:
chunk = r.bytes
if not chunk:
break
f.write(chunk)
offset += limit
if progress:
func = functools.partial(
progress,
min(offset, file_size) if file_size != 0 else
offset, file_size, *progress_args)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(
self.executor, func)
r = await session.send(raw.functions.upload.GetFile(
location=location, offset=offset, limit=limit),
sleep_threshold=30)
elif isinstance(r, raw.types.upload.FileCdnRedirect):
async with self.media_sessions_lock:
cdn_session = self.media_sessions.get(r.dc_id, None)
if cdn_session is None:
cdn_session = Session(
self,
r.dc_id,
await Auth(self, r.dc_id, await
self.storage.test_mode()).create(),
await self.storage.test_mode(),
is_media=True,
is_cdn=True)
await cdn_session.start()
self.media_sessions[r.dc_id] = cdn_session
try:
with tempfile.NamedTemporaryFile("wb", delete=False) as f:
file_name = f.name
while True:
r2 = await cdn_session.send(
raw.functions.upload.GetCdnFile(
file_token=r.file_token,
offset=offset,
limit=limit))
if isinstance(
r2,
raw.types.upload.CdnFileReuploadNeeded):
try:
await session.send(
raw.functions.upload.ReuploadCdnFile(
file_token=r.file_token,
request_token=r2.request_token))
except VolumeLocNotFound:
break
else:
continue
chunk = r2.bytes
# https://core.telegram.org/cdn#decrypting-files
decrypted_chunk = aes.ctr256_decrypt(
chunk, r.encryption_key,
bytearray(r.encryption_iv[:-4] +
(offset // 16).to_bytes(4, "big")))
hashes = await session.send(
raw.functions.upload.GetCdnFileHashes(
file_token=r.file_token, offset=offset))
# https://core.telegram.org/cdn#verifying-files
for i, h in enumerate(hashes):
cdn_chunk = decrypted_chunk[h.limit *
i:h.limit *
(i + 1)]
assert h.hash == sha256(cdn_chunk).digest(
), f"Invalid CDN hash part {i}"
f.write(decrypted_chunk)
offset += limit
if progress:
func = functools.partial(
progress,
min(offset, file_size) if file_size != 0
else offset, file_size, *progress_args)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(
self.executor, func)
if len(chunk) < limit:
break
except Exception as e:
raise e
except Exception as e:
if not isinstance(e, pyrogram.StopTransmission):
log.error(e, exc_info=True)
try:
os.remove(file_name)
except OSError:
pass
return ""
else:
return file_name
async def get_file(
self,
file_id: FileId,
file_size: int = 0,
limit: int = 0,
offset: int = 0,
progress: Callable = None,
progress_args: tuple = ()
) -> Optional[AsyncGenerator[bytes, None]]:
dc_id = file_id.dc_id
async with self.media_sessions_lock:
session = self.media_sessions.get(dc_id, None)
if session is None:
if dc_id != await self.storage.dc_id():
session = Session(self,
dc_id,
await
Auth(self, dc_id, await
self.storage.test_mode()).create(),
await self.storage.test_mode(),
is_media=True)
await session.start()
for _ in range(3):
exported_auth = await self.invoke(
raw.functions.auth.ExportAuthorization(dc_id=dc_id)
)
try:
await session.invoke(
raw.functions.auth.ImportAuthorization(
id=exported_auth.id,
bytes=exported_auth.bytes))
except AuthBytesInvalid:
continue
else:
break
else:
await session.stop()
raise AuthBytesInvalid
else:
session = Session(self,
dc_id,
await self.storage.auth_key(),
await self.storage.test_mode(),
is_media=True)
await session.start()
self.media_sessions[dc_id] = session
file_type = file_id.file_type
if file_type == FileType.CHAT_PHOTO:
if file_id.chat_id > 0:
peer = raw.types.InputPeerUser(
user_id=file_id.chat_id,
access_hash=file_id.chat_access_hash)
else:
if file_id.chat_access_hash == 0:
peer = raw.types.InputPeerChat(chat_id=-file_id.chat_id)
else:
peer = raw.types.InputPeerChannel(
channel_id=utils.get_channel_id(file_id.chat_id),
access_hash=file_id.chat_access_hash)
location = raw.types.InputPeerPhotoFileLocation(
peer=peer,
photo_id=file_id.media_id,
big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG)
elif file_type == FileType.PHOTO:
location = raw.types.InputPhotoFileLocation(
id=file_id.media_id,
access_hash=file_id.access_hash,
file_reference=file_id.file_reference,
thumb_size=file_id.thumbnail_size)
else:
location = raw.types.InputDocumentFileLocation(
id=file_id.media_id,
access_hash=file_id.access_hash,
file_reference=file_id.file_reference,
thumb_size=file_id.thumbnail_size)
current = 0
total = abs(limit) or (1 << 31) - 1
chunk_size = 1024 * 1024
offset_bytes = abs(offset) * chunk_size
try:
r = await session.invoke(raw.functions.upload.GetFile(
location=location, offset=offset_bytes, limit=chunk_size),
sleep_threshold=30)
if isinstance(r, raw.types.upload.File):
while True:
chunk = r.bytes
yield chunk
current += 1
offset_bytes += chunk_size
if progress:
func = functools.partial(
progress,
min(offset_bytes, file_size) if file_size != 0 else
offset_bytes, file_size, *progress_args)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(
self.executor, func)
if len(chunk) < chunk_size or current >= total:
break
r = await session.invoke(raw.functions.upload.GetFile(
location=location,
offset=offset_bytes,
limit=chunk_size),
sleep_threshold=30)
elif isinstance(r, raw.types.upload.FileCdnRedirect):
async with self.media_sessions_lock:
cdn_session = self.media_sessions.get(r.dc_id, None)
if cdn_session is None:
cdn_session = Session(
self,
r.dc_id,
await Auth(self, r.dc_id, await
self.storage.test_mode()).create(),
await self.storage.test_mode(),
is_media=True,
is_cdn=True)
await cdn_session.start()
self.media_sessions[r.dc_id] = cdn_session
try:
while True:
r2 = await cdn_session.invoke(
raw.functions.upload.GetCdnFile(
file_token=r.file_token,
offset=offset_bytes,
limit=chunk_size))
if isinstance(r2,
raw.types.upload.CdnFileReuploadNeeded):
try:
await session.invoke(
raw.functions.upload.ReuploadCdnFile(
file_token=r.file_token,
request_token=r2.request_token))
except VolumeLocNotFound:
break
else:
continue
chunk = r2.bytes
# https://core.telegram.org/cdn#decrypting-files
decrypted_chunk = aes.ctr256_decrypt(
chunk, r.encryption_key,
bytearray(r.encryption_iv[:-4] +
(offset_bytes // 16).to_bytes(4, "big")))
hashes = await session.invoke(
raw.functions.upload.GetCdnFileHashes(
file_token=r.file_token, offset=offset_bytes))
# https://core.telegram.org/cdn#verifying-files
for i, h in enumerate(hashes):
cdn_chunk = decrypted_chunk[h.limit * i:h.limit *
(i + 1)]
CDNFileHashMismatch.check(
h.hash == sha256(cdn_chunk).digest())
yield decrypted_chunk
current += 1
offset_bytes += chunk_size
if progress:
func = functools.partial(
progress,
min(offset_bytes, file_size) if file_size != 0
else offset_bytes, file_size, *progress_args)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(
self.executor, func)
if len(chunk) < chunk_size or current >= total:
break
except Exception as e:
raise e
except pyrogram.StopTransmission:
raise
except Exception as e:
log.error(e, exc_info=True)
async def get_streaming_file(
self,
file_id: FileId,
file_size: int,
progress: callable,
progress_args: tuple = ()
) -> AsyncGenerator:
dc_id = file_id.dc_id
session = await self.get_dc_session(dc_id)
location = await self.get_file_location(file_id)
limit = 1024 * 1024
offset = 0
try:
r = await session.send(
raw.functions.upload.GetFile(
location=location,
offset=offset,
limit=limit
),
sleep_threshold=30
)
if isinstance(r, raw.types.upload.File):
while True:
chunk = r.bytes
if not chunk:
return
yield chunk
offset += limit
if progress:
func = functools.partial(
progress,
min(offset, file_size)
if file_size != 0
else offset,
file_size,
*progress_args
)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(self.executor, func)
r = await session.send(
raw.functions.upload.GetFile(
location=location,
offset=offset,
limit=limit
),
sleep_threshold=30
)
elif isinstance(r, raw.types.upload.FileCdnRedirect):
async with self.media_sessions_lock:
cdn_session = self.media_sessions.get(r.dc_id, None)
if cdn_session is None:
cdn_session = Session(
self, r.dc_id, await Auth(self, r.dc_id, await self.storage.test_mode()).create(),
await self.storage.test_mode(), is_media=True, is_cdn=True
)
await cdn_session.start()
self.media_sessions[r.dc_id] = cdn_session
while True:
r2 = await cdn_session.send(
raw.functions.upload.GetCdnFile(
file_token=r.file_token,
offset=offset,
limit=limit
)
)
if isinstance(r2, raw.types.upload.CdnFileReuploadNeeded):
try:
await session.send(
raw.functions.upload.ReuploadCdnFile(
file_token=r.file_token,
request_token=r2.request_token
)
)
except VolumeLocNotFound:
return
else:
continue
chunk = r2.bytes
# https://core.telegram.org/cdn#decrypting-files
decrypted_chunk = aes.ctr256_decrypt(
chunk,
r.encryption_key,
bytearray(
r.encryption_iv[:-4]
+ (offset // 16).to_bytes(4, "big")
)
)
hashes = await session.send(
raw.functions.upload.GetCdnFileHashes(
file_token=r.file_token,
offset=offset
)
)
# https://core.telegram.org/cdn#verifying-files
for i, h in enumerate(hashes):
cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)]
assert h.hash == sha256(cdn_chunk).digest(), f"Invalid CDN hash part {i}"
yield decrypted_chunk
offset += limit
if progress:
func = functools.partial(
progress,
min(offset, file_size) if file_size != 0 else offset,
file_size,
*progress_args
)
if inspect.iscoroutinefunction(progress):
await func()
else:
await self.loop.run_in_executor(self.executor, func)
if len(chunk) < limit:
return
except Exception as e:
if not isinstance(e, pyrogram.StopTransmission):
log.error(e, exc_info=True)
return
async def get_file(
self,
file_id: FileId,
filename: str,
) -> str:
dc_id = file_id.dc_id
async with self.client.media_sessions_lock:
session = self.client.media_sessions.get(dc_id, None)
if session is None:
if dc_id != await self.client.storage.dc_id():
session = Session(
self.client,
dc_id,
await Auth(self.client, dc_id, await
self.client.storage.test_mode()).create(),
await self.client.storage.test_mode(),
is_media=True)
await session.start()
for _ in range(3):
exported_auth = await self.client.send(
raw.functions.auth.ExportAuthorization(dc_id=dc_id)
)
try:
await session.send(
raw.functions.auth.ImportAuthorization(
id=exported_auth.id,
bytes=exported_auth.bytes))
except AuthBytesInvalid:
continue
else:
break
else:
await session.stop()
raise AuthBytesInvalid
else:
session = Session(self.client,
dc_id,
await self.client.storage.auth_key(),
await self.client.storage.test_mode(),
is_media=True)
await session.start()
self.client.media_sessions[dc_id] = session
file_type = file_id.file_type
if file_type == FileType.CHAT_PHOTO:
if file_id.chat_id > 0:
peer = raw.types.InputPeerUser(
user_id=file_id.chat_id,
access_hash=file_id.chat_access_hash)
else:
if file_id.chat_access_hash == 0:
peer = raw.types.InputPeerChat(chat_id=-file_id.chat_id)
else:
peer = raw.types.InputPeerChannel(
channel_id=utils.get_channel_id(file_id.chat_id),
access_hash=file_id.chat_access_hash)
location = raw.types.InputPeerPhotoFileLocation(
peer=peer,
volume_id=file_id.volume_id,
local_id=file_id.local_id,
big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG)
elif file_type == FileType.PHOTO:
location = raw.types.InputPhotoFileLocation(
id=file_id.media_id,
access_hash=file_id.access_hash,
file_reference=file_id.file_reference,
thumb_size=file_id.thumbnail_size)
else:
location = raw.types.InputDocumentFileLocation(
id=file_id.media_id,
access_hash=file_id.access_hash,
file_reference=file_id.file_reference,
thumb_size=file_id.thumbnail_size)
limit = 1024 * 1024
offset = 0
file_name = ""
try:
r = await session.send(raw.functions.upload.GetFile(
location=location, offset=offset, limit=limit),
sleep_threshold=30)
if isinstance(r, raw.types.upload.File):
#with tempfile.NamedTemporaryFile("wb", delete=False) as f:
with open(filename, 'wb') as f:
file_name = filename
while True:
chunk = r.bytes
if not chunk:
break
f.write(chunk)
offset += limit
r = await session.send(raw.functions.upload.GetFile(
location=location, offset=offset, limit=limit),
sleep_threshold=30)
elif isinstance(r, raw.types.upload.FileCdnRedirect):
async with self.client.media_sessions_lock:
cdn_session = self.client.media_sessions.get(r.dc_id, None)
if cdn_session is None:
cdn_session = Session(
self.client,
r.dc_id,
await
Auth(self.client, r.dc_id, await
self.client.storage.test_mode()).create(),
await self.client.storage.test_mode(),
is_media=True,
is_cdn=True)
await cdn_session.start()
self.client.media_sessions[r.dc_id] = cdn_session
try:
with open(filename, 'wb') as f:
file_name = f
while True:
r2 = await cdn_session.send(
raw.functions.upload.GetCdnFile(
file_token=r.file_token,
offset=offset,
limit=limit))
if isinstance(
r2,
raw.types.upload.CdnFileReuploadNeeded):
try:
await session.send(
raw.functions.upload.ReuploadCdnFile(
file_token=r.file_token,
request_token=r2.request_token))
except VolumeLocNotFound:
break
else:
continue
chunk = r2.bytes
# https://core.telegram.org/cdn#decrypting-files
decrypted_chunk = aes.ctr256_decrypt(
chunk, r.encryption_key,
bytearray(r.encryption_iv[:-4] +
(offset // 16).to_bytes(4, "big")))
hashes = await session.send(
raw.functions.upload.GetCdnFileHashes(
file_token=r.file_token, offset=offset))
# https://core.telegram.org/cdn#verifying-files
for i, h in enumerate(hashes):
cdn_chunk = decrypted_chunk[h.limit *
i:h.limit *
(i + 1)]
assert h.hash == sha256(cdn_chunk).digest(
), f"Invalid CDN hash part {i}"
f.write(decrypted_chunk)
offset += limit
if len(chunk) < limit:
break
except Exception as e:
LOGGER.error(e, exc_info=True)
raise e
except Exception as e:
if not isinstance(e, pyrogram.StopTransmission):
LOGGER.error(str(e), exc_info=True)
try:
os.remove(file_name)
except OSError:
pass
return ""
else:
return file_name
