mir/pyrogram
2
0
Fork 0
mirror of https://github.com/pyrogram/pyrogram synced 2025-08-28 12:57:52 +00:00
Code Issues Releases Wiki Activity

Rewrite get_file

Browse Source
This commit is contained in:
Dan 2018-02-18 15:03:33 +01:00
parent aaf1b5c862
commit 3bd3d99e6a
1 changed files with 62 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

96
pyrogram/client/client.py
View File

@ -1723,11 +1723,7 @@ class Client:
volume_id: int = None, volume_id: int = None,
local_id: int = None, local_id: int = None,
secret: int = None, secret: int = None,
version: int = 0): version: int = 0) -> str:
# TODO: Refine
# TODO: Use proper file name and extension
# TODO: Remove redundant code
if dc_id != self.dc_id: if dc_id != self.dc_id:
exported_auth = self.send( exported_auth = self.send(
functions.auth.ExportAuthorization( functions.auth.ExportAuthorization(
@ -1763,19 +1759,23 @@ class Client:
session.start() session.start()
if volume_id: # Photos are accessed by volume_id, local_id, secret if volume_id: # Photos are accessed by volume_id, local_id, secret
file_name = "_".join(str(i) for i in [dc_id, volume_id, local_id, secret])
location = types.InputFileLocation( location = types.InputFileLocation(
volume_id=volume_id, volume_id=volume_id,
local_id=local_id, local_id=local_id,
secret=secret secret=secret
) )
else: # Any other file can be more easily accessed by id and access_hash else: # Any other file can be more easily accessed by id and access_hash
file_name = "_".join(str(i) for i in [dc_id, id, access_hash, version])
location = types.InputDocumentFileLocation( location = types.InputDocumentFileLocation(
id=id, id=id,
access_hash=access_hash, access_hash=access_hash,
version=version version=version
) )
limit = 512 * 1024 limit = 1024 * 1024
offset = 0 offset = 0
try: try:
@ -1788,7 +1788,7 @@ class Client:
) )
if isinstance(r, types.upload.File): if isinstance(r, types.upload.File):
with open("_".join([str(id), str(access_hash), str(version)]) + ".jpg", "wb") as f: with open(file_name, "wb") as f:
while True: while True:
chunk = r.bytes chunk = r.bytes
@ -1796,6 +1796,9 @@ class Client:
break break
f.write(chunk) f.write(chunk)
f.flush()
os.fsync(f.fileno())
offset += limit offset += limit
r = session.send( r = session.send(
@ -1805,6 +1808,7 @@ class Client:
limit=limit limit=limit
) )
) )
if isinstance(r, types.upload.FileCdnRedirect): if isinstance(r, types.upload.FileCdnRedirect):
cdn_session = Session( cdn_session = Session(
r.dc_id, r.dc_id,
@ -1818,39 +1822,63 @@ class Client:
cdn_session.start() cdn_session.start()
try: try:
with open("_".join([str(id), str(access_hash), str(version)]) + ".jpg", "wb") as f: r2 = cdn_session.send(
while True: functions.upload.GetCdnFile(
r2 = cdn_session.send( location=location,
functions.upload.GetCdnFile( file_token=r.file_token,
location=location, offset=offset,
file_token=r.file_token, limit=limit
offset=offset, )
limit=limit )
)
if isinstance(r2, types.upload.CdnFileReuploadNeeded):
session.send(
functions.upload.ReuploadCdnFile(
file_token=r.file_token,
request_token=r2.request_token
) )
)
if isinstance(r2, types.upload.CdnFileReuploadNeeded): else:
session.send( with open(file_name, "wb") as f:
functions.upload.ReuploadCdnFile( while True:
file_token=r.file_token, if not isinstance(r2, types.upload.CdnFile):
request_token=r2.request_token
)
)
continue
elif isinstance(r2, types.upload.CdnFile):
chunk = r2.bytes
if not chunk:
break break
# https://core.telegram.org/cdn#decrypting-files chunk = r2.bytes
decrypted_chunk = AES.ctr_decrypt(chunk, r.encryption_key, r.encryption_iv, offset)
# TODO: https://core.telegram.org/cdn#verifying-files # https://core.telegram.org/cdn#decrypting-files
# TODO: Save to temp file, flush each chunk, rename to full if everything is ok decrypted_chunk = AES.ctr_decrypt(
chunk,
r.encryption_key,
r.encryption_iv,
offset
)
hashes = session.send(
functions.upload.GetCdnFileHashes(
r.file_token,
offset
)
)
for i, h in enumerate(hashes):
cdn_chunk = decrypted_chunk[h.limit * i: h.limit * (i + 1)]
assert h.hash == sha256(cdn_chunk).digest(), "Invalid CDN hash part {}".format(i)
f.write(decrypted_chunk) f.write(decrypted_chunk)
f.flush()
os.fsync(f.fileno())
offset += limit offset += limit
r2 = cdn_session.send(
functions.upload.GetCdnFile(
location=location,
file_token=r.file_token,
offset=offset,
limit=limit
)
)
except Exception as e: except Exception as e:
log.error(e) log.error(e)
finally: finally:
@ -1858,7 +1886,7 @@ class Client:
except Exception as e: except Exception as e:
log.error(e) log.error(e)
else: else:
return True return file_name
finally: finally:
session.stop() session.stop()