MP-Server/web_server.py

# FastAPI web server for MacroPad

import io
import os
import sys
import uuid
import hmac
import asyncio
from typing import List, Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, HTTPException, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from pydantic import BaseModel
from PIL import Image
import uvicorn

from config import DEFAULT_PORT, VERSION

# Max accepted upload size for images (bytes).
MAX_IMAGE_UPLOAD_BYTES = 5 * 1024 * 1024

# Map Pillow image formats to safe file extensions we're willing to persist.
_IMAGE_FORMAT_EXT = {
    "PNG": ".png",
    "JPEG": ".jpg",
    "GIF": ".gif",
    "BMP": ".bmp",
    "WEBP": ".webp",
}


def get_resource_path(relative_path):
    """Get the path to a bundled resource file."""
    if getattr(sys, 'frozen', False):
        base_path = sys._MEIPASS
    else:
        base_path = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(base_path, relative_path)


class Command(BaseModel):
    """Single command in a macro sequence."""
    type: str  # text, key, hotkey, wait, app
    value: Optional[str] = None
    keys: Optional[List[str]] = None
    ms: Optional[int] = None
    command: Optional[str] = None


class MacroCreate(BaseModel):
    """Request model for creating a macro."""
    name: str
    commands: List[Command]
    category: Optional[str] = ""


class MacroUpdate(BaseModel):
    """Request model for updating a macro."""
    name: str
    commands: List[Command]
    category: Optional[str] = ""


class ExecuteRequest(BaseModel):
    """Request model for executing a macro."""
    macro_id: str


class TabCreate(BaseModel):
    """Request model for creating a tab."""
    name: str


class ConnectionManager:
    """Manages WebSocket connections for real-time updates."""

    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        """Send message to all connected clients, pruning dead sockets."""
        dead = []
        # Iterate over a snapshot so the list isn't mutated during iteration.
        for connection in list(self.active_connections):
            try:
                await connection.send_json(message)
            except Exception:
                dead.append(connection)
        for connection in dead:
            self.disconnect(connection)


class WebServer:
    """FastAPI-based web server for MacroPad."""

    def __init__(self, macro_manager, app_dir: str, port: int = DEFAULT_PORT,
                 auth_token: str = "", host: str = "127.0.0.1"):
        self.macro_manager = macro_manager
        self.app_dir = app_dir
        self.port = port
        self.auth_token = auth_token
        self.host = host
        self.app = None
        self.manager = ConnectionManager()
        self.server = None

    def create_app(self) -> FastAPI:
        """Create FastAPI application."""

        @asynccontextmanager
        async def lifespan(app: FastAPI):
            yield

        app = FastAPI(
            title="MacroPad Server",
            version=VERSION,
            lifespan=lifespan
        )

        # Serve static files from web directory (bundled with app)
        web_dir = get_resource_path("web")
        if os.path.exists(web_dir):
            app.mount("/static", StaticFiles(directory=web_dir), name="static")

        # Serve macro images (user data directory)
        images_dir = os.path.join(self.app_dir, "macro_images")
        if os.path.exists(images_dir):
            app.mount("/images", StaticFiles(directory=images_dir), name="images")

        # Paths served without a token (PWA shell + static assets)
        public_paths = {"/", "/manifest.json", "/service-worker.js"}

        def _add_security_headers(response):
            # Anti-clickjacking / MIME-sniffing. frame-ancestors can only be set
            # via a real response header (it is ignored inside a <meta> CSP), so
            # we enforce it here rather than in index.html.
            response.headers["X-Frame-Options"] = "DENY"
            response.headers["X-Content-Type-Options"] = "nosniff"
            response.headers["Referrer-Policy"] = "no-referrer"
            return response

        @app.middleware("http")
        async def auth_middleware(request: Request, call_next):
            """Require a valid auth token for all non-public paths."""
            path = request.url.path
            if path in public_paths or path.startswith("/static/"):
                return _add_security_headers(await call_next(request))

            # Fail-open only when no token is configured (shouldn't happen in prod)
            if self.auth_token:
                provided = request.headers.get("X-MacroPad-Token") \
                    or request.query_params.get("token") or ""
                if not hmac.compare_digest(provided, self.auth_token):
                    return _add_security_headers(
                        JSONResponse({"detail": "Unauthorized"}, status_code=401)
                    )

            return _add_security_headers(await call_next(request))

        @app.get("/", response_class=HTMLResponse)
        async def index():
            """Serve the main PWA page."""
            index_path = os.path.join(web_dir, "index.html")
            if os.path.exists(index_path):
                return FileResponse(index_path, media_type="text/html")
            return HTMLResponse("<h1>MacroPad Server</h1><p>Web interface not found.</p>")

        @app.get("/manifest.json")
        async def manifest():
            """Serve PWA manifest."""
            manifest_path = os.path.join(web_dir, "manifest.json")
            if os.path.exists(manifest_path):
                return FileResponse(manifest_path, media_type="application/json")
            raise HTTPException(status_code=404, detail="Manifest not found")

        @app.get("/service-worker.js")
        async def service_worker():
            """Serve service worker."""
            sw_path = os.path.join(web_dir, "service-worker.js")
            if os.path.exists(sw_path):
                return FileResponse(sw_path, media_type="application/javascript")
            raise HTTPException(status_code=404, detail="Service worker not found")

        @app.get("/api/tabs")
        async def get_tabs():
            """Get available tab categories."""
            return {"tabs": self.macro_manager.get_unique_tabs()}

        @app.get("/api/macros")
        async def get_macros():
            """Get all macros."""
            macros = self.macro_manager.get_all_macros()
            return {"macros": macros}

        @app.get("/api/macros/{tab}")
        async def get_macros_by_tab(tab: str):
            """Get macros filtered by tab/category."""
            all_macros = self.macro_manager.get_sorted_macros()
            filtered = self.macro_manager.filter_macros_by_tab(all_macros, tab)
            return {"macros": dict(filtered)}

        @app.get("/api/macro/{macro_id}")
        async def get_macro(macro_id: str):
            """Get a single macro by ID."""
            macro = self.macro_manager.get_macro(macro_id)
            if macro:
                return {"macro": macro}
            raise HTTPException(status_code=404, detail="Macro not found")

        @app.post("/api/execute")
        async def execute_macro(request: ExecuteRequest):
            """Execute a macro by ID."""
            # execute_macro is synchronous and may sleep (wait steps); run it in
            # a thread executor so it doesn't block the event loop / other traffic.
            loop = asyncio.get_event_loop()
            success = await loop.run_in_executor(
                None, self.macro_manager.execute_macro, request.macro_id
            )
            if success:
                # Broadcast execution to all connected clients
                await self.manager.broadcast({
                    "type": "executed",
                    "macro_id": request.macro_id
                })
                return {"success": True}
            raise HTTPException(status_code=404, detail="Macro not found or execution failed")

        @app.post("/api/macros")
        async def create_macro(macro: MacroCreate):
            """Create a new macro."""
            commands = [cmd.model_dump(exclude_none=True) for cmd in macro.commands]
            macro_id = self.macro_manager.add_macro(
                name=macro.name,
                commands=commands,
                category=macro.category or ""
            )
            # Broadcast update
            await self.manager.broadcast({"type": "macro_created", "macro_id": macro_id})
            return {"success": True, "macro_id": macro_id}

        @app.put("/api/macros/{macro_id}")
        async def update_macro(macro_id: str, macro: MacroUpdate):
            """Update an existing macro."""
            commands = [cmd.model_dump(exclude_none=True) for cmd in macro.commands]
            success = self.macro_manager.update_macro(
                macro_id=macro_id,
                name=macro.name,
                commands=commands,
                category=macro.category or ""
            )
            if success:
                await self.manager.broadcast({"type": "macro_updated", "macro_id": macro_id})
                return {"success": True}
            raise HTTPException(status_code=404, detail="Macro not found")

        @app.delete("/api/macros/{macro_id}")
        async def delete_macro(macro_id: str):
            """Delete a macro."""
            success = self.macro_manager.delete_macro(macro_id)
            if success:
                await self.manager.broadcast({"type": "macro_deleted", "macro_id": macro_id})
                return {"success": True}
            raise HTTPException(status_code=404, detail="Macro not found")

        @app.post("/api/upload-image")
        async def upload_image(file: UploadFile = File(...)):
            """Upload an image for a macro.

            Hardened: enforces a size cap, validates the bytes are a real image
            with Pillow (content_type is untrusted), and saves into the app's
            macro_images/ directory under a server-generated uuid name. Returns
            only a relative reference, never an absolute server path.
            """
            # Read with a hard size cap (reject oversize before buffering it all).
            data = b""
            chunk_size = 64 * 1024
            while True:
                chunk = await file.read(chunk_size)
                if not chunk:
                    break
                data += chunk
                if len(data) > MAX_IMAGE_UPLOAD_BYTES:
                    raise HTTPException(status_code=413, detail="Image too large")

            if not data:
                raise HTTPException(status_code=400, detail="Empty upload")

            # Validate it's really an image (don't trust content_type).
            try:
                Image.open(io.BytesIO(data)).verify()
                # verify() leaves the image unusable; reopen to read the format.
                img_format = (Image.open(io.BytesIO(data)).format or "").upper()
            except Exception:
                raise HTTPException(status_code=400, detail="File is not a valid image")

            ext = _IMAGE_FORMAT_EXT.get(img_format)
            if not ext:
                raise HTTPException(status_code=400, detail="Unsupported image format")

            images_dir = os.path.join(self.app_dir, "macro_images")
            os.makedirs(images_dir, exist_ok=True)

            filename = f"{uuid.uuid4().hex}{ext}"
            dest_path = os.path.join(images_dir, filename)
            rel_path = os.path.join("macro_images", filename)

            try:
                with open(dest_path, "wb") as out:
                    out.write(data)
            except Exception:
                # Clean up a partial file on failure.
                try:
                    if os.path.exists(dest_path):
                        os.remove(dest_path)
                except OSError:
                    pass
                raise HTTPException(status_code=500, detail="Failed to save image")

            # Return only a relative, opaque reference.
            return {"path": rel_path.replace(os.sep, "/")}

        @app.get("/api/image/{image_path:path}")
        async def get_image(image_path: str):
            """Get macro image (legacy compatibility, path-traversal safe)."""
            images_dir = os.path.join(self.app_dir, "macro_images")
            requested = os.path.realpath(os.path.join(self.app_dir, image_path))
            images_real = os.path.realpath(images_dir)
            # Only serve files that resolve to inside the macro_images directory
            if (os.path.commonpath([requested, images_real]) == images_real
                    and os.path.isfile(requested)):
                return FileResponse(requested)
            raise HTTPException(status_code=404, detail="Image not found")

        @app.websocket("/ws")
        async def websocket_endpoint(websocket: WebSocket):
            """WebSocket for real-time updates."""
            # Require a valid token before accepting the connection
            if self.auth_token:
                provided = websocket.query_params.get("token") or ""
                if not hmac.compare_digest(provided, self.auth_token):
                    await websocket.close(code=1008)
                    return
            await self.manager.connect(websocket)
            try:
                while True:
                    # receive_json raises on malformed input / non-JSON frames;
                    # any such error must still tear the connection down cleanly.
                    data = await websocket.receive_json()
                    # Handle incoming messages if needed
                    if isinstance(data, dict) and data.get("type") == "ping":
                        await websocket.send_json({"type": "pong"})
            except WebSocketDisconnect:
                pass
            except Exception:
                # Malformed input or any other error: don't leak the connection.
                pass
            finally:
                self.manager.disconnect(websocket)

        self.app = app
        return app

    def run(self):
        """Run the web server (blocking)."""
        if not self.app:
            self.create_app()

        config = uvicorn.Config(
            self.app,
            host=self.host,
            port=self.port,
            log_level="warning",
            log_config=None  # Disable default logging config for PyInstaller compatibility
        )
        self.server = uvicorn.Server(config)
        self.server.run()

    def stop(self):
        """Stop the web server."""
        if self.server:
            self.server.should_exit = True

    async def run_async(self):
        """Run the web server asynchronously."""
        if not self.app:
            self.create_app()

        config = uvicorn.Config(
            self.app,
            host=self.host,
            port=self.port,
            log_level="warning",
            log_config=None  # Disable default logging config for PyInstaller compatibility
        )
        self.server = uvicorn.Server(config)
        await self.server.serve()