siftlode/backend/app/routes/messages.py

"""Direct messaging (notification module, phase 2): end-to-end encrypted user-to-user chat
plus server-authored system messages, with live WebSocket delivery.

E2EE model: every private message is encrypted IN THE BROWSER under a key both parties derive
from ECDH(their private key, the other's public key). The server only ever stores ciphertext +
iv and acts as a key directory (public keys) and an opaque blob store (each user's private key,
wrapped client-side with a passphrase the server never sees). So not even an admin can read a
user-to-user conversation. `kind="system"` messages (the welcome; future announcements) are
plain server boilerplate — non-private, readable without any key setup, shown as a "Siftlode"
conversation.

Every endpoint is gated by `require_human` (the shared demo account can't message). Metadata
(who messaged whom, when, read state) is NOT encrypted — standard for E2EE and needed for
routing and unread counts.
"""
from datetime import datetime, timezone

from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from sqlalchemy import and_, case, func, or_, select
from sqlalchemy.orm import Session

from app.auth import require_human, resolved_user_id
from app.db import SessionLocal, get_db
from app.models import Message, MessageKey, User
from app.ratelimit import RateLimiter
from app.realtime import manager

router = APIRouter(prefix="/api/messages", tags=["messages"])

MAX_CIPHERTEXT = 8000  # base64 of a generous plaintext + AES-GCM overhead
SYSTEM_PARTNER_ID = 0  # synthetic conversation partner for system messages
SYSTEM_NAME = "Siftlode"
_send_limiter = RateLimiter(max_events=30, window_seconds=60)

# The welcome is rendered server-side at creation time, so it carries its own translations
# (one per supported UI language). Kept short and non-private (every new user gets the same text).
WELCOME = {
    "en": (
        "Welcome to Siftlode! 👋\n\n"
        "This is your private message inbox. Conversations with other members are end-to-end "
        "encrypted — only you and the person you're chatting with can read them, not even an "
        "admin. To start messaging someone, set up secure messaging with a passphrase first.\n\n"
        "Enjoy your feed!"
    ),
    "hu": (
        "Üdv a Siftlode-on! 👋\n\n"
        "Ez a privát üzenet-postafiókod. A többi taggal folytatott beszélgetések végpontig "
        "titkosítottak — csak te és a beszélgetőpartnered olvashatja őket, még az admin sem. "
        "Ha üzenni szeretnél valakinek, előbb állítsd be a biztonságos üzenetküldést egy "
        "jelszóval.\n\n"
        "Jó böngészést!"
    ),
    "de": (
        "Willkommen bei Siftlode! 👋\n\n"
        "Das ist dein privater Nachrichten-Posteingang. Unterhaltungen mit anderen Mitgliedern "
        "sind Ende-zu-Ende-verschlüsselt — nur du und dein Gegenüber könnt sie lesen, nicht "
        "einmal ein Admin. Um jemandem zu schreiben, richte zuerst die sichere "
        "Nachrichtenübermittlung mit einem Passwort ein.\n\n"
        "Viel Spaß mit deinem Feed!"
    ),
}


class KeySetupIn(BaseModel):
    public_key: str
    wrapped_private_key: str
    salt: str
    wrap_iv: str
    key_check: str


class SendIn(BaseModel):
    recipient_id: int
    ciphertext: str
    iv: str


def _user_label(u: User) -> str:
    return u.display_name or u.email.split("@")[0]


def _serialize_user(u: User) -> dict:
    return {"id": u.id, "name": _user_label(u), "avatar_url": u.avatar_url}


def _system_partner() -> dict:
    return {"id": SYSTEM_PARTNER_ID, "name": SYSTEM_NAME, "avatar_url": None}


def _serialize_msg(m: Message) -> dict:
    return {
        "id": m.id,
        "kind": m.kind,
        "sender_id": m.sender_id,
        "recipient_id": m.recipient_id,
        "body": m.body,  # only set for system messages
        "ciphertext": m.ciphertext,
        "iv": m.iv,
        "read_at": m.read_at.isoformat() if m.read_at else None,
        "created_at": m.created_at.isoformat() if m.created_at else None,
    }


def _messageable() -> list:
    """SQL clauses for "a real, active, non-suspended human" — for directory/recipient queries."""
    return [
        User.is_demo.is_(False),
        User.is_active.is_(True),
        User.is_suspended.is_(False),
    ]


def is_messageable_user(u: User | None) -> bool:
    """Python mirror of `_messageable()` for an already-loaded row (WS auth, send recipient)."""
    return bool(u and not u.is_demo and u.is_active and not u.is_suspended)


def ensure_welcome(db: Session, user: User) -> None:
    """Create the one-time system welcome message for a (human) user the first time it's needed.
    Idempotent: guarded on whether any system message already exists for them."""
    if user.is_demo:
        return
    exists = db.scalar(
        select(Message.id)
        .where(Message.recipient_id == user.id, Message.kind == "system")
        .limit(1)
    )
    if exists:
        return
    lang = (user.preferences or {}).get("language")
    body = WELCOME.get(lang, WELCOME["en"])
    m = Message(kind="system", sender_id=None, recipient_id=user.id, body=body)
    db.add(m)
    db.commit()
    db.refresh(m)
    manager.push(user.id, {"type": "message", "message": _serialize_msg(m)})


# --- E2EE key directory / blob store -----------------------------------------

@router.get("/keys/me")
def my_key(user: User = Depends(require_human), db: Session = Depends(get_db)) -> dict:
    """My own key bundle. `configured` says whether secure messaging is set up; if so the
    wrapped blob + params are returned so this device can unlock with the passphrase."""
    k = db.get(MessageKey, user.id)
    if k is None:
        return {"configured": False}
    return {
        "configured": True,
        "public_key": k.public_key,
        "wrapped_private_key": k.wrapped_private_key,
        "salt": k.salt,
        "wrap_iv": k.wrap_iv,
        "key_check": k.key_check,
    }


@router.post("/keys")
def setup_keys(
    payload: KeySetupIn,
    user: User = Depends(require_human),
    db: Session = Depends(get_db),
) -> dict:
    """Store this user's public key + passphrase-wrapped private key. Set-once: changing the
    key would orphan every existing conversation key, so re-setup is rejected here."""
    if db.get(MessageKey, user.id) is not None:
        raise HTTPException(status_code=409, detail="Secure messaging is already set up.")
    db.add(
        MessageKey(
            user_id=user.id,
            public_key=payload.public_key,
            wrapped_private_key=payload.wrapped_private_key,
            salt=payload.salt,
            wrap_iv=payload.wrap_iv,
            key_check=payload.key_check,
        )
    )
    db.commit()
    return {"configured": True}


@router.get("/keys/{user_id}")
def public_key(
    user_id: int, user: User = Depends(require_human), db: Session = Depends(get_db)
) -> dict:
    """A member's public key, for deriving the pairwise conversation key."""
    k = db.get(MessageKey, user_id)
    if k is None:
        raise HTTPException(status_code=404, detail="That member hasn't set up messaging yet.")
    return {"user_id": user_id, "public_key": k.public_key}


# --- directory + conversations -----------------------------------------------

@router.get("/users")
def list_recipients(
    user: User = Depends(require_human), db: Session = Depends(get_db)
) -> list[dict]:
    """Member directory for the recipient picker: messageable members (except me) who have set
    up secure messaging (so a conversation key can be derived)."""
    rows = (
        db.execute(
            select(User)
            .join(MessageKey, MessageKey.user_id == User.id)
            .where(User.id != user.id, *_messageable())
            .order_by(func.lower(func.coalesce(User.display_name, User.email)))
        )
        .scalars()
        .all()
    )
    return [_serialize_user(u) for u in rows]


@router.get("/conversations")
def list_conversations(
    user: User = Depends(require_human), db: Session = Depends(get_db)
) -> dict:
    """One entry per conversation: each user partner (E2EE) plus the system "Siftlode" thread,
    with the latest message + my unread count, newest first. Lazily creates the welcome on the
    first open. User-message previews are ciphertext — the client decrypts them."""
    ensure_welcome(db, user)
    # The conversation partner for a message, from this user's perspective (system → 0).
    partner_expr = case(
        (Message.kind == "system", SYSTEM_PARTNER_ID),
        (Message.sender_id == user.id, Message.recipient_id),
        else_=Message.sender_id,
    )
    mine = or_(Message.sender_id == user.id, Message.recipient_id == user.id)
    # Latest message per conversation in ONE query (DISTINCT ON, Postgres) instead of loading the
    # user's whole message history into Python — bounded by the number of conversations.
    latest = (
        db.execute(
            select(Message)
            .where(mine)
            .distinct(partner_expr)
            .order_by(partner_expr, Message.created_at.desc())
        )
        .scalars()
        .all()
    )
    # Unread incoming, grouped per conversation.
    unread = dict(
        db.execute(
            select(partner_expr, func.count())
            .where(mine, Message.recipient_id == user.id, Message.read_at.is_(None))
            .group_by(partner_expr)
        ).all()
    )

    def _pid(m: Message) -> int:
        if m.kind == "system":
            return SYSTEM_PARTNER_ID
        return m.recipient_id if m.sender_id == user.id else m.sender_id

    user_ids = [_pid(m) for m in latest if m.kind != "system"]
    partners = {}
    if user_ids:
        partners = {
            u.id: u
            for u in db.execute(select(User).where(User.id.in_(user_ids))).scalars().all()
        }
    items = []
    for m in latest:
        pid = _pid(m)
        if pid == SYSTEM_PARTNER_ID:
            partner = _system_partner()
        elif pid in partners:
            partner = _serialize_user(partners[pid])
        else:
            continue
        items.append(
            {"partner": partner, "last_message": _serialize_msg(m), "unread": int(unread.get(pid, 0))}
        )
    items.sort(key=lambda x: x["last_message"]["created_at"] or "", reverse=True)
    return {"items": items}


@router.get("/conversations/{partner_id}")
def get_thread(
    partner_id: int,
    mark_read: bool = True,
    user: User = Depends(require_human),
    db: Session = Depends(get_db),
) -> dict:
    """Full thread with a partner (or the system thread when partner_id == 0), oldest first.
    Stamps unread incoming messages read unless mark_read=false."""
    if partner_id == SYSTEM_PARTNER_ID:
        ensure_welcome(db, user)
        partner = _system_partner()
        cond = and_(Message.recipient_id == user.id, Message.kind == "system")
    else:
        p = db.get(User, partner_id)
        cond = and_(
            Message.kind == "user",
            or_(
                and_(Message.sender_id == user.id, Message.recipient_id == partner_id),
                and_(Message.sender_id == partner_id, Message.recipient_id == user.id),
            ),
        )
        # MB2: don't leak arbitrary user identities. Resolve a partner only if they're messageable
        # OR we already share a thread (so a conversation with a since-suspended/deactivated user
        # still opens). Otherwise the same 404 as a non-existent id → no enumeration oracle.
        has_history = db.scalar(select(Message.id).where(cond).limit(1))
        if p is None or (not is_messageable_user(p) and not has_history):
            raise HTTPException(status_code=404, detail="User not found")
        partner = _serialize_user(p)
    msgs = (
        db.execute(select(Message).where(cond).order_by(Message.created_at.asc()))
        .scalars()
        .all()
    )
    if mark_read:
        now = datetime.now(timezone.utc)
        changed = False
        for m in msgs:
            if m.recipient_id == user.id and m.read_at is None:
                m.read_at = now
                changed = True
        if changed:
            db.commit()
            # MB3: ping this user's OTHER tabs/devices that their unread dropped so the badge
            # refreshes live instead of waiting for the next 20s poll. The client re-fetches the
            # count on this signal, so no number is carried here.
            manager.push(user.id, {"type": "unread"})
    return {"partner": partner, "items": [_serialize_msg(m) for m in msgs]}


@router.post("")
def send_message(
    payload: SendIn,
    user: User = Depends(require_human),
    db: Session = Depends(get_db),
) -> dict:
    """Send an end-to-end-encrypted message. The server stores only the ciphertext the client
    produced, then pushes it live to the recipient (and the sender's other devices)."""
    if not _send_limiter.allow(f"msg:{user.id}"):
        raise HTTPException(status_code=429, detail="Too many messages — slow down a moment.")
    ct = (payload.ciphertext or "").strip()
    iv = (payload.iv or "").strip()
    if not ct or not iv:
        raise HTTPException(status_code=400, detail="Message is empty.")
    if len(ct) > MAX_CIPHERTEXT:
        raise HTTPException(status_code=400, detail="Message is too long.")
    if payload.recipient_id == user.id:
        raise HTTPException(status_code=400, detail="You can't message yourself.")
    recipient = db.get(User, payload.recipient_id)
    if not is_messageable_user(recipient):
        raise HTTPException(status_code=404, detail="Recipient not available.")
    if db.get(MessageKey, recipient.id) is None:
        raise HTTPException(status_code=404, detail="That member hasn't set up messaging yet.")
    m = Message(
        kind="user", sender_id=user.id, recipient_id=recipient.id, ciphertext=ct, iv=iv
    )
    db.add(m)
    db.commit()
    db.refresh(m)
    # Carry both parties so each recipient can open/flash the right dock window (the partner is
    # whichever of from/to isn't them).
    event = {
        "type": "message",
        "message": _serialize_msg(m),
        "from": _serialize_user(user),
        "to": _serialize_user(recipient),
    }
    manager.push(recipient.id, event)
    manager.push(user.id, event)  # the sender's other open devices
    return _serialize_msg(m)


@router.get("/unread_count")
def unread_count(
    user: User = Depends(require_human), db: Session = Depends(get_db)
) -> dict:
    n = db.scalar(
        select(func.count())
        .select_from(Message)
        .where(Message.recipient_id == user.id, Message.read_at.is_(None))
    )
    return {"count": n or 0}


# --- live delivery -----------------------------------------------------------

@router.websocket("/ws")
async def messages_ws(ws: WebSocket) -> None:
    """Live push channel: authenticated by the session cookie, registers the connection so
    sent messages reach this user's open tabs instantly. We don't consume client→server frames
    (sending goes through POST /api/messages); the receive loop only detects disconnect."""
    # Which signed-in account this socket acts as — the SHARED per-tab resolution (wallet-gated
    # ?account= override, session default otherwise). A browser WebSocket can't send the
    # X-Siftlode-Account header, so resolved_user_id falls through to the ?account= query param.
    uid, _ = resolved_user_id(ws)
    if not uid:
        await ws.close(code=1008)
        return
    db = SessionLocal()
    try:
        user = db.get(User, uid)
        # SA4: honour server-side session revocation on the live channel too — reject a cookie whose
        # recorded epoch is behind the account's current one (mirrors current_user). Without this a
        # copied cookie could keep receiving pushes after a password reset / "log out everywhere".
        epochs = ws.session.get("epochs") or {}
        epoch_ok = user is not None and int(epochs.get(str(uid), 0)) == (user.session_epoch or 0)
        ok = epoch_ok and is_messageable_user(user)
    finally:
        db.close()
    if not ok:
        await ws.close(code=1008)
        return
    await manager.connect(uid, ws)
    try:
        while True:
            await ws.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        await manager.disconnect(uid, ws)