mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2024-09-21 19:28:54 +03:00
83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
from hashlib import sha256
|
|
|
|
from ..tl import Session
|
|
from ..tl.functions.upload import GetCdnFileRequest, ReuploadCdnFileRequest
|
|
from ..tl.types.upload import CdnFileReuploadNeeded, CdnFile
|
|
from ..crypto import AESModeCTR
|
|
from ..errors import CdnFileTamperedError
|
|
|
|
|
|
class CdnDecrypter:
|
|
"""Used when downloading a file results in a 'FileCdnRedirect' to
|
|
both prepare the redirect, decrypt the file as it downloads, and
|
|
ensure the file hasn't been tampered. https://core.telegram.org/cdn
|
|
"""
|
|
def __init__(self, cdn_client, file_token, cdn_aes, cdn_file_hashes):
|
|
self.client = cdn_client
|
|
self.file_token = file_token
|
|
self.cdn_aes = cdn_aes
|
|
self.cdn_file_hashes = cdn_file_hashes
|
|
|
|
@staticmethod
|
|
def prepare_decrypter(client, cdn_client, cdn_redirect):
|
|
"""Prepares a CDN decrypter, returning (decrypter, file data).
|
|
'client' should be an existing client not connected to a CDN.
|
|
'cdn_client' should be an already-connected TelegramBareClient
|
|
with the auth key already created.
|
|
"""
|
|
cdn_aes = AESModeCTR(
|
|
key=cdn_redirect.encryption_key,
|
|
# 12 first bytes of the IV..4 bytes of the offset (0, big endian)
|
|
iv=cdn_redirect.encryption_iv[:12] + bytes(4)
|
|
)
|
|
|
|
# We assume that cdn_redirect.cdn_file_hashes are ordered by offset,
|
|
# and that there will be enough of these to retrieve the whole file.
|
|
decrypter = CdnDecrypter(
|
|
cdn_client, cdn_redirect.file_token,
|
|
cdn_aes, cdn_redirect.cdn_file_hashes
|
|
)
|
|
|
|
cdn_file = cdn_client(GetCdnFileRequest(
|
|
file_token=cdn_redirect.file_token,
|
|
offset=cdn_redirect.cdn_file_hashes[0].offset,
|
|
limit=cdn_redirect.cdn_file_hashes[0].limit
|
|
))
|
|
if isinstance(cdn_file, CdnFileReuploadNeeded):
|
|
# We need to use the original client here
|
|
client(ReuploadCdnFileRequest(
|
|
file_token=cdn_redirect.file_token,
|
|
request_token=cdn_file.request_token
|
|
))
|
|
|
|
# We want to always return a valid upload.CdnFile
|
|
cdn_file = decrypter.get_file()
|
|
else:
|
|
cdn_file.bytes = decrypter.cdn_aes.encrypt(cdn_file.bytes)
|
|
cdn_hash = decrypter.cdn_file_hashes.pop(0)
|
|
decrypter.check(cdn_file.bytes, cdn_hash)
|
|
|
|
return decrypter, cdn_file
|
|
|
|
def get_file(self):
|
|
"""Calls GetCdnFileRequest and decrypts its bytes.
|
|
Also ensures that the file hasn't been tampered.
|
|
"""
|
|
if self.cdn_file_hashes:
|
|
cdn_hash = self.cdn_file_hashes.pop(0)
|
|
cdn_file = self.client(GetCdnFileRequest(
|
|
self.file_token, cdn_hash.offset, cdn_hash.limit
|
|
))
|
|
cdn_file.bytes = self.cdn_aes.encrypt(cdn_file.bytes)
|
|
self.check(cdn_file.bytes, cdn_hash)
|
|
else:
|
|
cdn_file = CdnFile(bytes(0))
|
|
|
|
return cdn_file
|
|
|
|
@staticmethod
|
|
def check(data, cdn_hash):
|
|
"""Checks the integrity of the given data"""
|
|
if sha256(data).digest() != cdn_hash.hash:
|
|
raise CdnFileTamperedError()
|