Update app.py

app.py

@@ -105,7 +105,7 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # ------------------------------------------------------------------------
-# Database helper (new connection each time)
+# Database helper
 # ------------------------------------------------------------------------
 async def get_db() -> aiosqlite.Connection:
     conn = await aiosqlite.connect(DATABASE_URL)
@@ -115,7 +115,7 @@ async def get_db() -> aiosqlite.Connection:
     return conn
 
 # ------------------------------------------------------------------------
-# Encryption helpers
+# Encryption helpers (AES-256-GCM)
 # ------------------------------------------------------------------------
 def encrypt_message(plaintext: str) -> str:
     aesgcm = AESGCM(MESSAGE_KEY)
@@ -133,7 +133,7 @@ def decrypt_message(encrypted_b64: str) -> str:
         return "[Decryption Error]"
 
 # ------------------------------------------------------------------------
-# Password hashing
+# Password hashing (PBKDF2)
 # ------------------------------------------------------------------------
 def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
     if salt is None:
@@ -175,7 +175,7 @@ async def health_check():
     return {"status": "healthy", "timestamp": int(time.time()), "version": "1.0.0"}
 
 @app.post("/api/auth/signup")
-async def signup(username: str = Query(...), password: str = Query(...), display_name: str = Query(None)):
+async def signup(username: str = Query(..., min_length=3, max_length=30), password: str = Query(..., min_length=6), display_name: str = Query(None, max_length=50)):
     if not username.isalnum() and '_' not in username:
         raise HTTPException(400, "Invalid username")
     db = await get_db()
@@ -191,7 +191,7 @@ async def signup(username: str = Query(...), password: str = Query(...), display
             (username.lower(), display, pwd_hash, salt, token)
         )
         await db.commit()
-        return {"token": token, "user": {"username": username, "display_name": display, "avatar_path": None, "status": "online"}}
+        return {"token": token, "user": {"id": None, "username": username, "display_name": display, "avatar_path": None, "status": "online"}}
     finally:
         await db.close()
 
@@ -206,7 +206,16 @@ async def login(username: str = Query(...), password: str = Query(...)):
         token = generate_token()
         await db.execute("UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?", (token, int(time.time()), user['id']))
         await db.commit()
-        return {"token": token, "user": {"id": user['id'], "username": user['username'], "display_name": user['display_name'], "avatar_path": user['avatar_path'], "status": "online"}}
+        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()
 
@@ -272,15 +281,24 @@ async def upload_avatar(file: UploadFile = File(...), token: str = Header(..., a
 upload_sessions: Dict[str, dict] = {}
 
 @app.post("/api/upload/chunk")
-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")):
+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)
     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]
     if session['user_id'] != user['id']: raise HTTPException(403)
     chunk_data = await file.read()
@@ -374,12 +392,21 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
     # Send initial data
     await ws.send_json({
         "type": "connection_established",
-        "user_id": uid, "username": username, "display_name": display_name,
-        "avatar_path": avatar_path, "online_users": manager.get_online_users()
+        "user_id": uid,
+        "username": username,
+        "display_name": display_name,
+        "avatar_path": avatar_path,
+        "online_users": manager.get_online_users()
     })
+
+    # Broadcast new online status
     await manager.broadcast({
-        "type": "user_status", "user_id": uid, "username": username,
-        "display_name": display_name, "avatar_path": avatar_path, "status": "online"
+        "type": "user_status",
+        "user_id": uid,
+        "username": username,
+        "display_name": display_name,
+        "avatar_path": avatar_path,
+        "status": "online"
     }, exclude_user_id=uid)
 
     try:
@@ -387,6 +414,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
             data = await ws.receive_json()
             mtype = data.get("type")
 
+            # -------------------- SEND MESSAGE --------------------
             if mtype == "send_message":
                 content = data.get("content", "").strip()
                 reply_to_id = data.get("reply_to_id")
@@ -418,16 +446,37 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     await db.close()
 
                 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,
+                    "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"
                 }
 
-                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)
+                # Confirm to sender with delivered status
+                await ws.send_json({
+                    "type": "new_message",
+                    "message": {**msg_obj, "status": "delivered"},
+                    "client_id": client_id
+                })
+
+                # Broadcast to others
+                await manager.broadcast({
+                    "type": "new_message",
+                    "message": msg_obj
+                }, exclude_user_id=uid)
 
+            # -------------------- LOAD MESSAGES --------------------
             elif mtype == "load_messages":
                 cursor_id = data.get("cursor")
                 limit = min(data.get("limit", 50), 100)
@@ -451,66 +500,120 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     for row in reversed(await rows.fetchall()):
                         try:
                             decrypted = decrypt_message(row['encrypted_content'])
-                        except Exception:
+                        except:
                             decrypted = "[Decryption Error]"
-                        read_cursor = await db.execute("SELECT COUNT(*) as cnt FROM read_receipts WHERE message_id = ?", (row['id'],))
+
+                        # Determine 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,
+                            "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"
                         })
+
                     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})
+                    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 mtype == "edit_message":
                 mid = data.get("message_id")
                 new_content = data.get("content", "").strip()
-                if not mid or not new_content: continue
+                if not mid or not new_content:
+                    continue
                 db = await get_db()
                 try:
-                    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():
+                    async with db.execute(
+                        "SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0",
+                        (mid, uid)
+                    ) as cursor:
+                        if not await cursor.fetchone():
                             await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
                             continue
                     new_enc = encrypt_message(new_content)
-                    await db.execute("UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?", (new_enc, mid))
+                    await db.execute(
+                        "UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?",
+                        (new_enc, mid)
+                    )
                     await db.commit()
-                    await manager.broadcast({"type": "message_edited", "message_id": mid, "content": new_content, "editor_id": uid})
+                    await manager.broadcast({
+                        "type": "message_edited",
+                        "message_id": mid,
+                        "content": new_content,
+                        "editor_id": uid
+                    })
                 finally:
                     await db.close()
 
+            # -------------------- DELETE MESSAGE --------------------
             elif mtype == "delete_message":
                 mid = data.get("message_id")
-                if not mid: continue
+                if not mid:
+                    continue
                 db = await get_db()
                 try:
-                    async with db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid)) as c:
-                        msg = await c.fetchone()
+                    async with db.execute(
+                        "SELECT * FROM messages WHERE id = ? AND sender_id = ?",
+                        (mid, uid)
+                    ) as cursor:
+                        msg = await cursor.fetchone()
                     if not msg:
-                        await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Not found"})
+                        await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
                         continue
                     if msg['file_path']:
-                        try: delete_file(msg['file_path'])
-                        except: pass
+                        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()
-                    await manager.broadcast({"type": "message_deleted", "message_id": mid})
+                    await manager.broadcast({
+                        "type": "message_deleted",
+                        "message_id": mid
+                    })
                 finally:
                     await db.close()
 
+            # -------------------- TYPING --------------------
             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)
-
+                is_typing = bool(data.get("is_typing", False))
+                await manager.broadcast({
+                    "type": "typing_indicator",
+                    "user_id": uid,
+                    "username": username,
+                    "display_name": display_name,
+                    "is_typing": is_typing
+                }, exclude_user_id=uid)
+
+            # -------------------- MARK READ --------------------
             elif mtype == "mark_read":
                 up_to = data.get("up_to_message_id")
-                if not up_to: continue
+                if not up_to:
+                    continue
                 db = await get_db()
                 try:
                     rows = await db.execute(
@@ -521,16 +624,27 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     )
                     new_reads = []
                     async for r in rows:
-                        await db.execute("INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?,?)", (uid, r['id']))
+                        await db.execute(
+                            "INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?,?)",
+                            (uid, r['id'])
+                        )
                         new_reads.append(r)
                     await db.commit()
                     for r in new_reads:
-                        await manager.send_to_user(r['sender_id'], {"type": "message_read", "message_id": r['id'], "reader_username": username})
+                        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 mtype == "get_online_users":
-                await ws.send_json({"type": "online_users", "users": manager.get_online_users()})
+                await ws.send_json({
+                    "type": "online_users",
+                    "users": manager.get_online_users()
+                })
 
     except WebSocketDisconnect:
         pass
@@ -540,11 +654,21 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
         manager.disconnect(ws, uid)
         db = await get_db()
         try:
-            await db.execute("UPDATE users SET status = 'offline', last_seen = ? WHERE id = ?", (int(time.time()), uid))
+            await db.execute(
+                "UPDATE users SET status = 'offline', last_seen = ? WHERE id = ?",
+                (int(time.time()), uid)
+            )
             await db.commit()
         finally:
             await db.close()
-        await manager.broadcast({"type": "user_status", "user_id": uid, "username": username, "display_name": display_name, "status": "offline"})
+        await manager.broadcast({
+            "type": "user_status",
+            "user_id": uid,
+            "username": username,
+            "display_name": display_name,
+            "avatar_path": avatar_path,
+            "status": "offline"
+        })
 
 @app.get("/")
 async def root():