TGDRIVE / utils /streamer /__init__.py
dragxd's picture
Initial commit
e1a2a92
import math, mimetypes
from fastapi.responses import StreamingResponse, Response
from utils.logger import Logger
from utils.streamer.custom_dl import ByteStreamer
from utils.streamer.file_properties import get_name
from utils.clients import (
get_client,
)
from urllib.parse import quote
logger = Logger(__name__)
class_cache = {}
async def media_streamer(channel: int, message_id: int, file_name: str, request):
global class_cache
range_header = request.headers.get("Range", 0)
faster_client = get_client()
if faster_client in class_cache:
tg_connect = class_cache[faster_client]
else:
tg_connect = ByteStreamer(faster_client)
class_cache[faster_client] = tg_connect
file_id = await tg_connect.get_file_properties(channel, message_id)
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 = 0
until_bytes = file_size - 1
if (until_bytes > file_size) or (from_bytes < 0) or (until_bytes < from_bytes):
return Response(
status_code=416,
content="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
)
disposition = "attachment"
mime_type = mimetypes.guess_type(file_name.lower())[0] or "application/octet-stream"
if (
"video/" in mime_type
or "audio/" in mime_type
or "image/" in mime_type
or "/html" in mime_type
):
disposition = "inline"
return StreamingResponse(
status_code=206 if range_header else 200,
content=body,
headers={
"Content-Type": f"{mime_type}",
"Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}",
"Content-Length": str(req_length),
"Content-Disposition": f'{disposition}; filename="{quote(file_name)}"',
"Accept-Ranges": "bytes",
},
media_type=mime_type,
)