Update main.py

main.py

@@ -25,17 +25,12 @@ app = FastAPI(title="Static Site Server")
 # Add security middlewares
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Configure as needed
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-app.add_middleware(
-    TrustedHostMiddleware,
-    allowed_hosts=["*"]  # Configure as needed
-)
-
 # Constants
 MAX_UPLOAD_SIZE = 100 * 1024 * 1024  # 100MB
 ALLOWED_EXTENSIONS = {'.html', '.css', '.js', '.jpg', '.jpeg', '.png', '.gif', '.svg', '.ico', '.woff', '.woff2', '.ttf', '.eot'}
@@ -162,6 +157,13 @@ class SiteManager:
             return True
         return False
 
+    def get_site_path(self, site_id: str) -> Optional[Path]:
+        """Get the path for a site if it exists"""
+        site_path = self.sites_dir / site_id
+        if site_path.is_dir() and (site_path / 'index.html').exists():
+            return site_path
+        return None
+
 # Initialize site manager
 site_manager = SiteManager()
 
@@ -191,14 +193,31 @@ async def health_check():
     """Health check endpoint"""
     return {"status": "healthy", "sites_count": len(site_manager.active_sites)}
 
-# Mount static file handlers for each site
-@app.on_event("startup")
-async def startup_event():
-    """Configure static file handlers for existing sites"""
-    logger.info("Starting up server...")
-    for site_id in site_manager.active_sites:
-        site_path = site_manager.sites_dir / site_id
-        app.mount(f"/{site_id}", StaticFiles(directory=str(site_path), html=True), name=site_id)
+@app.get("/{site_id}/{file_path:path}")
+async def serve_site(site_id: str, file_path: str = ""):
+    """Serve files from the site directory"""
+    site_path = site_manager.get_site_path(site_id)
+    if not site_path:
+        raise HTTPException(status_code=404, detail="Site not found")
+    
+    # Default to index.html if no file specified
+    if not file_path:
+        file_path = "index.html"
+    
+    file_full_path = site_path / file_path
+    
+    # Prevent directory traversal
+    try:
+        file_full_path = file_full_path.resolve()
+        if not str(file_full_path).startswith(str(site_path)):
+            raise HTTPException(status_code=403, detail="Access denied")
+    except (RuntimeError, ValueError):
+        raise HTTPException(status_code=400, detail="Invalid path")
+    
+    if not file_full_path.is_file():
+        raise HTTPException(status_code=404, detail="File not found")
+    
+    return FileResponse(file_full_path)
 
 if __name__ == "__main__":
     import uvicorn