Update main.py

main.py

@@ -165,8 +165,8 @@ async def hentai_random(apikey: Optional[str] = Query(None), limit: str = Query(
     key = x_api_key or apikey
     return await ht_random_app.get_random(key, API_KEY, limit)
 
-@app.get("/api/v1/tiktok/download")
-async def tiktok_download(
+@app.get("/api/v1/tiktok/dl")
+async def tiktok_dl(
     url: str = Query(...), 
     hd: int = Query(1), 
     apikey: Optional[str] = Query(None), 
@@ -181,92 +181,86 @@ async def tiktok_download(
         raise HTTPException(status_code=400, detail=res["error"])
     return res
 
-@app.get("/api/v1/tiktok/dl-direct")
-async def tiktok_direct(
-    url: str = Query(None),
-    token: str = Query(None),
-    type: str = Query("mp4"),
+@app.get("/api/v1/tiktok/cdn")
+async def tiktok_cdn(
+    token: str = Query(None, description="Token generated from /dl endpoint"),
+    url: str = Query(None, description="Direct URL to proxy (Requires API Key)"),
+    type: str = Query("mp4", description="File extension type"),
     apikey: Optional[str] = Query(None),
     x_api_key: Optional[str] = Header(None, alias="X-API-Key")
 ):
     if url:
         key = x_api_key or apikey
         if not API_KEY or key != API_KEY:
-            raise HTTPException(status_code=403, detail="API key required to generate token")
-        
-        res = await tik_direct.generate_token(url, type)
-        if not res["ok"]:
-            raise HTTPException(status_code=400, detail=res["error"])
-        return res
-
-    if token:
+            raise HTTPException(status_code=403, detail="Proxying via URL requires a valid API Key")
+        data = {"url": url, "type": type}
+    
+    elif token:
         data = tik_direct.consume_token(token)
         if not data:
-            raise HTTPException(status_code=404, detail="Token invalid, expired hoặc đã dùng quá 3 lần")
-
-        if data["type"] in ["mp3_convert", "wav_convert"]:
-            is_wav = data["type"] == "wav_convert"
-            ext = "wav" if is_wav else "mp3"
-            acodec = "pcm_s16le" if is_wav else "libmp3lame"
-            fmt = "wav" if is_wav else "mp3"
-            media_type = "audio/wav" if is_wav else "audio/mpeg"
-            filename = f"Celeskry_Audio_{int(time.time())}.{ext}"
-            
-            try:
-                process = (
-                    ffmpeg
-                    .input(data["url"])
-                    .output('pipe:', format=fmt, acodec=acodec, ar='44100')
-                    .run_async(pipe_stdout=True, pipe_stderr=True)
-                )
-
-                async def stream_ffmpeg():
-                    try:
-                        while True:
-                            chunk = await asyncio.to_thread(process.stdout.read, 4096)
-                            if not chunk:
-                                break
-                            yield chunk
-                    finally:
-                        try:
-                            process.kill()
-                        except:
-                            pass
-
-                return StreamingResponse(
-                    stream_ffmpeg(),
-                    media_type=media_type,
-                    headers={"Content-Disposition": f"attachment; filename={filename}"}
-                )
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Lỗi convert {ext}: {str(e)}")
-
-        async def stream_file():
-            async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as client:
-                async with client.stream("GET", data["url"]) as r:
-                    if r.status_code != 200:
-                        yield b"Error: Could not fetch file from provider"
-                        return
-                    async for chunk in r.aiter_bytes(chunk_size=1024*1024):
+            raise HTTPException(status_code=404, detail="Token invalid, expired or limit reached")
+    
+    else:
+        raise HTTPException(status_code=400, detail="Either 'token' or 'url' must be provided")
+
+    target_url = data["url"]
+    media_type_req = data["type"]
+
+    if media_type_req in ["mp3_convert", "wav_convert"]:
+        is_wav = media_type_req == "wav_convert"
+        ext = "wav" if is_wav else "mp3"
+        acodec = "pcm_s16le" if is_wav else "libmp3lame"
+        fmt = "wav" if is_wav else "mp3"
+        media_header = "audio/wav" if is_wav else "audio/mpeg"
+        
+        try:
+            process = (
+                ffmpeg
+                .input(target_url)
+                .output('pipe:', format=fmt, acodec=acodec, ar='44100')
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+
+            async def stream_ffmpeg():
+                try:
+                    while True:
+                        chunk = await asyncio.to_thread(process.stdout.read, 8192)
+                        if not chunk: break
                         yield chunk
-
-        ext = data["type"] 
-        if ext == "mp4":
-            media_type = "video/mp4"
-        elif ext == "mp3":
-            media_type = "audio/mpeg"
-        elif ext in ["jpg", "jpeg", "png"]:
-            media_type = f"image/{ext}"
-        else:
-            media_type = "application/octet-stream"
-
-        filename = f"Celeste_{int(time.time())}.{ext}"
-
-        return StreamingResponse(
-            stream_file(),
-            media_type=media_type,
-            headers={"Content-Disposition": f"attachment; filename={filename}"}
-        )
+                finally:
+                    try: process.kill()
+                    except: pass
+
+            return StreamingResponse(
+                stream_ffmpeg(),
+                media_type=media_header,
+                headers={"Content-Disposition": f"attachment; filename=Celeskry_{int(time.time())}.{ext}"}
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Conversion Error: {str(e)}")
+
+    async def stream_file():
+        async with httpx.AsyncClient(follow_redirects=True, timeout=60.0) as client:
+            async with client.stream("GET", target_url) as r:
+                if r.status_code != 200:
+                    yield b"Error: Resource not found"
+                    return
+                async for chunk in r.aiter_bytes(chunk_size=1024*1024):
+                    yield chunk
+                    
+    content_types = {
+        "mp4": "video/mp4",
+        "mp3": "audio/mpeg",
+        "jpg": "image/jpeg",
+        "png": "image/png"
+    }
+    final_type = content_types.get(media_type_req, "application/octet-stream")
+    
+    return StreamingResponse(
+        stream_file(),
+        media_type=final_type,
+        headers={"Content-Disposition": f"attachment; filename=Celeste_{int(time.time())}.{media_type_req}"}
+    )
 
     raise HTTPException(status_code=400, detail="Provide either 'url' or 'token'")