diff --git a/WebStreamer/server/stream_routes.py b/WebStreamer/server/stream_routes.py index 6d560de..93a1249 100644 --- a/WebStreamer/server/stream_routes.py +++ b/WebStreamer/server/stream_routes.py @@ -74,36 +74,20 @@ async def media_streamer(request: web.Request, db_id: str): file_size = file_id.file_size - if range_header: - from_bytes, until_bytes = range_header.replace("bytes=", "").split("-") - from_bytes = int(from_bytes) - until_bytes = int(until_bytes) if until_bytes else file_size - 1 - else: - from_bytes = request.http_range.start or 0 - until_bytes = (request.http_range.stop or file_size) - 1 - - logging.debug(f"from_bytes: {from_bytes} until_bytes: {until_bytes}") - if from_bytes <10 and until_bytes >200: - await db.increment_dl_count(file_id.org_id) - - if (until_bytes > file_size) or (from_bytes < 0) or (until_bytes < from_bytes): + try: + offset = request.http_range.start or 0 + limit = request.http_range.stop or file_size + if (limit > file_size) or (offset < 0) or (limit < offset): + raise ValueError("range not in acceptable format") + except ValueError: return web.Response( status=416, body="416: Range not satisfiable", headers={"Content-Range": f"bytes */{file_size}"}, ) - chunk_size = 1024 * 1024 - until_bytes = min(until_bytes, file_size - 1) - - offset = from_bytes - (from_bytes % chunk_size) - first_part_cut = from_bytes - offset - last_part_cut = until_bytes % chunk_size + 1 - - req_length = until_bytes - from_bytes + 1 - part_count = math.ceil(until_bytes / chunk_size) - math.floor(offset / chunk_size) body = tg_connect.yield_file( - file_id, offset, first_part_cut, last_part_cut, part_count, chunk_size, multi_clients + file_id, offset, limit, multi_clients ) mime_type = file_id.mime_type @@ -121,8 +105,8 @@ async def media_streamer(request: web.Request, db_id: str): body=body, headers={ "Content-Type": f"{mime_type}", - "Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}", - "Content-Length": str(req_length), + "Content-Range": f"bytes {offset}-{limit}/{file_size}", + "Content-Length": str(limit - offset), "Content-Disposition": f'{disposition}; filename="{file_name}"', "Accept-Ranges": "bytes", }, diff --git a/WebStreamer/utils/custom_dl.py b/WebStreamer/utils/custom_dl.py index d1d9bb1..f8df959 100644 --- a/WebStreamer/utils/custom_dl.py +++ b/WebStreamer/utils/custom_dl.py @@ -3,8 +3,9 @@ import asyncio import logging # import heapq +import math from typing import Dict, Union -from WebStreamer.bot import work_loads +from WebStreamer.bot import StreamBot, work_loads from pyrogram import Client, utils, raw from .file_properties import get_file_ids from pyrogram.session import Session, Auth @@ -164,10 +165,7 @@ async def yield_file( self, file_id: FileId, offset: int, - first_part_cut: int, - last_part_cut: int, - part_count: int, - chunk_size: int, + limit: int, multi_clients ) -> Union[str, None]: """ @@ -178,47 +176,20 @@ async def yield_file( client = multi_clients[file_id.index] work_loads[file_id.index] += 1 logging.debug(f"Starting to yielding file with client {file_id.index}.") - media_session = await self.generate_media_session(client, file_id) + try: + if offset < 0: + if file_id.file_size == 0: + raise ValueError("Negative offsets are not supported for file ids, pass a Message object instead") - current_part = 1 + chunks = math.ceil(file_id.file_size / 1024 / 1024) + offset += chunks - location = await self.get_location(file_id) + async for chunk in client.get_file(file_id, file_id.file_size, limit, offset): + yield chunk - try: - r = await media_session.invoke( - raw.functions.upload.GetFile( - location=location, offset=offset, limit=chunk_size - ), - ) - if isinstance(r, raw.types.upload.File): - while True: - chunk = r.bytes - if not chunk: - break - elif part_count == 1: - yield chunk[first_part_cut:last_part_cut] - elif current_part == 1: - yield chunk[first_part_cut:] - elif current_part == part_count: - yield chunk[:last_part_cut] - else: - yield chunk - - current_part += 1 - offset += chunk_size - - if current_part > part_count: - break - - r = await media_session.invoke( - raw.functions.upload.GetFile( - location=location, offset=offset, limit=chunk_size - ), - ) except (TimeoutError, AttributeError): pass finally: - logging.debug(f"Finished yielding file with {current_part} parts.") work_loads[file_id.index] -= 1