mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2025-08-02 11:10:18 +03:00
Add file upload and download
This commit is contained in:
painor
parent
4a51587c68
commit
5ee5c1dd7c
|
|
@ -1,5 +1,6 @@
|
|||
import hashlib
|
||||
import random
|
||||
from hashlib import sha1, sha256
|
||||
from hashlib import sha1, sha256, md5
|
||||
from time import time
|
||||
|
||||
from .. import utils
|
||||
|
|
@ -7,10 +8,11 @@ from ..crypto import AES
|
|||
from ..errors import SecurityError, EncryptionAlreadyDeclinedError
|
||||
from ..extensions import BinaryReader
|
||||
from ..network.mtprotostate import MTProtoState
|
||||
from ..tl.functions.messages import AcceptEncryptionRequest
|
||||
from ..tl.functions.messages import AcceptEncryptionRequest, SendEncryptedFileRequest
|
||||
from ..tl.functions.messages import GetDhConfigRequest, RequestEncryptionRequest, SendEncryptedServiceRequest, \
|
||||
DiscardEncryptionRequest, SendEncryptedRequest
|
||||
from ..tl.types import InputEncryptedChat, TypeEncryptedChat
|
||||
from ..tl.types import InputEncryptedChat, TypeEncryptedChat, EncryptedFile, InputEncryptedFileLocation, \
|
||||
InputEncryptedFile, InputFileBig, InputFile, InputEncryptedFileBigUploaded, InputEncryptedFileUploaded
|
||||
from ..tl.types.messages import DhConfigNotModified
|
||||
from ..tl.types.secret import *
|
||||
|
||||
|
|
@ -105,7 +107,7 @@ class SecretChatMethods:
|
|||
|
||||
async def rekey(self, peer):
|
||||
peer = self.get_secret_chat(peer)
|
||||
self._log.debug(f'Rekeying secret chat {peer}')
|
||||
self._log[__name__].debug(f'Rekeying secret chat {peer}')
|
||||
dh_config = await self.get_dh_config()
|
||||
a = int.from_bytes(os.urandom(256), 'big', signed=False)
|
||||
g_a = pow(dh_config.g, a, dh_config.p)
|
||||
|
|
@ -132,7 +134,7 @@ class SecretChatMethods:
|
|||
if my_exchange_id == other_exchange_id:
|
||||
peer.rekeying = [0]
|
||||
return
|
||||
self._log.debug(f'Accepting rekeying secret chat {peer}')
|
||||
self._log[__name__].debug(f'Accepting rekeying secret chat {peer}')
|
||||
dh_config = await self.get_dh_config()
|
||||
random_bytes = os.urandom(256)
|
||||
b = int.from_bytes(random_bytes, byteorder="big", signed=False)
|
||||
|
|
@ -159,7 +161,7 @@ class SecretChatMethods:
|
|||
if peer.rekeying[0] != 1 or not self.temp_rekeyed_secret_chats.get(action.exchange_id, None):
|
||||
peer.rekeying = [0]
|
||||
return
|
||||
self._log.debug(f'Committing rekeying secret chat {peer}')
|
||||
self._log[__name__].debug(f'Committing rekeying secret chat {peer}')
|
||||
dh_config = await self.get_dh_config()
|
||||
g_b = int.from_bytes(action.g_b, 'big', signed=False)
|
||||
self.check_g_a(g_b, dh_config.p)
|
||||
|
|
@ -198,7 +200,7 @@ class SecretChatMethods:
|
|||
await self(SendEncryptedServiceRequest(InputEncryptedChat(peer.id, peer.access_hash), message))
|
||||
raise SecurityError("Invalid Key fingerprint")
|
||||
|
||||
self._log.debug(f'Completing rekeying secret chat {peer}')
|
||||
self._log[__name__].debug(f'Completing rekeying secret chat {peer}')
|
||||
peer.rekeying = [0]
|
||||
peer.key = self.temp_rekeyed_secret_chats[action.exchange_id]
|
||||
peer.ttr = 100
|
||||
|
|
@ -207,9 +209,9 @@ class SecretChatMethods:
|
|||
message = DecryptedMessageService(action=DecryptedMessageActionNoop())
|
||||
message = await self.encrypt_secret_message(peer, message)
|
||||
await self(SendEncryptedServiceRequest(InputEncryptedChat(peer.id, peer.access_hash), message))
|
||||
self._log.debug(f'Secret chat {peer} rekeyed succrsfully')
|
||||
self._log[__name__].debug(f'Secret chat {peer} rekeyed succrsfully')
|
||||
|
||||
async def handle_decrypted_message(self, decrypted_message, peer: Chats):
|
||||
async def handle_decrypted_message(self, decrypted_message, peer: Chats, file):
|
||||
if isinstance(decrypted_message, (DecryptedMessageService, DecryptedMessageService8)):
|
||||
if isinstance(decrypted_message.action, DecryptedMessageActionRequestKey):
|
||||
await self.accept_rekey(peer, decrypted_message.action)
|
||||
|
|
@ -237,7 +239,7 @@ class SecretChatMethods:
|
|||
decrypted_message.action.end_seq_no -= peer.out_seq_no_x
|
||||
decrypted_message.action.start_seq_no //= 2
|
||||
decrypted_message.action.end_seq_no //= 2
|
||||
self._log.warning(f"Resending messages for {peer.id}")
|
||||
self._log[__name__].warning(f"Resending messages for {peer.id}")
|
||||
for seq, message in peer.outgoing:
|
||||
if decrypted_message.action.start_seq_no <= seq <= decrypted_message.action.end_seq_no:
|
||||
await self.send_secret_message(peer.id, message.message)
|
||||
|
|
@ -246,6 +248,7 @@ class SecretChatMethods:
|
|||
return decrypted_message
|
||||
elif isinstance(decrypted_message,
|
||||
(DecryptedMessage8, DecryptedMessage23, DecryptedMessage46, DecryptedMessage)):
|
||||
decrypted_message.file = file
|
||||
return decrypted_message
|
||||
elif isinstance(decrypted_message, DecryptedMessageLayer):
|
||||
# TODO add checks
|
||||
|
|
@ -255,13 +258,16 @@ class SecretChatMethods:
|
|||
if decrypted_message.layer >= 17 and time() - peer.created > 15:
|
||||
await self.notify_layer(peer)
|
||||
decrypted_message = decrypted_message.message
|
||||
return await self.handle_decrypted_message(decrypted_message, peer)
|
||||
return await self.handle_decrypted_message(decrypted_message, peer, file)
|
||||
|
||||
async def handle_encrypted_update(self, event):
|
||||
if not self.secret_chats.get(event.message.chat_id):
|
||||
self._log.debug("Secret chat not saved. skipping")
|
||||
self._log[__name__].debug("Secret chat not saved. skipping")
|
||||
return False
|
||||
message = event.message
|
||||
|
||||
file = getattr(message, 'file', None)
|
||||
|
||||
auth_key_id = struct.unpack('<q', message.bytes[:8])[0]
|
||||
peer = self.get_secret_chat(message.chat_id)
|
||||
if not peer.key.fingerprint or \
|
||||
|
|
@ -279,7 +285,7 @@ class SecretChatMethods:
|
|||
decrypted_message = self.decrypt_mtproto1(bytes.fromhex(message_key.hex()), message.chat_id,
|
||||
bytes.fromhex(encrypted_data.hex()))
|
||||
peer.mtproto = 1
|
||||
self._log.debug(f"Used MTProto 1 with chat {message.chat_id}")
|
||||
self._log[__name__].debug(f"Used MTProto 1 with chat {message.chat_id}")
|
||||
|
||||
else:
|
||||
try:
|
||||
|
|
@ -290,12 +296,12 @@ class SecretChatMethods:
|
|||
decrypted_message = self.decrypt_mtproto2(bytes.fromhex(message_key.hex()), message.chat_id,
|
||||
bytes.fromhex(encrypted_data.hex()))
|
||||
peer.mtproto = 2
|
||||
self._log.debug(f"Used MTProto 2 with chat {message.chat_id}")
|
||||
self._log[__name__].debug(f"Used MTProto 2 with chat {message.chat_id}")
|
||||
peer.ttr -= 1
|
||||
if (peer.ttr <= 0 or (time() - peer.updated) > 7 * 24 * 60 * 60) and peer.rekeying[0] == 0:
|
||||
await self.rekey(peer)
|
||||
peer.incoming[peer.in_seq_no] = message
|
||||
return await self.handle_decrypted_message(decrypted_message, peer)
|
||||
return await self.handle_decrypted_message(decrypted_message, peer, file)
|
||||
|
||||
async def encrypt_secret_message(self, peer, message):
|
||||
peer = self.get_secret_chat(peer)
|
||||
|
|
@ -335,6 +341,24 @@ class SecretChatMethods:
|
|||
aes_iv)
|
||||
return message
|
||||
|
||||
async def download_secret_media(self, message):
|
||||
if not message.file or not isinstance(message.file, EncryptedFile):
|
||||
return b""
|
||||
key_fingerprint = message.file.key_fingerprint
|
||||
key = message.media.key
|
||||
iv = message.media.iv
|
||||
digest = md5(key + iv).digest()
|
||||
|
||||
fingerprint = int.from_bytes(digest[:4], byteorder="little", signed=True) ^ int.from_bytes(digest[4:8],
|
||||
byteorder="little",
|
||||
signed=True)
|
||||
if fingerprint != key_fingerprint:
|
||||
raise SecurityError("Wrong fingerprint")
|
||||
media = await self.download_file(InputEncryptedFileLocation(message.file.id, message.file.access_hash),
|
||||
file=bytes)
|
||||
decrypted_data = AES.decrypt_ige(media, message.media.key, message.media.iv)
|
||||
return decrypted_data
|
||||
|
||||
async def send_secret_message(self, peer_id, message, ttl=0, reply_to_id=None):
|
||||
peer = self.get_secret_chat(peer_id)
|
||||
if peer.layer == 8:
|
||||
|
|
@ -348,6 +372,104 @@ class SecretChatMethods:
|
|||
SendEncryptedRequest(peer=peer.input_chat, data=data))
|
||||
return res
|
||||
|
||||
async def upload_secret_file(self, file):
|
||||
key = os.urandom(32)
|
||||
iv = os.urandom(32)
|
||||
digest = md5(key + iv).digest()
|
||||
fingerprint = int.from_bytes(digest[:4], byteorder="little", signed=True) ^ int.from_bytes(digest[4:8],
|
||||
byteorder="little",
|
||||
signed=True)
|
||||
|
||||
file = await self.upload_file(file, key=key, iv=iv)
|
||||
if isinstance(file, InputFileBig):
|
||||
file = InputEncryptedFileBigUploaded(file.id, file.parts, fingerprint)
|
||||
elif isinstance(file, InputFile):
|
||||
file = InputEncryptedFileUploaded(file.id, file.parts, "", fingerprint)
|
||||
|
||||
return file, fingerprint, key, iv
|
||||
|
||||
async def send_secret_document(self, peer, document, thumb: bytes, thumb_w: int, thumb_h: int, file_name: str,
|
||||
mime_type: str, size: int, attributes=None, ttl=0, caption=""):
|
||||
if attributes is None:
|
||||
attributes = []
|
||||
peer = self.get_secret_chat(peer)
|
||||
file, fingerprint, key, iv = await self.upload_secret_file(document)
|
||||
if peer.layer == 8:
|
||||
message = DecryptedMessage8(os.urandom(8), caption,
|
||||
DecryptedMessageMediaDocument23(thumb, thumb_w, thumb_h, file_name, mime_type,
|
||||
size, key, iv))
|
||||
elif peer.layer == 46:
|
||||
message = DecryptedMessage46(ttl, caption,
|
||||
media=DecryptedMessageMediaDocument(thumb, thumb_w, thumb_h, mime_type,
|
||||
size, key, iv, attributes, caption))
|
||||
else:
|
||||
message = DecryptedMessage(ttl, caption,
|
||||
media=DecryptedMessageMediaDocument(thumb, thumb_w, thumb_h, mime_type,
|
||||
size, key, iv, attributes, caption))
|
||||
data = await self.encrypt_secret_message(peer, message)
|
||||
res = await self(SendEncryptedFileRequest(peer.input_chat, data, file=file))
|
||||
return res
|
||||
|
||||
async def send_secret_audio(self, peer, audio, duration, mime_type, size, ttl=0, caption=""):
|
||||
peer = self.get_secret_chat(peer)
|
||||
file, fingerprint, key, iv = await self.upload_secret_file(audio)
|
||||
if peer.layer == 8:
|
||||
message = DecryptedMessage8(os.urandom(8), caption,
|
||||
DecryptedMessageMediaAudio8(duration, size, key, iv))
|
||||
elif peer.layer == 46:
|
||||
message = DecryptedMessage46(ttl, caption,
|
||||
media=DecryptedMessageMediaAudio(duration, mime_type, size, key, iv))
|
||||
else:
|
||||
message = DecryptedMessage(ttl, caption,
|
||||
media=DecryptedMessageMediaAudio(duration, mime_type, size, key, iv))
|
||||
data = await self.encrypt_secret_message(peer, message)
|
||||
res = await self(SendEncryptedFileRequest(peer.input_chat, data, file=file))
|
||||
return res
|
||||
|
||||
async def send_secret_video(self, peer, video, thumb: bytes, thumb_w: int, thumb_h: int, duration: int,
|
||||
mime_type: str, w: int,
|
||||
h: int, size, ttl=0, caption=""):
|
||||
peer = self.get_secret_chat(peer)
|
||||
file, fingerprint, key, iv = await self.upload_secret_file(video)
|
||||
|
||||
if peer.layer == 8:
|
||||
message = DecryptedMessage8(os.urandom(8), caption,
|
||||
DecryptedMessageMediaVideo8(thumb, thumb_w, thumb_h, duration, w, h, size, key,
|
||||
iv))
|
||||
elif peer.layer == 46:
|
||||
message = DecryptedMessage46(ttl, caption,
|
||||
media=DecryptedMessageMediaVideo(thumb, thumb_w, thumb_h, duration, mime_type,
|
||||
w, h, size, key, iv,
|
||||
caption))
|
||||
else:
|
||||
message = DecryptedMessage(ttl, caption,
|
||||
media=DecryptedMessageMediaVideo(thumb, thumb_w, thumb_h, duration, mime_type,
|
||||
w, h, size, key, iv,
|
||||
caption))
|
||||
data = await self.encrypt_secret_message(peer, message)
|
||||
res = await self(SendEncryptedFileRequest(peer.input_chat, data, file=file))
|
||||
return res
|
||||
|
||||
async def send_secret_photo(self, peer, image, thumb, thumb_w, thumb_h, w, h, size, caption="",
|
||||
ttl=0):
|
||||
peer = self.get_secret_chat(peer)
|
||||
|
||||
file, fingerprint, key, iv = await self.upload_secret_file(image)
|
||||
if peer.layer == 8:
|
||||
message = DecryptedMessage8(os.urandom(8), caption,
|
||||
DecryptedMessageMediaPhoto23(thumb, thumb_w, thumb_h, w, h, size, key, iv))
|
||||
elif peer.layer == 46:
|
||||
message = DecryptedMessage46(ttl, caption,
|
||||
media=DecryptedMessageMediaPhoto(thumb, thumb_w, thumb_h, w, h, size, key, iv,
|
||||
caption))
|
||||
else:
|
||||
message = DecryptedMessage(ttl, caption,
|
||||
media=DecryptedMessageMediaPhoto(thumb, thumb_w, thumb_h, w, h, size, key, iv,
|
||||
caption))
|
||||
data = await self.encrypt_secret_message(peer, message)
|
||||
res = await self(SendEncryptedFileRequest(peer.input_chat, data, file=file))
|
||||
return res
|
||||
|
||||
async def notify_layer(self, peer):
|
||||
if isinstance(peer, int):
|
||||
peer = self.secret_chats[peer]
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import typing
|
|||
import inspect
|
||||
from io import BytesIO
|
||||
|
||||
from ..crypto import AES
|
||||
from .. import utils, helpers, hints
|
||||
from ..tl import types, functions, custom
|
||||
|
||||
|
|
@ -17,13 +18,13 @@ try:
|
|||
except ImportError:
|
||||
PIL = None
|
||||
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from .telegramclient import TelegramClient
|
||||
|
||||
|
||||
class _CacheType:
|
||||
"""Like functools.partial but pretends to be the wrapped class."""
|
||||
|
||||
def __init__(self, cls):
|
||||
self._cls = cls
|
||||
|
||||
|
|
@ -327,7 +328,7 @@ class UploadMethods:
|
|||
file_handle, media, image = await self._file_to_media(
|
||||
file, force_document=force_document,
|
||||
progress_callback=progress_callback,
|
||||
attributes=attributes, allow_cache=allow_cache, thumb=thumb,
|
||||
attributes=attributes, allow_cache=allow_cache, thumb=thumb,
|
||||
voice_note=voice_note, video_note=video_note,
|
||||
supports_streaming=supports_streaming
|
||||
)
|
||||
|
|
@ -415,7 +416,9 @@ class UploadMethods:
|
|||
part_size_kb: float = None,
|
||||
file_name: str = None,
|
||||
use_cache: type = None,
|
||||
progress_callback: 'hints.ProgressCallback' = None) -> 'types.TypeInputFile':
|
||||
progress_callback: 'hints.ProgressCallback' = None,
|
||||
key: bytes = b"",
|
||||
iv: bytes = b"") -> 'types.TypeInputFile':
|
||||
"""
|
||||
Uploads a file to Telegram's servers, without sending it.
|
||||
|
||||
|
|
@ -452,6 +455,13 @@ class UploadMethods:
|
|||
A callback function accepting two parameters:
|
||||
``(sent bytes, total)``.
|
||||
|
||||
iv ('bytes', optional):
|
||||
In case of an encrypted upload (secret chats) an iv is supplied
|
||||
|
||||
key ('bytes', optional):
|
||||
In case of an encrypted upload (secret chats) a key is supplied
|
||||
|
||||
|
||||
Returns
|
||||
:tl:`InputFileBig` if the file size is larger than 10MB,
|
||||
`InputSizedFile <telethon.tl.custom.inputsizedfile.InputSizedFile>`
|
||||
|
|
@ -558,6 +568,10 @@ class UploadMethods:
|
|||
# Read the file by in chunks of size part_size
|
||||
part = stream.read(part_size)
|
||||
|
||||
# encryption part if needed
|
||||
if key and iv:
|
||||
part = AES.encrypt_ige(part, key, iv)
|
||||
|
||||
# The SavePartRequest is different depending on whether
|
||||
# the file is too large or not (over or less than 10MB)
|
||||
if is_large:
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user