transient.chat/blueprints/chat.py

from flask import Blueprint, request, jsonify, Response
from datetime import datetime, timezone
from typing import Dict, Any, Tuple
import time

from blueprints.auth import verify_user_key
from models.state import (
    messages,
    room_users,
    authorized_room_users,
    kicked_users,
    room_owners,
    banned_ips,
    banned_usernames,
    user_ips,
    user_sessions,
    nuked_rooms,
)

chat_bp = Blueprint("chat", __name__)


@chat_bp.route("/chat/<room>", methods=["GET", "POST"])
def chat_room(room: str) -> Tuple[Response, int] | Response:
    room = room.strip()

    if request.method == "POST":
        data: Dict[str, Any] = request.get_json() or {}
        username: str = str(data.get("username", "")).strip()
        key: str = str(data.get("key", "")).strip()
        ciphertext: str = str(data.get("ciphertext", "")).strip()
        iv: str = str(data.get("iv", "")).strip()
        user_ip = request.remote_addr

        # Verify authentication
        if not username or not key or not verify_user_key(username, key):
            return jsonify({"error": "Invalid user credentials"}), 403

        # Verify room authorization
        if username not in authorized_room_users.get(room, set()):
            return jsonify({"error": "Room password authorization required", "code": "AUTH_REQUIRED"}), 401

        # Check bans and kicks
        if (
            user_ip in banned_ips
            or username in banned_usernames
            or (username in user_ips and user_ips[username] in banned_ips)
        ):
            return jsonify({"error": "banned"}), 403

        if username in kicked_users.get(room, set()):
            return jsonify({"error": "kicked"}), 403

        # Encryption-first enforcement: all messages stored in memory MUST be encrypted
        if not ciphertext or not iv:
            return jsonify({"error": "Encryption required: ciphertext and IV missing"}), 400

        # Ephemeral encrypted message in server memory
        message = {
            "username": username,
            "ciphertext": ciphertext,
            "iv": iv,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

        messages.setdefault(room, []).append(message)
        if len(messages[room]) > 100:
            messages[room] = messages[room][-100:]

        return jsonify({"status": "sent"})

    # GET endpoint (optional polling fallback)
    since = request.args.get("since")
    room_messages = messages.setdefault(room, [])

    if since:
        try:
            since_time = datetime.fromisoformat(since)
            new_msgs = [
                m
                for m in room_messages
                if datetime.fromisoformat(m["timestamp"]) > since_time
            ]
            return jsonify(new_msgs)
        except ValueError:
            pass

    return jsonify(room_messages)


@chat_bp.route("/chat/<room>/users", methods=["GET", "POST"])
def room_user_list(room: str) -> Tuple[Response, int] | Response:
    room = room.strip()

    if room in nuked_rooms:
        if time.time() < nuked_rooms[room]:
            return jsonify({"error": "Room is unavailable"}), 403
        else:
            nuked_rooms.pop(room, None)

    if request.method == "POST":
        data: Dict[str, Any] = request.get_json() or {}
        username: str = str(data.get("username", "")).strip()
        key: str = str(data.get("key", "")).strip()

        if not username:
            return jsonify({"error": "No username"}), 400

        if key and not verify_user_key(username, key):
            return jsonify({"error": "Invalid user credentials"}), 403

        if username not in authorized_room_users.get(room, set()):
            return jsonify({"error": "Room authorization required", "code": "AUTH_REQUIRED"}), 401

        if username in kicked_users.get(room, set()):
            return jsonify({"error": "Kicked from room"}), 403

        room_users.setdefault(room, {})
        room_users[room][username] = datetime.now(timezone.utc).isoformat()

    users = room_users.get(room, {})
    response: Dict[str, Any] = {"users": list(users.keys())}
    if room in room_owners:
        response["owner"] = room_owners[room]
    return jsonify(response)


@chat_bp.route("/chat/<room>/kick", methods=["POST"])
def kick_user(room: str) -> Tuple[Response, int] | Response:
    room = room.strip()
    data: Dict[str, Any] = request.get_json() or {}
    owner: str = str(data.get("owner", "")).strip()
    key: str = str(data.get("key", "")).strip()
    target: str = str(data.get("target", "")).strip()

    if not owner or not target:
        return jsonify({"error": "Missing owner or target"}), 400

    if key and not verify_user_key(owner, key):
        return jsonify({"error": "Invalid credentials"}), 403

    if room_owners.get(room) != owner:
        return jsonify({"error": "Not authorized"}), 403

    kicked_users.setdefault(room, set()).add(target)
    if room in authorized_room_users:
        authorized_room_users[room].discard(target)
    if room in room_users:
        room_users[room].pop(target, None)

    return jsonify({"kicked": target})


@chat_bp.route("/chat/<room>/poll", methods=["POST"])
def poll_room(room: str) -> Tuple[Response, int] | Response:
    room = room.strip()
    data: Dict[str, Any] = request.get_json() or {}
    username: str = str(data.get("username", "")).strip()
    key: str = str(data.get("key", "")).strip()
    since: str = data.get("since")

    if not username or not key or not verify_user_key(username, key):
        return jsonify({"error": "Invalid credentials"}), 403

    if room in nuked_rooms:
        if time.time() < nuked_rooms[room]:
            return jsonify({"error": "Room is unavailable"}), 403
        else:
            nuked_rooms.pop(room, None)

    if username in kicked_users.get(room, set()):
        return jsonify({"error": "Kicked from room"}), 403

    # Check room authorization
    if username not in authorized_room_users.get(room, set()):
        return jsonify({"error": "Room authorization required", "code": "AUTH_REQUIRED"}), 401

    # Update user presence
    room_users.setdefault(room, {})
    room_users[room][username] = datetime.now(timezone.utc).isoformat()

    user_sessions[username] = {
        "key": key,
        "last_active": datetime.now(timezone.utc).isoformat(),
    }

    # Fetch encrypted messages from memory
    room_messages = messages.setdefault(room, [])
    filtered_messages = room_messages
    if since:
        try:
            since_time = datetime.fromisoformat(since)
            filtered_messages = [
                m
                for m in room_messages
                if datetime.fromisoformat(m["timestamp"]) > since_time
            ]
        except ValueError:
            pass

    users_list = room_users.get(room, {})
    response: Dict[str, Any] = {
        "messages": filtered_messages,
        "users": list(users_list.keys()),
    }
    if room in room_owners:
        response["owner"] = room_owners[room]

    return jsonify(response)