Update app.py

app.py

@@ -3,26 +3,23 @@ import json
 import time
 import uuid
 import base64
-from typing import Optional, Dict, Set, List, Any
-import asyncio
-import aiosqlite
 import logging
+from typing import Optional, Dict, Set, List, Any
 from contextlib import asynccontextmanager
-from datetime import datetime
-from urllib.parse import unquote
 
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, Request, UploadFile, File
+import aiosqlite
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, UploadFile, File
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import Response, JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.middleware.trustedhost import TrustedHostMiddleware
 
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.backends import default_backend
 
-from storage_handler import store_file, retrieve_file, delete_file, get_bucket_info, file_exists
+from storage_handler import store_file, retrieve_file, delete_file, file_exists
 
 # ------------------------------------------------------------------------
 # Configuration
@@ -31,27 +28,23 @@ DATABASE_URL = os.environ.get("DATABASE_URL", "/data/infinitychat.db")
 MESSAGE_KEY_B64 = os.environ.get("SECRET_KEY", None)
 if MESSAGE_KEY_B64 is None:
     MESSAGE_KEY_B64 = base64.urlsafe_b64encode(os.urandom(32)).decode()
-    print("WARNING: SECRET_KEY not set – generated random key. Set it permanently in production!")
+    print("WARNING: SECRET_KEY not set – generated random key. Set it permanently!")
 MESSAGE_KEY = base64.urlsafe_b64decode(MESSAGE_KEY_B64)
-assert len(MESSAGE_KEY) == 32, "SECRET_KEY must decode to exactly 32 bytes"
+assert len(MESSAGE_KEY) == 32
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("InfinityChat")
 
 # ------------------------------------------------------------------------
-# Database initialization
+# Database initialization (once at startup)
 # ------------------------------------------------------------------------
 async def init_database():
-    """Initialize SQLite database with all required tables and indexes."""
     async with aiosqlite.connect(DATABASE_URL) as db:
         db.row_factory = aiosqlite.Row
         await db.execute("PRAGMA journal_mode=WAL;")
         await db.execute("PRAGMA foreign_keys=ON;")
         await db.execute("PRAGMA synchronous=NORMAL;")
-        await db.execute("PRAGMA cache_size=-20000;")  # 20MB cache
-        await db.execute("PRAGMA busy_timeout=5000;")
-        
-        # Users table
+        await db.execute("PRAGMA cache_size=-20000;")
         await db.execute("""
             CREATE TABLE IF NOT EXISTS users (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -66,8 +59,6 @@ async def init_database():
                 status TEXT DEFAULT 'offline'
             )
         """)
-        
-        # Messages table with optimized indexes
         await db.execute("""
             CREATE TABLE IF NOT EXISTS messages (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -84,8 +75,6 @@ async def init_database():
                 FOREIGN KEY(sender_id) REFERENCES users(id)
             )
         """)
-        
-        # Read receipts
         await db.execute("""
             CREATE TABLE IF NOT EXISTS read_receipts (
                 user_id INTEGER NOT NULL,
@@ -96,620 +85,309 @@ async def init_database():
                 FOREIGN KEY(message_id) REFERENCES messages(id)
             )
         """)
-        
-        # Conversation index for faster queries
-        await db.execute("""
-            CREATE INDEX IF NOT EXISTS idx_messages_sender_time 
-            ON messages(sender_id, timestamp_ms)
-        """)
-        await db.execute("""
-            CREATE INDEX IF NOT EXISTS idx_messages_time_id 
-            ON messages(timestamp_ms, id)
-        """)
-        await db.execute("""
-            CREATE INDEX IF NOT EXISTS idx_messages_not_deleted 
-            ON messages(is_deleted, id)
-        """)
-        await db.execute("""
-            CREATE INDEX IF NOT EXISTS idx_receipts_user_msg 
-            ON read_receipts(user_id, message_id)
-        """)
-        
+        await db.execute("CREATE INDEX IF NOT EXISTS idx_messages_sender_time ON messages(sender_id, timestamp_ms)")
+        await db.execute("CREATE INDEX IF NOT EXISTS idx_messages_time_id ON messages(timestamp_ms, id)")
+        await db.execute("CREATE INDEX IF NOT EXISTS idx_messages_not_deleted ON messages(is_deleted, id)")
+        await db.execute("CREATE INDEX IF NOT EXISTS idx_receipts_user_msg ON read_receipts(user_id, message_id)")
         await db.commit()
-    
     logger.info("✅ Database initialized successfully")
 
 # ------------------------------------------------------------------------
-# FastAPI app with lifespan
+# FastAPI lifespan
 # ------------------------------------------------------------------------
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Handle startup and shutdown events."""
     await init_database()
     yield
 
-app = FastAPI(
-    title="InfinityChat",
-    version="1.0.0",
-    lifespan=lifespan,
-    docs_url=None,
-    redoc_url=None
-)
-
-# Middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Static files
+app = FastAPI(title="InfinityChat", version="1.0.0", lifespan=lifespan)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # ------------------------------------------------------------------------
-# Database connection pool
+# Database helper (new connection each time)
 # ------------------------------------------------------------------------
-class DatabasePool:
-    """Simple database connection management."""
-    
-    def __init__(self):
-        self._conn = None
-    
-    async def get_connection(self) -> aiosqlite.Connection:
-        if self._conn is None:
-            self._conn = await aiosqlite.connect(DATABASE_URL)
-            self._conn.row_factory = aiosqlite.Row
-            await self._conn.execute("PRAGMA journal_mode=WAL")
-            await self._conn.execute("PRAGMA foreign_keys=ON")
-            await self._conn.execute("PRAGMA synchronous=NORMAL")
-        return self._conn
-    
-    async def close(self):
-        if self._conn:
-            await self._conn.close()
-            self._conn = None
-
-db_pool = DatabasePool()
-
 async def get_db() -> aiosqlite.Connection:
-    """Get or create a database connection."""
-    return await db_pool.get_connection()
+    conn = await aiosqlite.connect(DATABASE_URL)
+    conn.row_factory = aiosqlite.Row
+    await conn.execute("PRAGMA journal_mode=WAL")
+    await conn.execute("PRAGMA foreign_keys=ON")
+    return conn
 
 # ------------------------------------------------------------------------
-# Encryption (AES-256-GCM)
+# Encryption helpers
 # ------------------------------------------------------------------------
 def encrypt_message(plaintext: str) -> str:
-    """Encrypt message content with AES-256-GCM."""
     aesgcm = AESGCM(MESSAGE_KEY)
     nonce = os.urandom(12)
     ciphertext = aesgcm.encrypt(nonce, plaintext.encode('utf-8'), None)
     return base64.urlsafe_b64encode(nonce + ciphertext).decode('utf-8')
 
 def decrypt_message(encrypted_b64: str) -> str:
-    """Decrypt message content."""
     try:
         raw = base64.urlsafe_b64decode(encrypted_b64)
         nonce, ciphertext = raw[:12], raw[12:]
         aesgcm = AESGCM(MESSAGE_KEY)
         return aesgcm.decrypt(nonce, ciphertext, None).decode('utf-8')
-    except Exception as e:
-        logger.error(f"Decryption failed: {e}")
+    except Exception:
         return "[Decryption Error]"
 
 # ------------------------------------------------------------------------
-# Password hashing (PBKDF2)
+# Password hashing
 # ------------------------------------------------------------------------
 def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
-    """Hash password with PBKDF2-SHA256."""
     if salt is None:
         salt = os.urandom(32).hex()
-    kdf = PBKDF2HMAC(
-        algorithm=hashes.SHA256(),
-        length=32,
-        salt=salt.encode('utf-8'),
-        iterations=600000,
-        backend=default_backend()
-    )
-    key = base64.urlsafe_b64encode(kdf.derive(password.encode('utf-8'))).decode('utf-8')
+    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), iterations=600000, backend=default_backend())
+    key = base64.urlsafe_b64encode(kdf.derive(password.encode())).decode()
     return key, salt
 
 def verify_password(password: str, salt: str, stored_hash: str) -> bool:
-    """Verify password against stored hash."""
     key, _ = hash_password(password, salt)
     return key == stored_hash
 
 def generate_token() -> str:
-    """Generate a secure random token."""
-    return base64.urlsafe_b64encode(os.urandom(32)).decode('utf-8')
+    return base64.urlsafe_b64encode(os.urandom(32)).decode()
 
 # ------------------------------------------------------------------------
-# Authentication helpers
+# Authentication
 # ------------------------------------------------------------------------
 async def authenticate_user(token: str) -> Optional[dict]:
-    """Validate token and return user data."""
     if not token or token in ('null', 'undefined'):
         return None
-    
     db = await get_db()
     try:
-        async with db.execute(
-            "SELECT id, username, display_name, avatar_path, status FROM users WHERE token = ?",
-            (token,)
-        ) as cursor:
-            user = await cursor.fetchone()
-            if user:
-                # Update last_seen
-                await db.execute(
-                    "UPDATE users SET last_seen = ? WHERE id = ?",
-                    (int(time.time()), user['id'])
-                )
-                await db.commit()
-                return dict(user)
-            return None
-    except Exception as e:
-        logger.error(f"Authentication error: {e}")
+        cursor = await db.execute("SELECT id, username, display_name, avatar_path, status FROM users WHERE token = ?", (token,))
+        user = await cursor.fetchone()
+        if user:
+            await db.execute("UPDATE users SET last_seen = ? WHERE id = ?", (int(time.time()), user['id']))
+            await db.commit()
+            return dict(user)
         return None
+    finally:
+        await db.close()
 
 # ------------------------------------------------------------------------
 # HTTP Endpoints
 # ------------------------------------------------------------------------
 @app.get("/api/health")
 async def health_check():
-    """Health check endpoint."""
-    return {
-        "status": "healthy",
-        "timestamp": int(time.time()),
-        "version": "1.0.0"
-    }
+    return {"status": "healthy", "timestamp": int(time.time()), "version": "1.0.0"}
 
 @app.post("/api/auth/signup")
-async def signup(
-    username: str = Query(..., min_length=3, max_length=30),
-    password: str = Query(..., min_length=6),
-    display_name: Optional[str] = Query(None, max_length=50)
-):
-    """Create a new user account."""
-    # Validate username format
+async def signup(username: str = Query(...), password: str = Query(...), display_name: str = Query(None)):
     if not username.isalnum() and '_' not in username:
-        raise HTTPException(400, "Username can only contain letters, numbers, and underscores")
-    
+        raise HTTPException(400, "Invalid username")
     db = await get_db()
     try:
-        # Check if username exists
-        async with db.execute(
-            "SELECT id FROM users WHERE username = ?", (username.lower(),)
-        ) as cursor:
+        async with db.execute("SELECT id FROM users WHERE username = ?", (username.lower(),)) as cursor:
             if await cursor.fetchone():
-                raise HTTPException(409, "Username is already taken")
-        
-        # Create user
+                raise HTTPException(409, "Username taken")
         pwd_hash, salt = hash_password(password)
         token = generate_token()
         display = display_name or username
-        
-        cursor = await db.execute(
-            """INSERT INTO users (username, display_name, password_hash, salt, token, status)
-               VALUES (?, ?, ?, ?, ?, 'online')""",
+        await db.execute(
+            "INSERT INTO users (username, display_name, password_hash, salt, token, status) VALUES (?,?,?,?,?, 'online')",
             (username.lower(), display, pwd_hash, salt, token)
         )
         await db.commit()
-        user_id = cursor.lastrowid
-        
-        logger.info(f"New user registered: {username} (ID: {user_id})")
-        
-        return {
-            "token": token,
-            "user": {
-                "id": user_id,
-                "username": username,
-                "display_name": display,
-                "avatar_path": None,
-                "status": "online"
-            }
-        }
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Signup error: {e}")
-        raise HTTPException(500, "Internal server error")
+        return {"token": token, "user": {"username": username, "display_name": display, "avatar_path": None, "status": "online"}}
     finally:
         await db.close()
 
 @app.post("/api/auth/login")
-async def login(
-    username: str = Query(...),
-    password: str = Query(...)
-):
-    """Authenticate user and return token."""
+async def login(username: str = Query(...), password: str = Query(...)):
     db = await get_db()
     try:
-        async with db.execute(
-            "SELECT * FROM users WHERE username = ?", (username.lower(),)
-        ) as cursor:
+        async with db.execute("SELECT * FROM users WHERE username = ?", (username.lower(),)) as cursor:
             user = await cursor.fetchone()
-        
-        if not user:
-            # Use constant-time comparison to prevent timing attacks
-            verify_password("dummy", os.urandom(32).hex(), 
-                          base64.urlsafe_b64encode(os.urandom(32)).decode())
-            raise HTTPException(401, "Invalid username or password")
-        
-        if not verify_password(password, user['salt'], user['password_hash']):
-            raise HTTPException(401, "Invalid username or password")
-        
-        # Generate new token
+        if not user or not verify_password(password, user['salt'], user['password_hash']):
+            raise HTTPException(401, "Invalid credentials")
         token = generate_token()
-        await db.execute(
-            "UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?",
-            (token, int(time.time()), user['id'])
-        )
+        await db.execute("UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?", (token, int(time.time()), user['id']))
         await db.commit()
-        
-        logger.info(f"User logged in: {user['username']}")
-        
-        return {
-            "token": token,
-            "user": {
-                "id": user['id'],
-                "username": user['username'],
-                "display_name": user['display_name'],
-                "avatar_path": user['avatar_path'],
-                "status": "online"
-            }
-        }
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Login error: {e}")
-        raise HTTPException(500, "Internal server error")
+        return {"token": token, "user": {"id": user['id'], "username": user['username'], "display_name": user['display_name'], "avatar_path": user['avatar_path'], "status": "online"}}
     finally:
         await db.close()
 
 @app.get("/api/auth/verify")
 async def verify_token(token: str = Header(..., alias="X-Auth-Token")):
-    """Verify authentication token."""
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401, "Invalid or expired token")
+    if not user: raise HTTPException(401)
     return {"user": user}
 
 @app.post("/api/auth/logout")
 async def logout(token: str = Header(..., alias="X-Auth-Token")):
-    """Logout user by invalidating token."""
     db = await get_db()
     try:
-        await db.execute(
-            "UPDATE users SET token = NULL, status = 'offline', last_seen = ? WHERE token = ?",
-            (int(time.time()), token)
-        )
+        await db.execute("UPDATE users SET token = NULL, status = 'offline', last_seen = ? WHERE token = ?", (int(time.time()), token))
         await db.commit()
         return {"status": "logged_out"}
     finally:
         await db.close()
 
 # ------------------------------------------------------------------------
-# Profile Endpoints
+# Profile
 # ------------------------------------------------------------------------
 @app.get("/api/profile")
 async def get_profile(token: str = Header(..., alias="X-Auth-Token")):
-    """Get current user profile."""
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401)
+    if not user: raise HTTPException(401)
     return {"user": user}
 
 @app.patch("/api/profile")
-async def update_profile(
-    display_name: Optional[str] = Query(None, max_length=50),
-    token: str = Header(..., alias="X-Auth-Token")
-):
-    """Update user profile."""
+async def update_profile(display_name: str = Query(...), token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401)
-    
+    if not user: raise HTTPException(401)
     db = await get_db()
     try:
-        if display_name:
-            await db.execute(
-                "UPDATE users SET display_name = ? WHERE id = ?",
-                (display_name, user['id'])
-            )
-            await db.commit()
-            user['display_name'] = display_name
-        
+        await db.execute("UPDATE users SET display_name = ? WHERE id = ?", (display_name, user['id']))
+        await db.commit()
+        user['display_name'] = display_name
         return {"user": user}
     finally:
         await db.close()
 
 @app.post("/api/profile/avatar")
-async def upload_avatar(
-    file: UploadFile = File(...),
-    token: str = Header(..., alias="X-Auth-Token")
-):
-    """Upload user avatar."""
+async def upload_avatar(file: UploadFile = File(...), token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401)
-    
-    # Validate file type
-    allowed_types = {'image/jpeg', 'image/png', 'image/gif', 'image/webp'}
-    if file.content_type not in allowed_types:
-        raise HTTPException(400, "Only JPEG, PNG, GIF, and WebP images are allowed")
-    
-    # Validate file size (max 5MB)
+    if not user: raise HTTPException(401)
     data = await file.read()
-    if len(data) > 5 * 1024 * 1024:
-        raise HTTPException(400, "Image must be less than 5MB")
-    
-    # Generate unique path
+    if len(data) > 5*1024*1024:
+        raise HTTPException(400, "Image too large")
     ext = os.path.splitext(file.filename)[1] if file.filename else '.jpg'
     remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
-    
-    try:
-        store_file(remote_path, data)
-    except Exception as e:
-        logger.error(f"Avatar upload failed: {e}")
-        raise HTTPException(500, "Failed to upload avatar")
-    
+    store_file(remote_path, data)
     db = await get_db()
     try:
-        await db.execute(
-            "UPDATE users SET avatar_path = ? WHERE id = ?",
-            (remote_path, user['id'])
-        )
+        await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote_path, user['id']))
         await db.commit()
-        user['avatar_path'] = remote_path
     finally:
         await db.close()
-    
     return {"avatar_path": remote_path}
 
 # ------------------------------------------------------------------------
-# File Upload (Chunked)
+# Chunked Upload
 # ------------------------------------------------------------------------
-upload_sessions: Dict[str, Dict] = {}
+upload_sessions: Dict[str, dict] = {}
 
 @app.post("/api/upload/chunk")
-async def upload_chunk(
-    file: UploadFile = File(...),
-    chunk_index: int = Query(..., ge=0),
-    total_chunks: int = Query(..., ge=1),
-    file_name: str = Query(..., max_length=255),
-    file_type: str = Query(..., max_length=100),
-    file_size: int = Query(..., ge=0),
-    upload_id: str = Query(..., min_length=8),
-    token: str = Header(..., alias="X-Auth-Token")
-):
-    """Handle chunked file upload."""
+async def upload_chunk(file: UploadFile = File(...), chunk_index: int = Query(...), total_chunks: int = Query(...),
+                       file_name: str = Query(...), file_type: str = Query(...), file_size: int = Query(...),
+                       upload_id: str = Query(...), token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401)
-    
-    # Validate file size (max 100MB total)
-    if file_size > 100 * 1024 * 1024:
-        raise HTTPException(400, "File must be less than 100MB")
-    
-    # Initialize upload session if new
+    if not user: raise HTTPException(401)
     if upload_id not in upload_sessions:
-        upload_sessions[upload_id] = {
-            "chunks": {},
-            "filename": file_name,
-            "file_type": file_type,
-            "file_size": file_size,
-            "total_chunks": total_chunks,
-            "received": 0,
-            "user_id": user['id'],
-            "created_at": time.time()
-        }
-    
+        upload_sessions[upload_id] = {"chunks": {}, "filename": file_name, "file_type": file_type,
+                                      "file_size": file_size, "total_chunks": total_chunks, "received": 0,
+                                      "user_id": user['id'], "created_at": time.time()}
     session = upload_sessions[upload_id]
-    
-    # Verify ownership
-    if session['user_id'] != user['id']:
-        raise HTTPException(403, "Upload session belongs to another user")
-    
-    # Store chunk
+    if session['user_id'] != user['id']: raise HTTPException(403)
     chunk_data = await file.read()
-    if chunk_index in session['chunks']:
-        return {"status": "duplicate", "received": session['received'], "total": total_chunks}
-    
     session['chunks'][chunk_index] = chunk_data
     session['received'] += 1
-    
-    # Check if upload complete
     if session['received'] == total_chunks:
-        # Reassemble file
-        full_data = b"".join(session['chunks'][i] for i in sorted(session['chunks'].keys()))
-        
-        # Generate storage path
+        full_data = b"".join(session['chunks'][i] for i in sorted(session['chunks']))
         remote_path = f"uploads/{user['username']}/{uuid.uuid4().hex}/{file_name}"
-        
-        try:
-            store_file(remote_path, full_data)
-        except Exception as e:
-            logger.error(f"Failed to store uploaded file: {e}")
-            del upload_sessions[upload_id]
-            raise HTTPException(500, "Failed to store file")
-        
-        # Cleanup
+        store_file(remote_path, full_data)
         del upload_sessions[upload_id]
-        
-        # Cleanup old sessions (older than 30 minutes)
-        current_time = time.time()
-        expired = [uid for uid, s in upload_sessions.items() 
-                  if current_time - s['created_at'] > 1800]
-        for uid in expired:
-            del upload_sessions[uid]
-        
-        return {
-            "status": "complete",
-            "file_path": remote_path,
-            "file_name": file_name,
-            "file_type": file_type,
-            "file_size": file_size
-        }
-    
-    return {
-        "status": "in_progress",
-        "received": session['received'],
-        "total": total_chunks
-    }
+        return {"status": "complete", "file_path": remote_path, "file_name": file_name, "file_type": file_type, "file_size": file_size}
+    return {"status": "in_progress", "received": session['received'], "total": total_chunks}
 
 # ------------------------------------------------------------------------
 # File Download
 # ------------------------------------------------------------------------
 @app.get("/api/download/{file_path:path}")
-async def download_file(
-    file_path: str,
-    token: str = Header(..., alias="X-Auth-Token")
-):
-    """Download a file with authentication."""
+async def download_file(file_path: str, token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user:
-        raise HTTPException(401)
-    
-    # Security: prevent path traversal
-    if '..' in file_path or file_path.startswith('/'):
-        raise HTTPException(400, "Invalid file path")
-    
+    if not user: raise HTTPException(401)
+    if '..' in file_path: raise HTTPException(400)
     try:
         data = retrieve_file(file_path)
     except FileNotFoundError:
-        raise HTTPException(404, "File not found")
-    except Exception as e:
-        logger.error(f"File download failed: {file_path} - {e}")
-        raise HTTPException(500, "Failed to retrieve file")
-    
+        raise HTTPException(404)
     import mimetypes
-    media_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-    filename = os.path.basename(file_path)
-    
-    return Response(
-        content=data,
-        media_type=media_type,
-        headers={
-            "Content-Disposition": f'inline; filename="{filename}"',
-            "Cache-Control": "private, max-age=3600"
-        }
-    )
+    mt = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
+    return Response(content=data, media_type=mt)
 
 # ------------------------------------------------------------------------
-# WebSocket Connection Manager
+# WebSocket Manager
 # ------------------------------------------------------------------------
 class ConnectionManager:
-    """Manage WebSocket connections with per-user tracking."""
-    
     def __init__(self):
-        self.active_connections: Dict[int, Dict[str, Any]] = {}  # user_id -> {username, connections: set}
-    
-    async def connect(self, websocket: WebSocket, user_id: int, username: str):
-        """Add a new connection."""
-        await websocket.accept()
-        
-        if user_id not in self.active_connections:
-            self.active_connections[user_id] = {
-                "username": username,
-                "connections": set()
-            }
-        
-        self.active_connections[user_id]["connections"].add(websocket)
-    
-    def disconnect(self, websocket: WebSocket, user_id: int):
-        """Remove a connection."""
-        if user_id in self.active_connections:
-            self.active_connections[user_id]["connections"].discard(websocket)
-            if not self.active_connections[user_id]["connections"]:
-                del self.active_connections[user_id]
-    
-    async def broadcast(self, message: dict, exclude_user_id: Optional[int] = None):
-        """Broadcast message to all connected users."""
-        for uid, info in self.active_connections.items():
-            if uid == exclude_user_id:
-                continue
+        self.active: Dict[int, Dict[str, Any]] = {}
+
+    async def connect(self, ws: WebSocket, user_id: int, username: str):
+        await ws.accept()
+        if user_id not in self.active:
+            self.active[user_id] = {"username": username, "connections": set()}
+        self.active[user_id]["connections"].add(ws)
+
+    def disconnect(self, ws: WebSocket, user_id: int):
+        if user_id in self.active:
+            self.active[user_id]["connections"].discard(ws)
+            if not self.active[user_id]["connections"]:
+                del self.active[user_id]
+
+    async def broadcast(self, msg: dict, exclude_user_id: Optional[int] = None):
+        for uid, info in list(self.active.items()):
+            if uid == exclude_user_id: continue
             for ws in list(info["connections"]):
                 try:
-                    await ws.send_json(message)
+                    await ws.send_json(msg)
                 except Exception:
                     self.disconnect(ws, uid)
-    
-    async def send_to_user(self, user_id: int, message: dict):
-        """Send message to a specific user."""
-        if user_id in self.active_connections:
-            for ws in list(self.active_connections[user_id]["connections"]):
+
+    async def send_to_user(self, user_id: int, msg: dict):
+        if user_id in self.active:
+            for ws in list(self.active[user_id]["connections"]):
                 try:
-                    await ws.send_json(message)
+                    await ws.send_json(msg)
                 except Exception:
                     self.disconnect(ws, user_id)
-    
+
     def get_online_users(self) -> List[Dict]:
-        """Get list of online users with their info."""
-        users = []
-        for uid, info in self.active_connections.items():
-            users.append({
-                "id": uid,
-                "username": info["username"]
-            })
-        return users
-    
+        return [{"id": uid, "username": info["username"]} for uid, info in self.active.items()]
+
     def is_online(self, user_id: int) -> bool:
-        """Check if a user is online."""
-        return user_id in self.active_connections
+        return user_id in self.active
 
 manager = ConnectionManager()
 
 # ------------------------------------------------------------------------
-# WebSocket Endpoint
+# WebSocket endpoint
 # ------------------------------------------------------------------------
 @app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket, token: str = Query(...)):
-    """WebSocket endpoint for real-time messaging."""
-    # Authenticate
+async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
     user = await authenticate_user(token)
     if not user:
-        await websocket.close(code=4001, reason="Authentication failed")
+        await ws.close(code=4001, reason="Invalid token")
         return
-    
-    user_id = user['id']
+
+    uid = user['id']
     username = user['username']
     display_name = user['display_name']
     avatar_path = user['avatar_path']
-    
-    # Connect
-    await manager.connect(websocket, user_id, username)
-    
+
+    await manager.connect(ws, uid, username)
+
     # Send initial data
-    online_users = manager.get_online_users()
-    await websocket.send_json({
+    await ws.send_json({
         "type": "connection_established",
-        "user_id": user_id,
-        "username": username,
-        "display_name": display_name,
-        "avatar_path": avatar_path,
-        "online_users": online_users
+        "user_id": uid, "username": username, "display_name": display_name,
+        "avatar_path": avatar_path, "online_users": manager.get_online_users()
     })
-    
-    # Broadcast user online status
     await manager.broadcast({
-        "type": "user_status",
-        "user_id": user_id,
-        "username": username,
-        "display_name": display_name,
-        "avatar_path": avatar_path,
-        "status": "online"
-    }, exclude_user_id=user_id)
-    
+        "type": "user_status", "user_id": uid, "username": username,
+        "display_name": display_name, "avatar_path": avatar_path, "status": "online"
+    }, exclude_user_id=uid)
+
     try:
         while True:
-            data = await websocket.receive_json()
-            message_type = data.get("type")
-            
-            # --------------------------------------------------------------------
-            # Send Message
-            # --------------------------------------------------------------------
-            if message_type == "send_message":
+            data = await ws.receive_json()
+            mtype = data.get("type")
+
+            if mtype == "send_message":
                 content = data.get("content", "").strip()
                 reply_to_id = data.get("reply_to_id")
                 file_path = data.get("file_path")
@@ -717,379 +395,157 @@ async def websocket_endpoint(websocket: WebSocket, token: str = Query(...)):
                 file_name = data.get("file_name")
                 file_size = data.get("file_size", 0)
                 client_id = data.get("client_id")
-                
-                # Validate content
+
                 if not content and not file_path:
-                    await websocket.send_json({
-                        "type": "error",
-                        "code": "INVALID_MESSAGE",
-                        "message": "Message must contain text or a file"
-                    })
+                    await ws.send_json({"type": "error", "code": "EMPTY", "message": "Message cannot be empty"})
                     continue
-                
-                # Validate content length
                 if len(content) > 10000:
-                    await websocket.send_json({
-                        "type": "error",
-                        "code": "MESSAGE_TOO_LONG",
-                        "message": "Message must be less than 10,000 characters"
-                    })
+                    await ws.send_json({"type": "error", "code": "TOO_LONG", "message": "Message too long"})
                     continue
-                
-                # Validate reply exists
-                if reply_to_id:
-                    db = await get_db()
-                    try:
-                        async with db.execute(
-                            "SELECT id FROM messages WHERE id = ? AND is_deleted = 0",
-                            (reply_to_id,)
-                        ) as cursor:
-                            if not await cursor.fetchone():
-                                reply_to_id = None
-                    finally:
-                        await db.close()
-                
-                timestamp_ms = int(time.time() * 1000)
-                
-                # Encrypt content
+
                 encrypted = encrypt_message(content)
-                
-                # Store in database
+                ts = int(time.time() * 1000)
+
                 db = await get_db()
                 try:
                     cursor = await db.execute(
-                        """INSERT INTO messages 
-                           (sender_id, encrypted_content, timestamp_ms, reply_to_id, 
-                            file_path, file_type, file_name, file_size)
-                           VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
-                        (user_id, encrypted, timestamp_ms, reply_to_id,
-                         file_path, file_type, file_name, file_size)
+                        "INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name, file_size) VALUES (?,?,?,?,?,?,?,?)",
+                        (uid, encrypted, ts, reply_to_id, file_path, file_type, file_name, file_size)
                     )
-                    message_id = cursor.lastrowid
+                    mid = cursor.lastrowid
                     await db.commit()
                 finally:
                     await db.close()
-                
-                # Build message object
-                message_obj = {
-                    "type": "new_message",
-                    "message": {
-                        "id": message_id,
-                        "sender_id": user_id,
-                        "sender_username": username,
-                        "sender_display_name": display_name,
-                        "sender_avatar_path": avatar_path,
-                        "content": content,
-                        "timestamp_ms": timestamp_ms,
-                        "reply_to_id": reply_to_id,
-                        "file_path": file_path,
-                        "file_type": file_type,
-                        "file_name": file_name,
-                        "file_size": file_size,
-                        "is_edited": False,
-                        "is_deleted": False,
-                        "status": "sent"
-                    }
+
+                msg_obj = {
+                    "id": mid, "sender_id": uid, "sender_username": username, "sender_display_name": display_name,
+                    "sender_avatar_path": avatar_path, "content": content, "timestamp_ms": ts,
+                    "reply_to_id": reply_to_id, "file_path": file_path, "file_type": file_type,
+                    "file_name": file_name, "file_size": file_size, "is_edited": False, "is_deleted": False,
+                    "status": "sent"
                 }
-                
-                # Send to sender with delivered status
-                await websocket.send_json({
-                    **message_obj,
-                    "message": {**message_obj["message"], "status": "delivered"},
-                    "client_id": client_id
-                })
-                
-                # Broadcast to other users
-                await manager.broadcast(message_obj, exclude_user_id=user_id)
-            
-            # --------------------------------------------------------------------
-            # Load Messages (Cursor-based Pagination)
-            # --------------------------------------------------------------------
-            elif message_type == "load_messages":
-                cursor_id = data.get("cursor")  # Message ID to load before
-                limit = min(int(data.get("limit", 50)), 100)
-                
+
+                await ws.send_json({"type": "new_message", "message": {**msg_obj, "status": "delivered"}, "client_id": client_id})
+                await manager.broadcast({"type": "new_message", "message": msg_obj}, exclude_user_id=uid)
+
+            elif mtype == "load_messages":
+                cursor_id = data.get("cursor")
+                limit = min(data.get("limit", 50), 100)
                 db = await get_db()
                 try:
                     if cursor_id:
-                        query = """
-                            SELECT m.*, u.username, u.display_name, u.avatar_path
-                            FROM messages m
-                            JOIN users u ON m.sender_id = u.id
-                            WHERE m.is_deleted = 0 AND m.id < ?
-                            ORDER BY m.id DESC
-                            LIMIT ?
-                        """
-                        params = (cursor_id, limit)
+                        rows = await db.execute(
+                            """SELECT m.*, u.username, u.display_name, u.avatar_path
+                               FROM messages m JOIN users u ON m.sender_id = u.id
+                               WHERE m.is_deleted = 0 AND m.id < ? ORDER BY m.id DESC LIMIT ?""",
+                            (cursor_id, limit)
+                        )
                     else:
-                        query = """
-                            SELECT m.*, u.username, u.display_name, u.avatar_path
-                            FROM messages m
-                            JOIN users u ON m.sender_id = u.id
-                            WHERE m.is_deleted = 0
-                            ORDER BY m.id DESC
-                            LIMIT ?
-                        """
-                        params = (limit,)
-                    
-                    rows = await (await db.execute(query, params)).fetchall()
-                    
-                    messages = []
-                    for row in reversed(rows):  # Reverse to get chronological order
+                        rows = await db.execute(
+                            """SELECT m.*, u.username, u.display_name, u.avatar_path
+                               FROM messages m JOIN users u ON m.sender_id = u.id
+                               WHERE m.is_deleted = 0 ORDER BY m.id DESC LIMIT ?""",
+                            (limit,)
+                        )
+                    msgs = []
+                    for row in reversed(await rows.fetchall()):
                         try:
                             decrypted = decrypt_message(row['encrypted_content'])
                         except Exception:
                             decrypted = "[Decryption Error]"
-                        
-                        # Check read status
-                        read_status = "sent"
-                        async with db.execute(
-                            "SELECT COUNT(*) as count FROM read_receipts WHERE message_id = ?",
-                            (row['id'],)
-                        ) as read_cursor:
-                            read_result = await read_cursor.fetchone()
-                            if read_result and read_result['count'] > 0:
-                                read_status = "read"
-                        
-                        messages.append({
-                            "id": row['id'],
-                            "sender_id": row['sender_id'],
-                            "sender_username": row['username'],
-                            "sender_display_name": row['display_name'],
-                            "sender_avatar_path": row['avatar_path'],
-                            "content": decrypted,
-                            "timestamp_ms": row['timestamp_ms'],
-                            "reply_to_id": row['reply_to_id'],
-                            "file_path": row['file_path'],
-                            "file_type": row['file_type'],
-                            "file_name": row['file_name'],
-                            "file_size": row['file_size'],
-                            "is_edited": bool(row['is_edited']),
-                            "is_deleted": False,
-                            "status": read_status
+                        read_cursor = await db.execute("SELECT COUNT(*) as cnt FROM read_receipts WHERE message_id = ?", (row['id'],))
+                        read_count = (await read_cursor.fetchone())['cnt']
+                        msgs.append({
+                            "id": row['id'], "sender_id": row['sender_id'], "sender_username": row['username'],
+                            "sender_display_name": row['display_name'], "sender_avatar_path": row['avatar_path'],
+                            "content": decrypted, "timestamp_ms": row['timestamp_ms'], "reply_to_id": row['reply_to_id'],
+                            "file_path": row['file_path'], "file_type": row['file_type'], "file_name": row['file_name'],
+                            "file_size": row['file_size'], "is_edited": bool(row['is_edited']), "is_deleted": False,
+                            "status": "read" if read_count > 0 else "sent"
                         })
-                    
-                    # Get the oldest message ID for next page
-                    next_cursor = None
-                    if len(messages) > 0:
-                        next_cursor = messages[0]['id']
-                    
-                    await websocket.send_json({
-                        "type": "messages_loaded",
-                        "messages": messages,
-                        "next_cursor": next_cursor,
-                        "has_more": len(messages) == limit
-                    })
+                    next_cursor = msgs[0]['id'] if msgs else None
+                    await ws.send_json({"type": "messages_loaded", "messages": msgs, "next_cursor": next_cursor, "has_more": len(msgs) == limit})
                 finally:
                     await db.close()
-            
-            # --------------------------------------------------------------------
-            # Edit Message
-            # --------------------------------------------------------------------
-            elif message_type == "edit_message":
-                message_id = data.get("message_id")
+
+            elif mtype == "edit_message":
+                mid = data.get("message_id")
                 new_content = data.get("content", "").strip()
-                
-                if not message_id or not new_content:
-                    await websocket.send_json({
-                        "type": "error",
-                        "code": "INVALID_EDIT",
-                        "message": "Message ID and new content are required"
-                    })
-                    continue
-                
-                if len(new_content) > 10000:
-                    await websocket.send_json({
-                        "type": "error",
-                        "code": "MESSAGE_TOO_LONG",
-                        "message": "Message must be less than 10,000 characters"
-                    })
-                    continue
-                
+                if not mid or not new_content: continue
                 db = await get_db()
                 try:
-                    # Verify ownership
-                    async with db.execute(
-                        "SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0",
-                        (message_id, user_id)
-                    ) as cursor:
-                        if not await cursor.fetchone():
-                            await websocket.send_json({
-                                "type": "error",
-                                "code": "NOT_FOUND",
-                                "message": "Message not found or you don't have permission to edit it"
-                            })
+                    async with db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0", (mid, uid)) as c:
+                        if not await c.fetchone():
+                            await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
                             continue
-                    
-                    new_encrypted = encrypt_message(new_content)
-                    await db.execute(
-                        "UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?",
-                        (new_encrypted, message_id)
-                    )
+                    new_enc = encrypt_message(new_content)
+                    await db.execute("UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?", (new_enc, mid))
                     await db.commit()
-                    
-                    # Broadcast edit
-                    await manager.broadcast({
-                        "type": "message_edited",
-                        "message_id": message_id,
-                        "content": new_content,
-                        "editor_id": user_id
-                    })
+                    await manager.broadcast({"type": "message_edited", "message_id": mid, "content": new_content, "editor_id": uid})
                 finally:
                     await db.close()
-            
-            # --------------------------------------------------------------------
-            # Delete Message
-            # --------------------------------------------------------------------
-            elif message_type == "delete_message":
-                message_id = data.get("message_id")
-                
-                if not message_id:
-                    continue
-                
+
+            elif mtype == "delete_message":
+                mid = data.get("message_id")
+                if not mid: continue
                 db = await get_db()
                 try:
-                    # Verify ownership
-                    async with db.execute(
-                        "SELECT * FROM messages WHERE id = ? AND sender_id = ?",
-                        (message_id, user_id)
-                    ) as cursor:
-                        message = await cursor.fetchone()
-                    
-                    if not message:
-                        await websocket.send_json({
-                            "type": "error",
-                            "code": "NOT_FOUND",
-                            "message": "Message not found or you don't have permission to delete it"
-                        })
+                    async with db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid)) as c:
+                        msg = await c.fetchone()
+                    if not msg:
+                        await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Not found"})
                         continue
-                    
-                    # Delete associated file if exists
-                    if message['file_path']:
-                        try:
-                            delete_file(message['file_path'])
-                        except Exception:
-                            pass
-                    
-                    # Remove message and read receipts
-                    await db.execute("DELETE FROM messages WHERE id = ?", (message_id,))
-                    await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (message_id,))
+                    if msg['file_path']:
+                        try: delete_file(msg['file_path'])
+                        except: pass
+                    await db.execute("DELETE FROM messages WHERE id = ?", (mid,))
+                    await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
                     await db.commit()
-                    
-                    # Broadcast deletion
-                    await manager.broadcast({
-                        "type": "message_deleted",
-                        "message_id": message_id
-                    })
+                    await manager.broadcast({"type": "message_deleted", "message_id": mid})
                 finally:
                     await db.close()
-            
-            # --------------------------------------------------------------------
-            # Typing Indicator
-            # --------------------------------------------------------------------
-            elif message_type == "typing":
-                is_typing = bool(data.get("is_typing", False))
-                
-                await manager.broadcast({
-                    "type": "typing_indicator",
-                    "user_id": user_id,
-                    "username": username,
-                    "display_name": display_name,
-                    "is_typing": is_typing
-                }, exclude_user_id=user_id)
-            
-            # --------------------------------------------------------------------
-            # Mark as Read
-            # --------------------------------------------------------------------
-            elif message_type == "mark_read":
-                up_to_message_id = data.get("up_to_message_id")
-                
-                if not up_to_message_id:
-                    continue
-                
+
+            elif mtype == "typing":
+                await manager.broadcast({"type": "typing_indicator", "user_id": uid, "username": username, "display_name": display_name, "is_typing": bool(data.get("is_typing", False))}, exclude_user_id=uid)
+
+            elif mtype == "mark_read":
+                up_to = data.get("up_to_message_id")
+                if not up_to: continue
                 db = await get_db()
                 try:
-                    # Get messages that haven't been read by this user
                     rows = await db.execute(
-                        """SELECT m.id, m.sender_id 
-                           FROM messages m 
-                           WHERE m.id <= ? 
-                           AND m.sender_id != ? 
-                           AND m.is_deleted = 0 
-                           AND m.id > (
-                               SELECT COALESCE(MAX(message_id), 0) 
-                               FROM read_receipts 
-                               WHERE user_id = ?
-                           )""",
-                        (up_to_message_id, user_id, user_id)
+                        """SELECT m.id, m.sender_id FROM messages m
+                           WHERE m.id <= ? AND m.sender_id != ? AND m.is_deleted = 0
+                           AND m.id NOT IN (SELECT message_id FROM read_receipts WHERE user_id = ?)""",
+                        (up_to, uid, uid)
                     )
-                    
                     new_reads = []
-                    async for row in rows:
-                        await db.execute(
-                            "INSERT OR IGNORE INTO read_receipts (user_id, message_id, read_at) VALUES (?, ?, ?)",
-                            (user_id, row['id'], int(time.time()))
-                        )
-                        new_reads.append({
-                            "message_id": row['id'],
-                            "sender_id": row['sender_id']
-                        })
-                    
+                    async for r in rows:
+                        await db.execute("INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?,?)", (uid, r['id']))
+                        new_reads.append(r)
                     await db.commit()
-                    
-                    # Notify message senders
-                    for read_info in new_reads:
-                        await manager.send_to_user(read_info['sender_id'], {
-                            "type": "message_read",
-                            "message_id": read_info['message_id'],
-                            "reader_id": user_id,
-                            "reader_username": username
-                        })
+                    for r in new_reads:
+                        await manager.send_to_user(r['sender_id'], {"type": "message_read", "message_id": r['id'], "reader_username": username})
                 finally:
                     await db.close()
-            
-            # --------------------------------------------------------------------
-            # Get Online Users
-            # --------------------------------------------------------------------
-            elif message_type == "get_online_users":
-                await websocket.send_json({
-                    "type": "online_users",
-                    "users": manager.get_online_users()
-                })
-    
+
+            elif mtype == "get_online_users":
+                await ws.send_json({"type": "online_users", "users": manager.get_online_users()})
+
     except WebSocketDisconnect:
         pass
     except Exception as e:
-        logger.error(f"WebSocket error for {username}: {e}")
+        logger.error(f"WS error for {username}: {e}")
     finally:
-        # Disconnect
-        manager.disconnect(websocket, user_id)
-        
-        # Update status in database
+        manager.disconnect(ws, uid)
         db = await get_db()
         try:
-            await db.execute(
-                "UPDATE users SET status = 'offline', last_seen = ? WHERE id = ?",
-                (int(time.time()), user_id)
-            )
+            await db.execute("UPDATE users SET status = 'offline', last_seen = ? WHERE id = ?", (int(time.time()), uid))
             await db.commit()
         finally:
             await db.close()
-        
-        # Broadcast offline status
-        await manager.broadcast({
-            "type": "user_status",
-            "user_id": user_id,
-            "username": username,
-            "display_name": display_name,
-            "avatar_path": avatar_path,
-            "status": "offline"
-        })
+        await manager.broadcast({"type": "user_status", "user_id": uid, "username": username, "display_name": display_name, "status": "offline"})
 
-# ------------------------------------------------------------------------
-# Root endpoint
-# ------------------------------------------------------------------------
 @app.get("/")
 async def root():
-    """Serve the main application."""
     return FileResponse("static/index.html")