feat: streams download response to client

Author: anonymous-org-za

functions/tinfoilFunctions.py

@@ -1,10 +1,12 @@
 import asyncio
-from functions.torboxFunctions import getDownloads
+from functions.torboxFunctions import getDownloads, getDownloadLink
 import logging
 from library.tinfoil import errorMessage
-from fastapi.responses import JSONResponse
+from fastapi import BackgroundTasks
+from fastapi.responses import JSONResponse, StreamingResponse
 import fnmatch
 import human_readable
+import httpx
 
 ACCEPTABLE_SWITCH_FILES = [".nsp", ".nsz", ".xci", ".xcz"]
 
@@ -23,7 +25,6 @@ async def generateIndex(base_url: str):
     success_message = "Welcome to your self-hosted TorBox Tinfoil Server! You are now able to directly download your files from TorBox to your switch.\n\n"
     files = []
 
-    
     try:
         # runs all requests in parallel for faster responses
         torrents, usenet_downloads, web_downloads = await asyncio.gather(getDownloads("torrents"), getDownloads("usenet"), getDownloads("webdl"))
@@ -63,6 +64,36 @@ async def generateIndex(base_url: str):
             content=errorMessage(f"There was an error generating the index. Error: {str(e)}", error_code="UNKOWN_ERROR")
         )
 
+async def serveFile(background_task: BackgroundTasks, download_type: str, download_id: int, file_id: int = 0):
+    """
+    Retrieves the TorBox download link and starts proxying the download through the server. This is necessary as generating a bunch of links through the index generation process can take some time, and is wasteful.
+
+    Requires:
+    - download_type: the download type of the file. Must be either 'torrents', 'usenet' or 'webdl'.
+    - download_id: an integer which represents the id of the download in the TorBox database.
+    - file_id: an integer which represents the id of the file which is inside of the download.
+    
+    Returns:
+    - Streaming Response: containing the download of the file to be served on the fly.
+    """
+
+    download_link = await getDownloadLink(download_type=download_type, download_id=download_id, file_id=file_id)
+
+    if not download_link:
+        return JSONResponse(
+            status_code=500,
+            content=errorMessage("There was an error retrieving the download link from TorBox. Please try again.", error_code="DATABASE_ERROR")
+        )
+
+    # now stream link and stream out
+    client = httpx.AsyncClient()
+    request = client.build_request(method="GET", url=download_link, timeout=90)
+    response = await client.send(request, stream=True)
+
+    cleanup = background_task.add_task(response.aclose)
+
+    return StreamingResponse(content=response.aiter_bytes(), background=cleanup)
+    
 
 
 

functions/torboxFunctions.py

@@ -36,7 +36,6 @@ async def getDownloads(type: str):
             # only get downloads that are completely ready for download
             if not download.get("download_present", False):
                 continue
-            type = "download"
             id = download.get("id", None)
             for file in download.get("files", []):
                 if not file.get("s3_path", None):
@@ -57,4 +56,39 @@ async def getDownloads(type: str):
         traceback.print_exc()
         logging.error(f"There was an error getting {type} downloads from TorBox. Error: {str(e)}")
         return []
+    
+async def getDownloadLink(download_type: str, download_id: int, file_id: int = 0):
+    if download_type not in ["torrents", "usenet", "webdl"]:
+        logging.error("Please provide a type of either 'torrents', 'usenet' or 'webdl'.")
+        return None
+    
+    try:
+        async with httpx.AsyncClient() as client:
+            if download_type == "torrents":
+                id_type = "torrent_id"
+            elif download_type == "usenet":
+                id_type = "usenet_id"
+            elif download_type == "webdl":
+                id_type = "web_id"
+            params = {
+                "token": TORBOX_API_KEY,
+                id_type: download_id,
+                file_id: file_id
+            }
+            response = await client.get(
+                url=f"{TORBOX_API_URL}/{download_type}/requestdl",
+                params=params,
+                headers={
+                    "User-Agent": "TorBox SelfHosted Tinfoil Server/1.0.0"
+                }
+            )
+        if response.status_code != httpx.codes.OK:
+            logging.error(f"Unable to retrieve TorBox {download_type} download link. Response Code: {response.status_code}. Response: {response.json()}")
+            return None
+        
+        json = response.json()
+        link = json.get("data", None)
+        return link
+    except Exception as e:
+        logging.error(f"There was an error retrieving the download url from TorBox. Error: {str(e)}")
 

main.py

@@ -1,13 +1,13 @@
 import uvicorn
-from fastapi import FastAPI, HTTPException, status, Depends, Header, Request
+from fastapi import FastAPI, HTTPException, status, Depends, Header, Request, BackgroundTasks
 from fastapi.responses import JSONResponse
 from typing_extensions import Annotated, Union
 from library.tinfoil import errorMessage
 from library.server import PORT
 from functions.authFunctions import checkCorrectCredentials
 from functions.serverFunctions import checkAllowed
 import logging
-from functions.tinfoilFunctions import generateIndex
+from functions.tinfoilFunctions import generateIndex, serveFile
 
 app = FastAPI()
 logging.basicConfig(level=logging.INFO)
@@ -36,7 +36,15 @@ async def get_user_files(
             status_code=401
         )
     return await generateIndex(base_url=request.base_url)
-    
+
+@app.get("/{download_type}/{download_id}/{file_id}")
+async def get_file(
+    background_task: BackgroundTasks, # background_task is used to clean up the httpx response afterwards to prevent leakage
+    download_type: str,
+    download_id: int,
+    file_id: int = 0
+):
+    return await serveFile(background_task=background_task, download_type=download_type, download_id=download_id, file_id=file_id)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=PORT)