Streamixph05/Rstream

1# webserver.py (FULL, COMPLETE CODE for the main.py structure)2 3import math4import traceback5import os6from fastapi import FastAPI, Request, HTTPException7from fastapi.responses import HTMLResponse, StreamingResponse8from fastapi.templating import Jinja2Templates9from pyrogram.file_id import FileId10from pyrogram import raw, Client11from pyrogram.session import Session, Auth12 13# Local imports from your project14from config import Config15from bot import multi_clients, work_loads, get_readable_file_size16from database import db17 18# FastAPI app instance, started by main.py19app = FastAPI()20templates = Jinja2Templates(directory="templates")21 22# A cache to store ByteStreamer instances to avoid re-creating them23class_cache = {}24 25@app.api_route("/", methods=["GET", "HEAD"])26async def root():27    """A simple health check route."""28    return {"status": "ok", "message": "Web server is healthy!"}29 30def mask_filename(name: str) -> str:31    """Obfuscates the filename to hide it in the URL/page."""32    if not name: return "Protected File"33    resolutions = ["216_p", "480p", "720p", "1080p", "2160p"]34    res_part = ""35    for res in resolutions:36        if res in name:37            res_part = f" {res}"38            name = name.replace(res, "")39            break40    base, ext = os.path.splitext(name)41    masked_base = ''.join(c if (i % 3 == 0 and c.isalnum()) else '*' for i, c in enumerate(base))42    return f"{masked_base}{res_part}{ext}"43 44class ByteStreamer:45    """Handles the low-level logic of fetching file parts from Telegram."""46    def __init__(self, client: Client):47        self.client = client48 49    @staticmethod50    async def get_location(file_id: FileId):51        return raw.types.InputDocumentFileLocation(52            id=file_id.media_id,53            access_hash=file_id.access_hash,54            file_reference=file_id.file_reference,55            thumb_size=file_id.thumbnail_size56        )57 58    async def yield_file(self, file_id: FileId, index: int, offset: int, first_part_cut: int, last_part_cut: int, part_count: int, chunk_size: int):59        client = self.client60        work_loads[index] += 161        62        media_session = client.media_sessions.get(file_id.dc_id)63        if media_session is None:64            if file_id.dc_id != await client.storage.dc_id():65                auth_key = await Auth(client, file_id.dc_id, await client.storage.test_mode()).create()66                media_session = Session(client, file_id.dc_id, auth_key, await client.storage.test_mode(), is_media=True)67                await media_session.start()68                exported_auth = await client.invoke(raw.functions.auth.ExportAuthorization(dc_id=file_id.dc_id))69                await media_session.invoke(raw.functions.auth.ImportAuthorization(id=exported_auth.id, bytes=exported_auth.bytes))70            else:71                media_session = client.session72            client.media_sessions[file_id.dc_id] = media_session73        74        location = await self.get_location(file_id)75        current_part = 176        try:77            while current_part <= part_count:78                r = await media_session.invoke(79                    raw.functions.upload.GetFile(location=location, offset=offset, limit=chunk_size),80                    retries=081                )82                if isinstance(r, raw.types.upload.File):83                    chunk = r.bytes84                    if not chunk: break85                    86                    if part_count == 1: yield chunk[first_part_cut:last_part_cut]87                    elif current_part == 1: yield chunk[first_part_cut:]88                    elif current_part == part_count: yield chunk[:last_part_cut]89                    else: yield chunk90                    91                    current_part += 192                    offset += chunk_size93                else:94                    break95        finally:96            work_loads[index] -= 197 98@app.get("/show/{unique_id}", response_class=HTMLResponse)99async def show_file_page(request: Request, unique_id: str):100    """The route that displays the download page to the user."""101    try:102        storage_msg_id = await db.get_link(unique_id)103        if not storage_msg_id:104            raise HTTPException(status_code=404, detail="Link expired or invalid.")105        106        # Use the main bot (client 0) to get message details107        main_bot = multi_clients.get(0)108        if not main_bot:109            raise HTTPException(status_code=503, detail="Bot is not ready yet. Please try again in a moment.")110        111        file_msg = await main_bot.get_messages(Config.STORAGE_CHANNEL, storage_msg_id)112        media = file_msg.document or file_msg.video or file_msg.audio113        if not media:114            raise HTTPException(status_code=404, detail="File not found in the message.")115        116        original_file_name = media.file_name or "file"117        safe_file_name = "".join(c for c in original_file_name if c.isalnum() or c in (' ', '.', '_', '-')).rstrip()118 119        context = {120            "request": request,121            "file_name": mask_filename(original_file_name),122            "file_size": get_readable_file_size(media.file_size),123            "is_media": (media.mime_type or "").startswith(("video/", "audio/")),124            "direct_dl_link": f"{Config.BASE_URL}/dl/{storage_msg_id}/{safe_file_name}",125            "mx_player_link": f"intent:{Config.BASE_URL}/dl/{storage_msg_id}/{safe_file_name}#Intent;action=android.intent.action.VIEW;type={media.mime_type};end",126            "vlc_player_link": f"vlc://{Config.BASE_URL}/dl/{storage_msg_id}/{safe_file_name}"127        }128        return templates.TemplateResponse("show.html", context)129 130    except HTTPException:131        raise132    except Exception as e:133        print(f"Error in /show route: {traceback.format_exc()}")134        raise HTTPException(status_code=500, detail="Internal server error.")135 136@app.get("/dl/{msg_id}/{file_name}")137async def stream_handler(request: Request, msg_id: int, file_name: str):138    """The route that handles the actual file streaming and download."""139    try:140        # Choose the client with the least workload141        index = min(work_loads, key=work_loads.get, default=0)142        client = multi_clients.get(index)143        if not client:144            raise HTTPException(status_code=503, detail="No available clients to handle the request.")145        146        tg_connect = class_cache.get(client)147        if not tg_connect:148            tg_connect = ByteStreamer(client)149            class_cache[client] = tg_connect150            151        message = await client.get_messages(Config.STORAGE_CHANNEL, msg_id)152        media = message.document or message.video or message.audio153        if not media or message.empty:154            raise FileNotFoundError155 156        file_id = FileId.decode(media.file_id)157        file_size = media.file_size158        159        range_header = request.headers.get("Range", 0)160        from_bytes, until_bytes = 0, file_size - 1161        if range_header:162            from_bytes_str, until_bytes_str = range_header.replace("bytes=", "").split("-")163            from_bytes = int(from_bytes_str)164            if until_bytes_str:165                until_bytes = int(until_bytes_str)166        167        if (until_bytes >= file_size) or (from_bytes < 0):168            raise HTTPException(status_code=416, detail="Requested range not satisfiable")169        170        req_length = until_bytes - from_bytes + 1171        chunk_size = 1024 * 1024  # 1 MB172        offset = (from_bytes // chunk_size) * chunk_size173        first_part_cut = from_bytes - offset174        last_part_cut = (until_bytes % chunk_size) + 1175        part_count = math.ceil(req_length / chunk_size)176        177        body = tg_connect.yield_file(file_id, index, offset, first_part_cut, last_part_cut, part_count, chunk_size)178        179        status_code = 206 if range_header else 200180        headers = {181            "Content-Type": media.mime_type or "application/octet-stream",182            "Accept-Ranges": "bytes",183            "Content-Disposition": f'inline; filename="{media.file_name}"',184            "Content-Length": str(req_length)185        }186        if range_header:187            headers["Content-Range"] = f"bytes {from_bytes}-{until_bytes}/{file_size}"188        189        return StreamingResponse(content=body, status_code=status_code, headers=headers)190        191    except FileNotFoundError:192        raise HTTPException(status_code=404, detail="File not found on Telegram.")193    except Exception as e:194        print(f"Error in /dl route: {traceback.format_exc()}")195        raise HTTPException(status_code=500, detail="Internal streaming error.")196