diff --git a/main.py b/main.py index 5691f1f..db098d1 100644 --- a/main.py +++ b/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 @@ -8,6 +8,8 @@ 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" @@ -25,8 +27,13 @@ app = FastAPI(title="Pythagoras", description="A proxy service handling HTTP and app.state.auto_polling = False app.state.polling_rate = 5 -# Store for connected websocket clients +# 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): self.active_connections: List[WebSocket] = [] @@ -141,6 +148,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