Spaces:
Running
Running
| import os | |
| import json | |
| import math | |
| import asyncio | |
| import logging | |
| import mimetypes | |
| from pathlib import Path | |
| from aiohttp import web | |
| from dotenv import load_dotenv | |
| logger = logging.getLogger(__name__) | |
| load_dotenv() | |
| STORAGE_FILE = "file_store.json" | |
| def load_store(): | |
| if Path(STORAGE_FILE).exists(): | |
| with open(STORAGE_FILE, "r") as f: | |
| return json.load(f) | |
| return {} | |
| def save_store(store): | |
| with open(STORAGE_FILE, "w") as f: | |
| json.dump(store, f, indent=2) | |
| async def serve_homepage(request): | |
| store = load_store() | |
| total_files = len(store) | |
| total_downloads = sum(v.get("downloads", 0) for v in store.values()) | |
| html = f"""<!DOCTYPE html> | |
| <html><head><meta charset="UTF-8"><title>TG Stream Server</title> | |
| <style>body{{font-family:monospace;background:#0a0a0a;color:#00ff88;display:flex;align-items:center;justify-content:center;min-height:100vh;margin:0}}</style> | |
| </head><body> | |
| <h1>⚡ TG Streaming Server (Pyrogram)</h1> | |
| <p>Files Hosted: {total_files} | Total Downloads: {total_downloads}</p> | |
| </body></html>""" | |
| return web.Response(text=html, content_type="text/html") | |
| MESSAGE_CACHE = {} | |
| CHUNK_SIZE = 1048576 # Pyrogram's fixed 1 MiB chunk size | |
| async def _pipe_stream(request, response, bot_app, target_media, | |
| first_chunk_idx, num_chunks, slice_start, content_length): | |
| """ | |
| Streams Pyrogram chunks to the aiohttp response. | |
| Isolated into its own function so the finally block can guarantee | |
| generator.aclose() is called regardless of how the loop exits | |
| (client abort, ConnectionResetError, CancelledError, or normal EOF). | |
| """ | |
| generator = bot_app.stream_media( | |
| target_media, | |
| offset=first_chunk_idx, | |
| limit=num_chunks, | |
| ) | |
| bytes_remaining = content_length | |
| is_first_chunk = True | |
| try: | |
| async for raw_chunk in generator: | |
| # Detect client disconnect before attempting a write. | |
| if request.transport is None or request.transport.is_closing(): | |
| logger.debug("Client disconnected mid-stream, cleaning up generator.") | |
| break | |
| # Slice the first chunk to honour the byte-level range start. | |
| if is_first_chunk: | |
| raw_chunk = raw_chunk[slice_start:] | |
| is_first_chunk = False | |
| # Trim the last chunk so we never overshoot the requested range end. | |
| if len(raw_chunk) > bytes_remaining: | |
| raw_chunk = raw_chunk[:bytes_remaining] | |
| if not raw_chunk: | |
| break | |
| try: | |
| await response.write(raw_chunk) | |
| except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError): | |
| # Client aborted (normal for movi-player's probing requests). | |
| logger.debug("Write interrupted by client, breaking stream loop.") | |
| break | |
| bytes_remaining -= len(raw_chunk) | |
| if bytes_remaining <= 0: | |
| break | |
| except asyncio.CancelledError: | |
| logger.debug("Stream handler task cancelled (server shutdown or client abort).") | |
| except Exception as e: | |
| logger.error(f"Unexpected error in _pipe_stream: {e}", exc_info=True) | |
| finally: | |
| # THE CRITICAL FIX: explicitly close the async generator. | |
| # This throws GeneratorExit into stream_media, causing Pyrogram to | |
| # cancel the in-flight GetFile RPC and release the MTProto session slot. | |
| # Without this, every aborted request leaks a slot until the pool deadlocks. | |
| try: | |
| await generator.aclose() | |
| except Exception as e: | |
| logger.warning(f"Error during generator.aclose(): {e}") | |
| async def serve_file(request, inline=False): | |
| token = request.match_info.get('token') | |
| print(f"[serve_file] HIT — token={token}, range={request.headers.get('Range')}") | |
| store = load_store() | |
| if token not in store: | |
| return web.Response(status=404, text="File not found") | |
| info = store[token] | |
| filename = info["filename"] | |
| file_size = info["file_size"] | |
| message_id = info["message_id"] | |
| chat_id = info["chat_id"] | |
| bot_app = request.app['bot'] | |
| cache_key = f"{chat_id}_{message_id}" | |
| msg = MESSAGE_CACHE.get(cache_key) | |
| if not msg: | |
| msg = await bot_app.get_messages(chat_id, message_id) | |
| if msg: | |
| MESSAGE_CACHE[cache_key] = msg | |
| if not msg: | |
| return web.Response(status=404, text="Message not found on Telegram") | |
| file_obj = msg.document or msg.video or msg.audio or msg.photo | |
| if not file_obj: | |
| return web.Response(status=404, text="File not found on Telegram") | |
| target_media = file_obj | |
| mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" | |
| disposition = "inline" if inline else "attachment" | |
| headers = { | |
| "Content-Disposition": f'{disposition}; filename="{filename}"', | |
| "Content-Type": mime_type, | |
| "Accept-Ranges": "bytes", | |
| "Cross-Origin-Resource-Policy": "cross-origin", | |
| "Access-Control-Allow-Origin": "*" | |
| } | |
| range_header = request.headers.get("Range") | |
| start_byte = 0 | |
| end_byte = file_size - 1 | |
| if range_header: | |
| try: | |
| s, e = range_header.replace("bytes=", "").split("-") | |
| if s == "": | |
| # bytes=-N (tail) | |
| end_byte = file_size - 1 | |
| start_byte = max(0, file_size - int(e)) | |
| else: | |
| start_byte = int(s) | |
| end_byte = int(e) if e else file_size - 1 | |
| headers["Content-Range"] = f"bytes {start_byte}-{end_byte}/{file_size}" | |
| response = web.StreamResponse(status=206, headers=headers) | |
| except Exception: | |
| response = web.StreamResponse(status=200, headers=headers) | |
| else: | |
| response = web.StreamResponse(status=200, headers=headers) | |
| content_length = end_byte - start_byte + 1 | |
| response.content_length = content_length | |
| first_chunk_idx = start_byte // CHUNK_SIZE | |
| slice_start = start_byte % CHUNK_SIZE | |
| num_chunks = math.ceil((slice_start + content_length) / CHUNK_SIZE) | |
| # Acquire semaphore BEFORE preparing the response. | |
| # If we prepare first and then block on the semaphore, the browser sees | |
| # an accepted connection with no data flowing — it hangs forever. | |
| async with request.app['pyrogram_sem']: | |
| await response.prepare(request) | |
| await _pipe_stream( | |
| request, response, bot_app, target_media, | |
| first_chunk_idx, num_chunks, slice_start, content_length | |
| ) | |
| try: | |
| store[token]["downloads"] = store[token].get("downloads", 0) + 1 | |
| save_store(store) | |
| except Exception as e: | |
| logger.warning(f"Could not update download counter: {e}") | |
| return response | |
| async def handle_download(request): | |
| return await serve_file(request, inline=False) | |
| def create_app(bot_app): | |
| app = web.Application() | |
| app['bot'] = bot_app | |
| app.router.add_get('/', serve_homepage) | |
| app.router.add_get('/download/{token}', handle_download) | |
| return app | |