Add media streaming capabilities

main.py

@@ -1,5 +1,5 @@
-from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, BackgroundTasks
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, BackgroundTasks, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 import logging
 import uvicorn
 from typing import Dict, List, Any
@@ -7,6 +7,9 @@ from dataclasses import dataclass
 import json
 import httpx
 import asyncio
+import struct
+import os
+from pathlib import Path
 
 # Some useful variables
 PEEHAITCHPEA_ENDPOINT = "http://localhost:9000/api.php?cmd=getselectedmessage"
@@ -24,6 +27,12 @@ app = FastAPI(title="Pythagoras", description="A proxy service handling HTTP and
 app.state.auto_polling = False
 app.state.polling_rate = 5
 
+# Define the media directory
+MEDIA_DIR = Path("./media")
+
+# Create the media directory if it doesn't exist
+os.makedirs(MEDIA_DIR, exist_ok=True)
+
 # Store for connected websocket clients
 class ConnectionManager:
     def __init__(self):
@@ -116,6 +125,28 @@ async def websocket_endpoint(websocket: WebSocket):
         logger.error(f"WebSocket error: {str(e)}")
         manager.disconnect(websocket)
 
+@app.get("/media/{file_path:path}")
+async def get_media(file_path: str):
+    """
+    Serve media files from the media directory.
+    """
+    full_path = MEDIA_DIR / file_path
+    absolute_path = full_path.resolve()
+    
+    # Security check: Make sure the path is within the media directory
+    if not str(absolute_path).startswith(str(MEDIA_DIR.resolve())):
+        logger.warning(f"Attempted directory traversal: {file_path}")
+        raise HTTPException(status_code=403, detail="Access denied")
+    
+    # Check if the file exists
+    if not absolute_path.is_file():
+        logger.warning(f"File not found: {absolute_path}")
+        raise HTTPException(status_code=404, detail="File not found")
+    
+    logger.info(f"Serving media file: {file_path}")
+    
+    return FileResponse(path=absolute_path, filename=absolute_path.name)
+
 
 
 # Functions