Upload 2 files

support upload of png, md and txt

app/main.py

@@ -1,176 +1,362 @@
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Header, HTTPException
-from fastapi.responses import HTMLResponse, JSONResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.encoders import jsonable_encoder
-
-from uuid import uuid4
-from datetime import datetime, timezone, timedelta
-from typing import Optional
-import os
-import asyncio
-
-from .models import Message, PostRequest, HistoryPayload, NewPostPayload, ErrorPayload
-from .storage import messages, connected_readers, storage_lock
-from .rate_limit import (
-    is_ip_blocked,
-    record_auth_failure,
-    record_auth_success,
-    get_client_ip,
-)
-from .logger_config import logger
-
-POSTER_KEY = os.getenv("POSTER_KEY", "")
-READER_KEY = os.getenv("READER_KEY", "")
-
-app = FastAPI()
-
-# Static HTML (root)
-app.mount("/static", StaticFiles(directory="app/static"), name="static")
-
-
-@app.get("/", response_class=HTMLResponse)
-async def root():
-    with open("app/static/index.html", "r", encoding="utf-8") as f:
-        return f.read()
-
-
-@app.get("/health")
-async def health():
-    async with storage_lock:
-        message_count = len(messages)
-        connected = len(connected_readers)
-    return {
-        "status": "healthy",
-        "message_count": message_count,
-        "connected_readers": connected,
-    }
-
-
-@app.post("/api/v1/posts")
-async def create_post(
-    request: Request,
-    payload: PostRequest,
-    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
-):
-    client_ip = get_client_ip(request)
-
-    # 1) IP block check
-    if is_ip_blocked(client_ip):
-        logger.warning(
-            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /api/v1/posts"
-        )
-        return JSONResponse(
-            status_code=429,
-            content={"error": "Too many failed authentication attempts"},
-        )
-
-    # 2) API key validation
-    if not x_api_key or x_api_key != POSTER_KEY:
-        record_auth_failure(client_ip)
-        logger.warning(
-            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_poster_key - Endpoint: /api/v1/posts"
-        )
-        raise HTTPException(status_code=401, detail="Invalid API key")
-
-    # 3) Defensive length check (pydantic also enforces this)
-    if len(payload.content) > 1000:
-        raise HTTPException(
-            status_code=413, detail="Content exceeds 1000 character limit"
-        )
-
-    now = datetime.now(timezone.utc)
-    msg = Message(
-        id=str(uuid4()),
-        poster_id=payload.poster_id,
-        content=payload.content,
-        timestamp=now,
-        category=payload.category,
-        metadata=payload.metadata or {},
-    )
-
-    async with storage_lock:
-        messages.append(msg)
-
-        # Retention: expire >48h (50-message maxlen still takes precedence)
-        cutoff = now - timedelta(hours=48)
-        tmp = [m for m in messages if m.timestamp >= cutoff]
-        messages.clear()
-        for m in tmp:
-            messages.append(m)
-
-        # Broadcast to all connected readers (JSON-safe)
-        payload_out = jsonable_encoder(NewPostPayload(type="new_post", message=msg))
-        stale = []
-        for ws in connected_readers:
-            try:
-                await ws.send_json(payload_out)
-            except Exception:
-                stale.append(ws)
-        for ws in stale:
-            connected_readers.discard(ws)
-
-    logger.info(
-        f"{datetime.utcnow().isoformat()} - AUTH_SUCCESS - IP: {client_ip} - UserType: poster - Action: post - PosterID: {payload.poster_id}"
-    )
-    record_auth_success(client_ip)
-
-    return JSONResponse(
-        status_code=201,
-        content={
-            "message_id": msg.id,
-            "status": "accepted",
-            "timestamp": msg.timestamp.isoformat(),
-        },
-    )
-
-
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    client_ip = websocket.client.host if websocket.client else "unknown"
-
-    # 1) IP block check
-    if is_ip_blocked(client_ip):
-        logger.warning(
-            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /ws"
-        )
-        await websocket.close(code=1008)
-        return
-
-    # 2) API key validation (query param)
-    api_key = websocket.query_params.get("api_key")
-    if not api_key or api_key != READER_KEY:
-        record_auth_failure(client_ip)
-        logger.warning(
-            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_reader_key - Endpoint: /ws"
-        )
-        await websocket.accept()
-        await websocket.send_json(
-            jsonable_encoder(ErrorPayload(type="error", error="Invalid API key"))
-        )
-        await websocket.close(code=1008)
-        return
-
-    await websocket.accept()
-    logger.info(
-        f"{datetime.utcnow().isoformat()} - AUTH_SUCCESS - IP: {client_ip} - UserType: reader - Action: connect"
-    )
-    record_auth_success(client_ip)
-
-    # Register + send initial history (JSON-safe)
-    async with storage_lock:
-        connected_readers.add(websocket)
-        history_payload = HistoryPayload(type="history", messages=list(messages))
-        await websocket.send_json(jsonable_encoder(history_payload))
-
-    try:
-        while True:
-            # Keepalive: wait for any client message; if none, send ping
-            try:
-                await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
-            except asyncio.TimeoutError:
-                await websocket.send_json({"type": "ping"})
-    except WebSocketDisconnect:
-        pass
-    finally:
-        async with storage_lock:
-            connected_readers.discard(websocket)
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Header, HTTPException, File, UploadFile, Form
+from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.encoders import jsonable_encoder
+
+from uuid import uuid4
+from datetime import datetime, timezone, timedelta
+from typing import Optional
+import os
+import asyncio
+import base64
+import json
+
+from .models import Message, PostRequest, HistoryPayload, NewPostPayload, ErrorPayload, FileContentPayload
+from .storage import messages, connected_readers, storage_lock
+from .rate_limit import (
+    is_ip_blocked,
+    record_auth_failure,
+    record_auth_success,
+    get_client_ip,
+)
+from .logger_config import logger
+
+POSTER_KEY = os.getenv("POSTER_KEY", "")
+READER_KEY = os.getenv("READER_KEY", "")
+
+app = FastAPI()
+
+# Static HTML (root)
+app.mount("/static", StaticFiles(directory="app/static"), name="static")
+
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    with open("app/static/index.html", "r", encoding="utf-8") as f:
+        return f.read()
+
+
+@app.get("/health")
+async def health():
+    async with storage_lock:
+        message_count = len(messages)
+        connected = len(connected_readers)
+    return {
+        "status": "healthy",
+        "message_count": message_count,
+        "connected_readers": connected,
+    }
+
+
+@app.post("/api/v1/posts")
+async def create_post(
+    request: Request,
+    payload: PostRequest,
+    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
+):
+    client_ip = get_client_ip(request)
+
+    # 1) IP block check
+    if is_ip_blocked(client_ip):
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /api/v1/posts"
+        )
+        return JSONResponse(
+            status_code=429,
+            content={"error": "Too many failed authentication attempts"},
+        )
+
+    # 2) API key validation
+    if not x_api_key or x_api_key != POSTER_KEY:
+        record_auth_failure(client_ip)
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_poster_key - Endpoint: /api/v1/posts"
+        )
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    # 3) Defensive length check (pydantic also enforces this)
+    if len(payload.content) > 1000:
+        raise HTTPException(
+            status_code=413, detail="Content exceeds 1000 character limit"
+        )
+
+    now = datetime.now(timezone.utc)
+    msg = Message(
+        id=str(uuid4()),
+        poster_id=payload.poster_id,
+        content=payload.content,
+        timestamp=now,
+        category=payload.category,
+        metadata=payload.metadata or {},
+    )
+
+    async with storage_lock:
+        messages.append(msg)
+
+        # Retention: expire >48h (50-message maxlen still takes precedence)
+        cutoff = now - timedelta(hours=48)
+        tmp = [m for m in messages if m.timestamp >= cutoff]
+        messages.clear()
+        for m in tmp:
+            messages.append(m)
+
+        # Broadcast to all connected readers (JSON-safe)
+        payload_out = jsonable_encoder(NewPostPayload(type="new_post", message=msg))
+        stale = []
+        for ws in connected_readers:
+            try:
+                await ws.send_json(payload_out)
+            except Exception:
+                stale.append(ws)
+        for ws in stale:
+            connected_readers.discard(ws)
+
+    logger.info(
+        f"{datetime.utcnow().isoformat()} - AUTH_SUCCESS - IP: {client_ip} - UserType: poster - Action: post - PosterID: {payload.poster_id}"
+    )
+    record_auth_success(client_ip)
+
+    return JSONResponse(
+        status_code=201,
+        content={
+            "message_id": msg.id,
+            "status": "accepted",
+            "timestamp": msg.timestamp.isoformat(),
+        },
+    )
+
+
+@app.post("/api/v1/upload")
+async def upload_file(
+    request: Request,
+    poster_id: str = Form(...),
+    category: Optional[str] = Form(default=None),
+    file: UploadFile = File(...),
+    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
+):
+    client_ip = get_client_ip(request)
+
+    # 1) IP block check
+    if is_ip_blocked(client_ip):
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /api/v1/upload"
+        )
+        return JSONResponse(
+            status_code=429,
+            content={"error": "Too many failed authentication attempts"},
+        )
+
+    # 2) API key validation
+    if not x_api_key or x_api_key != POSTER_KEY:
+        record_auth_failure(client_ip)
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_poster_key - Endpoint: /api/v1/upload"
+        )
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    # 3) Validate file extension and MIME type
+    filename = file.filename or "unknown"
+    file_ext = os.path.splitext(filename)[1].lower()
+
+    allowed_extensions = {".png", ".md", ".txt"}
+    if file_ext not in allowed_extensions:
+        raise HTTPException(status_code=400, detail="Only PNG, MD, and TXT files are allowed")
+
+    # MIME type validation
+    mime_type = file.content_type or ""
+    valid_mime_types = {
+        ".png": ["image/png"],
+        ".md": ["text/markdown", "text/plain", "application/octet-stream"],
+        ".txt": ["text/plain", "application/octet-stream"],
+    }
+
+    if mime_type and mime_type not in valid_mime_types.get(file_ext, []):
+        # Allow if no content-type provided, otherwise validate
+        pass  # Be lenient with MIME types as they can vary
+
+    # 4) Read file content and validate size
+    content = await file.read()
+    file_size = len(content)
+
+    if file_ext == ".png":
+        if file_size > 2 * 1024 * 1024:  # 2MB
+            raise HTTPException(status_code=413, detail="PNG file exceeds 2MB limit")
+    else:  # .md or .txt
+        if file_size > 100 * 1024:  # 100KB
+            raise HTTPException(status_code=413, detail="Text file exceeds 100KB limit")
+
+    # 5) Process file based on type
+    now = datetime.now(timezone.utc)
+    msg_id = str(uuid4())
+
+    if file_ext == ".png":
+        # Store PNG as base64 data URL
+        base64_data = base64.b64encode(content).decode("utf-8")
+        file_url = f"data:image/png;base64,{base64_data}"
+        msg = Message(
+            id=msg_id,
+            poster_id=poster_id,
+            content=filename,
+            timestamp=now,
+            category=category,
+            metadata={},
+            message_type="png",
+            file_url=file_url,
+            title=filename,
+        )
+    elif file_ext == ".md":
+        # Extract first line as title, store full content
+        text_content = content.decode("utf-8", errors="replace")
+        first_line = text_content.split("\n")[0].strip() or filename
+        msg = Message(
+            id=msg_id,
+            poster_id=poster_id,
+            content=filename,
+            timestamp=now,
+            category=category,
+            metadata={},
+            message_type="md",
+            title=first_line,
+            file_content=text_content,
+        )
+    else:  # .txt
+        # Extract first line as title, store full content
+        text_content = content.decode("utf-8", errors="replace")
+        first_line = text_content.split("\n")[0].strip() or filename
+        msg = Message(
+            id=msg_id,
+            poster_id=poster_id,
+            content=filename,
+            timestamp=now,
+            category=category,
+            metadata={},
+            message_type="txt",
+            title=first_line,
+            file_content=text_content,
+        )
+
+    # 6) Store message and broadcast (same as create_post)
+    async with storage_lock:
+        messages.append(msg)
+
+        # Retention: expire >48h (50-message maxlen still takes precedence)
+        cutoff = now - timedelta(hours=48)
+        tmp = [m for m in messages if m.timestamp >= cutoff]
+        messages.clear()
+        for m in tmp:
+            messages.append(m)
+
+        # Broadcast to all connected readers (JSON-safe)
+        payload_out = jsonable_encoder(NewPostPayload(type="new_post", message=msg))
+        stale = []
+        for ws in connected_readers:
+            try:
+                await ws.send_json(payload_out)
+            except Exception:
+                stale.append(ws)
+        for ws in stale:
+            connected_readers.discard(ws)
+
+    logger.info(
+        f"{datetime.utcnow().isoformat()} - AUTH_SUCCESS - IP: {client_ip} - UserType: poster - Action: upload - PosterID: {poster_id}"
+    )
+    record_auth_success(client_ip)
+
+    return JSONResponse(
+        status_code=201,
+        content={
+            "message_id": msg.id,
+            "status": "accepted",
+            "timestamp": msg.timestamp.isoformat(),
+        },
+    )
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    client_ip = websocket.client.host if websocket.client else "unknown"
+
+    # 1) IP block check
+    if is_ip_blocked(client_ip):
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /ws"
+        )
+        await websocket.close(code=1008)
+        return
+
+    # 2) API key validation (query param)
+    api_key = websocket.query_params.get("api_key")
+    if not api_key or api_key != READER_KEY:
+        record_auth_failure(client_ip)
+        logger.warning(
+            f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_reader_key - Endpoint: /ws"
+        )
+        await websocket.accept()
+        await websocket.send_json(
+            jsonable_encoder(ErrorPayload(type="error", error="Invalid API key"))
+        )
+        await websocket.close(code=1008)
+        return
+
+    await websocket.accept()
+    logger.info(
+        f"{datetime.utcnow().isoformat()} - AUTH_SUCCESS - IP: {client_ip} - UserType: reader - Action: connect"
+    )
+    record_auth_success(client_ip)
+
+    # Register + send initial history (JSON-safe)
+    async with storage_lock:
+        connected_readers.add(websocket)
+        history_payload = HistoryPayload(type="history", messages=list(messages))
+        await websocket.send_json(jsonable_encoder(history_payload))
+
+    try:
+        while True:
+            # Keepalive: wait for client message; if none, send ping
+            try:
+                data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
+                # Handle file content request
+                try:
+                    msg = json.loads(data)
+                    if msg.get("type") == "get_file_content":
+                        message_id = msg.get("message_id")
+                        async with storage_lock:
+                            target_msg = None
+                            for m in messages:
+                                if m.id == message_id:
+                                    target_msg = m
+                                    break
+
+                        if target_msg is None:
+                            await websocket.send_json(
+                                jsonable_encoder(ErrorPayload(
+                                    type="error",
+                                    error="Message not found"
+                                ))
+                            )
+                        elif target_msg.message_type not in ("md", "txt"):
+                            await websocket.send_json(
+                                jsonable_encoder(ErrorPayload(
+                                    type="error",
+                                    error="Invalid message type"
+                                ))
+                            )
+                        else:
+                            await websocket.send_json(
+                                jsonable_encoder(FileContentPayload(
+                                    type="file_content",
+                                    message_id=message_id,
+                                    content=target_msg.file_content or "",
+                                    content_type=target_msg.message_type
+                                ))
+                            )
+                except json.JSONDecodeError:
+                    # Ignore non-JSON messages (keepalive behavior)
+                    pass
+            except asyncio.TimeoutError:
+                await websocket.send_json({"type": "ping"})
+    except WebSocketDisconnect:
+        pass
+    finally:
+        async with storage_lock:
+            connected_readers.discard(websocket)

app/models.py

@@ -1,34 +1,46 @@
-from pydantic import BaseModel, Field
-from datetime import datetime
-from typing import Optional, Dict, Any, List
-
-
-class PostRequest(BaseModel):
-    poster_id: str = Field(..., min_length=1)
-    content: str = Field(..., min_length=1, max_length=1000)
-    category: Optional[str] = None
-    metadata: Optional[Dict[str, Any]] = None
-
-
-class Message(BaseModel):
-    id: str
-    poster_id: str
-    content: str
-    timestamp: datetime
-    category: Optional[str] = None
-    metadata: Dict[str, Any] = {}
-
-
-class HistoryPayload(BaseModel):
-    type: str = "history"
-    messages: List[Message]
-
-
-class NewPostPayload(BaseModel):
-    type: str = "new_post"
-    message: Message
-
-
-class ErrorPayload(BaseModel):
-    type: str = "error"
-    error: str
+from pydantic import BaseModel, Field
+from datetime import datetime
+from typing import Optional, Dict, Any, List, Literal
+
+
+class PostRequest(BaseModel):
+    poster_id: str = Field(..., min_length=1)
+    content: str = Field(..., min_length=1, max_length=1000)
+    category: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class Message(BaseModel):
+    id: str
+    poster_id: str
+    content: str
+    timestamp: datetime
+    category: Optional[str] = None
+    metadata: Dict[str, Any] = {}
+    # File-related fields
+    message_type: Literal["text", "png", "md", "txt"] = "text"
+    file_url: Optional[str] = None
+    title: Optional[str] = None
+    file_content: Optional[str] = None  # For MD/TXT - stored but not shown in list
+
+
+class HistoryPayload(BaseModel):
+    type: str = "history"
+    messages: List[Message]
+
+
+class NewPostPayload(BaseModel):
+    type: str = "new_post"
+    message: Message
+
+
+class ErrorPayload(BaseModel):
+    type: str = "error"
+    error: str
+
+
+class FileContentPayload(BaseModel):
+    type: str = "file_content"
+    message_id: str
+    content: str
+    content_type: Literal["md", "txt"]