import json from fastapi import FastAPI, HTTPException, Header, BackgroundTasks, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel import os import asyncio import aiohttp import aiofiles from pyrogram import Client import uuid from dotenv import load_dotenv # Load .env if it exists load_dotenv() app = FastAPI() @app.get("/") async def root(): return {"status": "Sonic Vault Active", "message": "Heavy Lifter Operational"} # Environment Variables API_ID = int(os.getenv("TELEGRAM_API_ID") or os.getenv("API_ID") or "0") API_HASH = os.getenv("TELEGRAM_API_HASH") or os.getenv("API_HASH") or "" BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "") CHAT_ID = int(os.getenv("TELEGRAM_CHAT_ID", "0")) HF_SECRET = os.getenv("HF_SECRET") or os.getenv("SECRET_TOKEN") or "" HF_TOKEN = os.getenv("HF_TOKEN", "") EXTERNAL_API_URL = os.getenv("EXTERNAL_API_URL", "https://api01.shrutibots.site") YT_API_KEY = os.getenv("YT_API_KEY", "") # Initialize Pyrogram Client (MTProto - Bot Only) bot = Client( "sonic_vault_bot", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN, in_memory=True ) class VaultRequest(BaseModel): videoId: str type: str # 'audio' or 'video' title: str downloadUrl: str = None callbackUrl: str = None @app.on_event("startup") async def startup(): await bot.start() from functools import wraps def require_auth(f): @wraps(f) async def decorated(*args, **kwargs): request = kwargs.get('request') if not request: for arg in args: if isinstance(arg, Request): request = arg break if not request: raise HTTPException(status_code=500, detail="Internal Server Error: Request object not found") auth = request.headers.get("Authorization") secret = request.headers.get("X-Secret-Token", "").strip().strip('"').strip("'") # Extract bearer and clean it bearer = None if auth and auth.startswith("Bearer "): bearer = auth.split(" ")[1].strip().strip('"').strip("'") # Clean environment variables clean_hf_token = HF_TOKEN.strip().strip('"').strip("'") clean_hf_secret = HF_SECRET.strip().strip('"').strip("'") # Dual-header authentication: # 1. Bearer token must match either clean_hf_token or clean_hf_secret # 2. X-Secret-Token must match clean_hf_secret is_bearer_valid = False if bearer: if clean_hf_token and bearer == clean_hf_token: is_bearer_valid = True elif clean_hf_secret and bearer == clean_hf_secret: is_bearer_valid = True elif not clean_hf_token and not clean_hf_secret: is_bearer_valid = True is_secret_valid = (secret == clean_hf_secret) if clean_hf_secret else True if not is_bearer_valid or not is_secret_valid: print(f"[AUTH] Denied. Bearer: {'Match' if is_bearer_valid else 'Fail'} (len:{len(bearer or '')}), Secret: {'Match' if is_secret_valid else 'Fail'} (len:{len(secret)})") # If Bearer failed, log if it matched the secret instead (common misconfiguration) if not is_bearer_valid and bearer == clean_hf_secret: print("[AUTH] Warning: Bearer token matched HF_SECRET instead of HF_TOKEN.") raise HTTPException(status_code=401, detail="Unauthorized") return await f(*args, **kwargs) return decorated async def fetch_download_url(video_id: str, media_type: str): try: async with aiohttp.ClientSession() as session: url = f"{EXTERNAL_API_URL}/download" params = {"url": video_id, "type": media_type, "api_key": YT_API_KEY} headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {} async with session.get(url, params=params, headers=headers) as resp: if resp.status == 200: return f"{EXTERNAL_API_URL}/download?url={video_id}&type={media_type}&api_key={YT_API_KEY}" except Exception as e: print(f"[Vault] Error fetching download URL: {str(e)}") return None async def process_vault(request: VaultRequest): base_file = f"/tmp/{request.videoId}_{uuid.uuid4().hex[:8]}" file_path = f"{base_file}.{'mp4' if request.type == 'video' else 'mp3'}" final_path = None try: # 1. Resolve Download URL via External API ONLY download_url = request.downloadUrl if not download_url: print(f"[Vault] Resolving {request.videoId} via {EXTERNAL_API_URL}...") download_url = await fetch_download_url(request.videoId, request.type) if not download_url: raise Exception("Could not resolve download URL from external API") print(f"[Vault] Downloading from {download_url}...") async with aiohttp.ClientSession() as session: async with session.get(download_url) as resp: if resp.status == 200: async with aiofiles.open(file_path, mode='wb') as f: async for chunk in resp.content.iter_chunked(1024 * 1024): await f.write(chunk) final_path = file_path print(f"[Vault] Downloaded: {os.path.getsize(final_path)} bytes") else: raise Exception(f"Download failed with status {resp.status}") # 2. Upload to Telegram (MTProto) print(f"[Vault] Uploading to chat {CHAT_ID}...") caption = f"{request.title}\nID: {request.videoId}" if request.type == 'video': msg = await bot.send_video( chat_id=CHAT_ID, video=final_path, caption=caption, supports_streaming=True, file_name=f"{request.title}.mp4" ) file_id = msg.video.file_id else: msg = await bot.send_audio( chat_id=CHAT_ID, audio=final_path, caption=caption, title=request.title, performer="Sonic Vault", file_name=f"{request.title}.mp3" ) file_id = msg.audio.file_id print(f"[Vault] Upload success! Message ID: {msg.id}") # 3. Callback if request.callbackUrl: payload = { "status": "success", "telegram_file_id": file_id, "message_id": msg.id, "videoId": request.videoId, "type": request.type, "title": request.title } async with aiohttp.ClientSession() as session: await session.post( request.callbackUrl, json=payload, headers={ "Authorization": f"Bearer {HF_SECRET}", "X-Secret-Token": HF_SECRET }, timeout=15 ) except Exception as e: print(f"[Vault] Process Error: {str(e)}") if request.callbackUrl: try: async with aiohttp.ClientSession() as session: await session.post( request.callbackUrl, json={"status": "error", "error": str(e), "videoId": request.videoId}, headers={ "Authorization": f"Bearer {HF_SECRET}", "X-Secret-Token": HF_SECRET }, timeout=10 ) except: pass finally: if final_path and os.path.exists(final_path): try: os.remove(final_path) except: pass @app.post("/vault") @require_auth async def vault_media(request: Request, vault_req: VaultRequest, background_tasks: BackgroundTasks): background_tasks.add_task(process_vault, vault_req) return {"status": "accepted", "message": "Vaulting started"} @app.get("/proxy-stream/{file_id}") @require_auth async def proxy_stream(request: Request, file_id: str, type: str = "audio"): async def stream_generator(): try: # Pyrogram's stream_media uses MTProto (Bot Only) async for chunk in bot.stream_media(file_id): yield chunk except Exception as e: print(f"[Stream] MTProto Error: {str(e)}") media_type = "video/mp4" if type == "video" else "audio/mpeg" return StreamingResponse(stream_generator(), media_type=media_type) @app.get("/health") async def health(): return {"status": "ok", "mtproto": "active" if bot.is_connected else "inactive"}