Update main.py

main.py

@@ -1,28 +1,19 @@
-import os
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pathlib import Path
 import shutil
 import zipfile
 import logging
-import tempfile
-import magic
-from pathlib import Path
-from typing import Set, Optional
-from fastapi import FastAPI, File, UploadFile, HTTPException, Request
-from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.middleware.trustedhost import TrustedHostMiddleware
+from typing import Set
 
 # Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize FastAPI app
-app = FastAPI(title="Static Site Server")
+app = FastAPI()
 
-# Add security middlewares
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -31,105 +22,40 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Constants
-MAX_UPLOAD_SIZE = 100 * 1024 * 1024  # 100MB
-ALLOWED_EXTENSIONS = {'.html', '.css', '.js', '.jpg', '.jpeg', '.png', '.gif', '.svg', '.ico', '.woff', '.woff2', '.ttf', '.eot'}
-
 class SiteManager:
     def __init__(self):
         self.sites_dir = Path("/app/sites")
-        self.temp_dir = Path("/app/temp")
-        self.active_sites: Set[str] = set()
-        
-        # Ensure directories exist
         self.sites_dir.mkdir(parents=True, exist_ok=True)
-        self.temp_dir.mkdir(parents=True, exist_ok=True)
-        
-        # Load existing sites
+        self.active_sites: Set[str] = set()
         self._load_existing_sites()
-        
+
     def _load_existing_sites(self):
         """Load existing sites from disk"""
-        logger.info("Loading existing sites...")
         for site_dir in self.sites_dir.iterdir():
             if site_dir.is_dir() and (site_dir / 'index.html').exists():
                 self.active_sites.add(site_dir.name)
+                # Mount the site directory
+                app.mount(f"/{site_dir.name}", StaticFiles(directory=str(site_dir), html=True), name=site_dir.name)
                 logger.info(f"Loaded site: {site_dir.name}")
 
-    def _validate_file_types(self, zip_path: Path) -> bool:
-        """Validate file types in ZIP archive"""
-        mime = magic.Magic(mime=True)
-        with zipfile.ZipFile(zip_path) as zip_ref:
-            for file_info in zip_ref.filelist:
-                if file_info.filename.endswith('/'):  # Skip directories
-                    continue
-                    
-                suffix = Path(file_info.filename).suffix.lower()
-                if suffix not in ALLOWED_EXTENSIONS:
-                    return False
-                
-                # Extract file to check MIME type
-                with tempfile.NamedTemporaryFile() as tmp:
-                    with zip_ref.open(file_info) as source:
-                        shutil.copyfileobj(source, tmp)
-                    tmp.flush()
-                    mime_type = mime.from_file(tmp.name)
-                    if mime_type.startswith('application/x-'):
-                        return False
-        return True
-
     async def deploy_site(self, unique_id: str, zip_file: UploadFile) -> dict:
         """Deploy a new site from a ZIP file"""
-        if await zip_file.read(1) == b'':
-            raise HTTPException(status_code=400, detail="Empty file")
-        await zip_file.seek(0)
-
-        # Create temporary file
-        temp_file = self.temp_dir / f"{unique_id}.zip"
+        site_path = self.sites_dir / unique_id
+        
         try:
-            # Save uploaded file
-            content = await zip_file.read()
-            if len(content) > MAX_UPLOAD_SIZE:
-                raise HTTPException(status_code=400, detail="File too large")
-            
-            temp_file.write_bytes(content)
-            
-            # Validate ZIP file
-            if not zipfile.is_zipfile(temp_file):
-                raise HTTPException(status_code=400, detail="Invalid ZIP file")
-                
-            # Validate file types
-            if not self._validate_file_types(temp_file):
-                raise HTTPException(status_code=400, detail="Invalid file types in ZIP")
-            
-            # Process the ZIP file
-            site_path = self.sites_dir / unique_id
-            with zipfile.ZipFile(temp_file) as zip_ref:
-                # Verify index.html exists
-                if not any(name.endswith('/index.html') or name == 'index.html' 
-                          for name in zip_ref.namelist()):
-                    raise HTTPException(
-                        status_code=400,
-                        detail="ZIP file must contain index.html in root directory"
-                    )
-                
+            # Save and extract ZIP
+            with zipfile.ZipFile(zip_file.file) as zip_ref:
                 # Clear existing site if present
                 if site_path.exists():
                     shutil.rmtree(site_path)
                 
-                # Extract files
-                zip_ref.extractall(self.temp_dir / unique_id)
-                
-                # Move to final location
-                extraction_path = self.temp_dir / unique_id
-                root_dir = next(
-                    (p for p in extraction_path.iterdir() if p.is_dir() 
-                     and (p / 'index.html').exists()),
-                    extraction_path
-                )
-                shutil.move(str(root_dir), str(site_path))
-                
+                # Extract ZIP contents
+                zip_ref.extractall(site_path)
+            
+            # Mount the new site
+            app.mount(f"/{unique_id}", StaticFiles(directory=str(site_path), html=True), name=unique_id)
             self.active_sites.add(unique_id)
+            
             return {
                 "status": "success",
                 "message": f"Site deployed at /{unique_id}",
@@ -138,14 +64,9 @@ class SiteManager:
             
         except Exception as e:
             logger.error(f"Error deploying site {unique_id}: {str(e)}")
+            if site_path.exists():
+                shutil.rmtree(site_path)
             raise HTTPException(status_code=500, detail=str(e))
-        finally:
-            # Cleanup
-            if temp_file.exists():
-                temp_file.unlink()
-            cleanup_path = self.temp_dir / unique_id
-            if cleanup_path.exists():
-                shutil.rmtree(cleanup_path)
 
     def remove_site(self, unique_id: str) -> bool:
         """Remove a deployed site"""
@@ -157,13 +78,6 @@ 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()
 
@@ -173,8 +87,7 @@ async def deploy_site(unique_id: str, file: UploadFile = File(...)):
     if not file.filename.endswith('.zip'):
         raise HTTPException(status_code=400, detail="File must be a ZIP archive")
     
-    result = await site_manager.deploy_site(unique_id, file)
-    return JSONResponse(content=result)
+    return await site_manager.deploy_site(unique_id, file)
 
 @app.delete("/site/{unique_id}")
 async def remove_site(unique_id: str):
@@ -193,32 +106,6 @@ async def health_check():
     """Health check endpoint"""
     return {"status": "healthy", "sites_count": len(site_manager.active_sites)}
 
-@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
     uvicorn.run(app, host="0.0.0.0", port=8000)