transfer-test / app.py
r3hab's picture
Update app.py
c918e99 verified
raw
history blame
3.12 kB
# app.py
import os
import time
import uuid
import asyncio
import libtorrent as lt
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import magic
app = FastAPI()
active_sessions = {}
def get_largest_file(torrent_info):
files = torrent_info.files()
file_sizes = [(i, files.file_size(i)) for i in range(files.num_files())]
if not file_sizes:
return None
return max(file_sizes, key=lambda x: x[1])
async def add_torrent(magnet_link: str, save_path: str):
ses = lt.session()
params = {
'save_path': save_path,
'storage_mode': lt.storage_mode_t(2),
}
handle = lt.add_magnet_uri(ses, magnet_link, params)
# Wait for metadata
while not handle.has_metadata():
await asyncio.sleep(1)
torrent_info = handle.get_torrent_info()
file_info = get_largest_file(torrent_info)
if not file_info:
raise Exception("No files found in torrent")
file_index, _ = file_info
file_path = os.path.join(save_path, torrent_info.files().file_path(file_index))
# Prioritize largest file
for i in range(torrent_info.num_files()):
handle.file_priority(i, 7 if i == file_index else 0)
handle.set_sequential_download(True)
ses.resume()
return {
"session": ses,
"handle": handle,
"file_path": file_path,
"downloading": True
}
@app.post("/start_stream")
async def start_stream(magnet_link: str):
torrent_id = str(uuid.uuid4())
save_path = f"./downloads/{torrent_id}"
os.makedirs(save_path, exist_ok=True)
try:
session_info = await add_torrent(magnet_link, save_path)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
active_sessions[torrent_id] = session_info
return {"stream_url": f"/stream/{torrent_id}"}
@app.get("/stream/{torrent_id}")
async def stream_torrent(torrent_id: str):
if torrent_id not in active_sessions:
raise HTTPException(status_code=404, detail="Torrent session not found")
session_info = active_sessions[torrent_id]
file_path = session_info["file_path"]
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="File not found")
mime = magic.Magic(mime=True)
mime_type = mime.from_file(file_path)
def file_generator():
last_position = 0
while True:
if not os.path.exists(file_path):
time.sleep(1)
continue
with open(file_path, "rb") as f:
f.seek(last_position)
data = f.read(1024 * 1024) # 1MB chunks
if not data:
if session_info["handle"].is_seed():
break
time.sleep(1)
continue
last_position = f.tell()
yield data
return StreamingResponse(file_generator(), media_type=mime_type)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)