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, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from sqlalchemy import and_, case, delete, 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
from app.security import verify_password
from app.userscope import is_messageable_user, messageable_clauses
from app.utils import iso
from app.refusal import refused
from app.routes._common import not_found

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)
_reset_key_limiter = RateLimiter(max_events=5, window_seconds=300)

# 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!"
    ),
}


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


class KeyResetIn(BaseModel):
    # Present only for password accounts; a Google-only account has no password to re-verify.
    current_password: str | None = None


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": iso(m.read_at),
        "created_at": iso(m.created_at),
    }


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 refused(409, "messaging_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 {"ok": True, "configured": True}


@router.post("/keys/reset")
def reset_keys(
    payload: KeyResetIn,
    user: User = Depends(require_human),
    db: Session = Depends(get_db),
) -> dict:
    """Forgotten-passphrase escape hatch: drop this user's secure-messaging key so they can set a
    NEW passphrase (the set-once guard in setup_keys otherwise makes a lost passphrase a permanent
    wall). Destroys the user's entire E2EE history BY DESIGN — every conversation key is derived
    from this keypair, so replacing it orphans every stored ciphertext in both directions. Those
    now-undecryptable `user` messages are deleted here (a clean slate for both parties) rather than
    left as permanent "can't decrypt" clutter; `system` messages (plaintext) are untouched.

    An account WITH a password must re-authenticate: a hijacked session must not be able to silently
    reset messaging and then read the victim's future messages under a new passphrase. A password-
    less (Google-only) account has nothing to re-verify, so it relies on its live session plus the
    client-side danger confirmation."""
    if not _reset_key_limiter.allow(str(user.id)):
        raise refused(429, "too_many_attempts")
    if user.password_hash and not verify_password(
        payload.current_password or "", user.password_hash
    ):
        raise refused(403, "wrong_password")
    # Delete the now-dead ciphertext first (both directions), then the wrapped keypair itself so a
    # fresh setup is allowed. Order doesn't matter (no FK between them), but do it in one commit.
    db.execute(
        delete(Message).where(
            Message.kind == "user",
            or_(Message.sender_id == user.id, Message.recipient_id == user.id),
        )
    )
    key = db.get(MessageKey, user.id)
    if key is not None:
        db.delete(key)
    db.commit()
    return {"ok": True, "configured": False}


@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. MB2 (like get_thread):
    resolve the target only if they're messageable OR we already share a thread, else the same
    "User not found" as a non-existent id — so this can't be probed to enumerate user_ids or to
    fetch keys for suspended/deactivated/demo accounts."""
    cond = and_(
        Message.kind == "user",
        or_(
            and_(Message.sender_id == user.id, Message.recipient_id == user_id),
            and_(Message.sender_id == user_id, Message.recipient_id == user.id),
        ),
    )
    p = db.get(User, user_id)
    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 not_found("user")
    k = db.get(MessageKey, user_id)
    if k is None:
        raise refused(404, "no_messaging")
    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_clauses())
            .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 not_found("user")
        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 refused(429, "message_rate")
    ct = (payload.ciphertext or "").strip()
    iv = (payload.iv or "").strip()
    if not ct or not iv:
        raise refused(400, "message_empty")
    if len(ct) > MAX_CIPHERTEXT:
        raise refused(400, "message_too_long")
    if payload.recipient_id == user.id:
        raise refused(400, "message_self")
    recipient = db.get(User, payload.recipient_id)
    if not is_messageable_user(recipient):
        raise not_found("recipient")
    if db.get(MessageKey, recipient.id) is None:
        raise refused(404, "no_messaging")
    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 (a wallet-gated
    # ?account=, the session default when none is named, nobody for an account no longer signed in
    # here). 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()
    # silent-ok: the ordinary way a chat socket ends — a closed tab. The `finally` below is the
    # whole handling, and nothing failed.
    except WebSocketDisconnect:
        pass
    finally:
        await manager.disconnect(uid, ws)