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