Fix files, delete, edit in DMs

backend/routes/messaging.py

@@ -4,14 +4,13 @@ from pathlib import Path
 import os
 import re
 import uuid
-from typing import Iterable
 from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File, Form
 from fastapi.responses import FileResponse
 from fastapi.security import HTTPAuthorizationCredentials
 from sqlalchemy.orm import Session
 from dependencies import get_current_user, get_db
 from constants import OWNER_USERNAME
-from models import Message, SendMessageRequest, EditMessageRequest, User, DMEnvelope, MessageFile
+from models import Message, SendMessageRequest, EditMessageRequest, User, DMEnvelope, MessageFile, DMFile
 from push_service import push_service
 from PIL import Image
 import io
@@ -150,11 +149,7 @@ async def send_message(
 
             mf = MessageFile(
                 message_id=new_message.id,
-                path=str(out_path),
-                encrypted=False,
-                filename=original_name,
-                content_type=up.content_type,
-                size=len(content),
+                path=str(out_path)
             )
             db.add(mf)
         db.commit()
@@ -203,7 +198,6 @@ async def dm_send(
     files: list[UploadFile] = File(default=[]),
     fileNames: str | None = Form(default=None),  # JSON array of filenames corresponding to files
 ):
-    import json
     if dm_payload and payload is None:
         try:
             payload = json.loads(dm_payload)
@@ -226,6 +220,7 @@ async def dm_send(
         salt_b64=payload["salt"],
         iv2_b64=payload["iv2"],
         wrapped_mk_b64=payload["wrappedMk"],
+        reply_to_id=payload.get("replyToId") if isinstance(payload.get("replyToId"), int) else None,
     )
     db.add(env)
     db.commit()
@@ -235,12 +230,12 @@ async def dm_send(
     if files:
         # Validate total size
         total_size = 0
-        for up in files:
-            if hasattr(up, "size") and up.size is not None:
-                total_size += int(up.size)
+        for file in files:
+            if hasattr(file, "size") and file.size is not None:
+                total_size += int(file.size)
             else:
-                data = await up.read()
-                up.file.seek(0)
+                data = await file.read()
+                file.file.seek(0)
                 total_size += len(data)
             if total_size > MAX_TOTAL_SIZE:
                 raise HTTPException(status_code=400, detail="Total attachments size exceeds 4GB")
@@ -254,21 +249,31 @@ async def dm_send(
             except Exception:
                 names = []
 
-        for idx, up in enumerate(files):
-            provided = names[idx] if idx < len(names) else None
+        for i, file in enumerate(files):
+            provided = names[i] if i < len(names) else None
             # Sanitize provided name to avoid path traversal
             if provided and not re.match(r"^[A-Za-z0-9._-]{1,200}$", provided):
                 provided = None
-            original_name = provided or Path(up.filename or "file").name
+            original_name = provided or Path(file.filename or "file").name
             # Save using provided/original name to allow client to reference path directly
             safe_name = original_name
-            out_path = FILES_ENCRYPTED_DIR / safe_name
+            out_name = f"{current_user.id}_{env.recipient_id}_{env.id}_{safe_name}"
+            out_path = FILES_ENCRYPTED_DIR / out_name
 
-            content = await up.read()
+            content = await file.read()
             with open(out_path, "wb") as f:
                 f.write(content)
 
-            # We do not store linkage to public messages for DMs; paths will be referenced inside encrypted JSON
+            # Save DM file record
+            df = DMFile(
+                message_id=env.id,
+                sender_id=current_user.id,
+                recipient_id=env.recipient_id,
+                path=f"/api/uploads/files/encrypted/{out_name}",
+                name=safe_name
+            )
+            db.add(df)
+        db.commit()
 
     # Send push notification for DM
     try:
@@ -278,7 +283,6 @@ async def dm_send(
 
     # Realtime notify both users for HTTP requests
     try:
-        from .messaging import messagingManager  # self import
         payload_ws = {
             "type": "dmNew",
             "data": {
@@ -291,6 +295,7 @@ async def dm_send(
                 "iv2": env.iv2_b64,
                 "wrappedMk": env.wrapped_mk_b64,
                 "timestamp": env.timestamp.isoformat(),
+                "replyToId": env.reply_to_id,
             }
         }
         await messagingManager.send_to_user(env.recipient_id, payload_ws)
@@ -300,13 +305,7 @@ async def dm_send(
 
     return {"status": "ok", "id": env.id}
 
-
-@router.get("/dm/fetch")
-async def dm_fetch(since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
-    q = db.query(DMEnvelope).filter(DMEnvelope.recipient_id == current_user.id)
-    if since:
-        q = q.filter(DMEnvelope.id > since)
-    envs = q.order_by(DMEnvelope.id.asc()).all()
+def convert_envelopes(envs: list[DMEnvelope]):
     return {
         "status": "ok",
         "messages": [
@@ -320,15 +319,23 @@ async def dm_fetch(since: int | None = None, current_user: User = Depends(get_cu
                 "iv2": e.iv2_b64,
                 "wrappedMk": e.wrapped_mk_b64,
                 "timestamp": e.timestamp.isoformat(),
+                "files": [{"name": file.name, "path": file.path, "id": file.id} for file in e.files]
             }
             for e in envs
         ]
     }
 
+@router.get("/dm/fetch")
+async def dm_fetch(since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
+    q = db.query(DMEnvelope).filter(DMEnvelope.recipient_id == current_user.id)
+    if since:
+        q = q.filter(DMEnvelope.id > since)
+    return convert_envelopes(q.order_by(DMEnvelope.id.asc()).all())
+
 
 @router.get("/dm/history/{other_user_id}")
 async def dm_history(other_user_id: int, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
-    envs = (
+    return convert_envelopes(
         db.query(DMEnvelope)
         .filter(
             ((DMEnvelope.sender_id == current_user.id) & (DMEnvelope.recipient_id == other_user_id))
@@ -337,23 +344,6 @@ async def dm_history(other_user_id: int, current_user: User = Depends(get_curren
         .order_by(DMEnvelope.id.asc())
         .all()
     )
-    return {
-        "status": "ok",
-        "messages": [
-            {
-                "id": e.id,
-                "senderId": e.sender_id,
-                "recipientId": e.recipient_id,
-                "iv": e.iv_b64,
-                "ciphertext": e.ciphertext_b64,
-                "salt": e.salt_b64,
-                "iv2": e.iv2_b64,
-                "wrappedMk": e.wrapped_mk_b64,
-                "timestamp": e.timestamp.isoformat(),
-            }
-            for e in envs
-        ]
-    }
 
 
 @router.put("/edit_message/{message_id}")
@@ -503,6 +493,7 @@ class MessaggingSocketManager:
                         salt_b64=payload["salt"],
                         iv2_b64=payload["iv2"],
                         wrapped_mk_b64=payload["wrappedMk"],
+                        reply_to_id=payload.get("replyToId") if isinstance(payload.get("replyToId"), int) else None,
                     )
                     db.add(env)
                     db.commit()
@@ -520,6 +511,7 @@ class MessaggingSocketManager:
                             "iv2": env.iv2_b64,
                             "wrappedMk": env.wrapped_mk_b64,
                             "timestamp": env.timestamp.isoformat(),
+                            "replyToId": env.reply_to_id,
                         }
                     }
 
@@ -552,6 +544,76 @@ class MessaggingSocketManager:
                     await websocket.send_json({"type": type, "data": response})
                 except HTTPException as e:
                     await self.send_error(websocket, type, e)
+            elif type == "dmEdit":
+                try:
+                    current_user = get_current_user_inner()
+                    if not current_user:
+                        raise HTTPException(401)
+
+                    payload = data["data"]
+                    env_id = int(payload["id"])
+                    env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first()
+                    if not env:
+                        raise HTTPException(status_code=404, detail="DM not found")
+                    if env.sender_id != current_user.id:
+                        raise HTTPException(status_code=403, detail="You can only edit your own messages")
+
+                    # Replace ciphertext and iv
+                    env.iv_b64 = payload["iv"]
+                    env.ciphertext_b64 = payload["ciphertext"]
+                    env.iv2_b64 = payload["iv2"]
+                    env.wrapped_mk_b64 = payload["wrappedMk"]
+                    env.salt_b64 = payload["salt"]
+                    db.commit()
+                    db.refresh(env)
+
+                    payload_ws = {
+                        "type": "dmEdited",
+                        "data": {
+                            "id": env.id,
+                            "iv": env.iv_b64,
+                            "ciphertext": env.ciphertext_b64,
+                            "iv2": env.iv2_b64,
+                            "wrappedMk": env.wrapped_mk_b64,
+                            "salt": env.salt_b64,
+                            "timestamp": env.timestamp.isoformat(),
+                        }
+                    }
+                    await self.send_to_user(env.recipient_id, payload_ws)
+                    await self.send_to_user(env.sender_id, payload_ws)
+                    await websocket.send_json({"type": type, "data": {"status": "ok", "id": env.id}})
+                except HTTPException as e:
+                    await self.send_error(websocket, type, e)
+            elif type == "dmDelete":
+                try:
+                    current_user = get_current_user_inner()
+                    if not current_user:
+                        raise HTTPException(401)
+
+                    payload = data["data"]
+                    env_id = int(payload["id"])
+                    env: DMEnvelope | None = db.query(DMEnvelope).filter(DMEnvelope.id == env_id).first()
+                    if not env:
+                        raise HTTPException(status_code=404, detail="DM not found")
+                    if env.sender_id != current_user.id:
+                        raise HTTPException(status_code=403, detail="You can only delete your own messages")
+
+                    db.delete(env)
+                    db.commit()
+
+                    payload_ws = {
+                        "type": "dmDeleted",
+                        "data": {
+                            "id": env_id,
+                            "senderId": current_user.id,
+                            "recipientId": payload.get("recipientId")
+                        }
+                    }
+                    await self.send_to_user(env.recipient_id, payload_ws)
+                    await websocket.send_json({"type": type, "data": {"status": "ok", "id": env_id}})
+                    await self.send_to_user(env.sender_id, payload_ws)
+                except HTTPException as e:
+                    await self.send_error(websocket, type, e)
             elif type == "deleteMessage":
                 try:
                     current_user = get_current_user_inner()
@@ -609,7 +671,7 @@ async def chat_websocket(
 
 
 # File serving endpoints
-@router.get("/files/normal/{filename}")
+@router.get("/uploads/files/normal/{filename}")
 async def get_file_normal(filename: str):
     if not re.match(r"^[A-Za-z0-9._-]+$", filename):
         raise HTTPException(status_code=400, detail="Invalid file name")
@@ -619,11 +681,22 @@ async def get_file_normal(filename: str):
     return FileResponse(str(path))
 
 
-@router.get("/files/encrypted/{filename}")
-async def get_file_encrypted(filename: str):
+@router.get("/uploads/files/encrypted/{filename}")
+async def get_file_encrypted(filename: str, current_user: User = Depends(get_current_user)):
     if not re.match(r"^[A-Za-z0-9._-]+$", filename):
         raise HTTPException(status_code=400, detail="Invalid file name")
     path = FILES_ENCRYPTED_DIR / filename
     if not path.exists():
         raise HTTPException(status_code=404, detail="File not found")
+
+    match = re.match(r"^(\d+)_(\d+)_(\d+)_.*$", path.resolve().name)
+    if match:
+        sender_id = int(match.group(1))
+        recipient_id = int(match.group(2))
+
+        if not current_user.id in [sender_id, recipient_id]:
+            raise HTTPException(403)
+    else:
+        raise HTTPException(500)
+
     return FileResponse(str(path))