transient.chat/blueprints/auth.py

from flask import Blueprint, request, jsonify, current_app
from datetime import datetime, timezone
import os
from itsdangerous import URLSafeTimedSerializer, SignatureExpired, BadSignature

from models.state import user_sessions, user_ips, banned_ips, banned_usernames

auth_bp = Blueprint("auth", __name__)


def generate_user_key(username):
    serializer = URLSafeTimedSerializer(current_app.config["SECRET_KEY"])
    return serializer.dumps({"u": username})


def verify_user_key(username, token, max_age=30*24*60*60):  # 30 days
    serializer = URLSafeTimedSerializer(current_app.config["SECRET_KEY"])
    try:
        data = serializer.loads(token, max_age=max_age)
        return data.get("u") == username
    except (SignatureExpired, BadSignature):
        return False


@auth_bp.route("/register", methods=["POST"])
def register():
    data = request.get_json()
    username = data.get("username")
    user_ip = request.remote_addr

    if not username:
        return jsonify({"error": "No username provided"}), 400
    if username in user_ips:  # Check global user_ips to see if name is taken
        return jsonify({"error": "Username taken"}), 409
    if user_ip in banned_ips:
        return jsonify({"error": "IP banned"}), 403
    if username in banned_usernames:
        return jsonify({"error": "Username banned"}), 403

    user_key = generate_user_key(username)
    user_sessions[username] = {
        "key": user_key,
        "last_active": datetime.now(timezone.utc).isoformat(),
    }
    user_ips[username] = user_ip
    
    return jsonify({"username": username, "key": user_key})


@auth_bp.route("/validate", methods=["POST"])
def validate():
    data = request.get_json()
    username = data.get("username")
    key = data.get("key")

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

    if verify_user_key(username, key):
        # Restore active user session in-memory if lost on restart
        user_sessions[username] = {
            "key": key,
            "last_active": datetime.now(timezone.utc).isoformat(),
        }
        user_ips[username] = request.remote_addr
        return jsonify({"status": "valid"})
        
    return jsonify({"error": "Invalid credentials"}), 403


@auth_bp.route("/logout", methods=["POST"])
def logout():
    data = request.get_json()
    username = data.get("username")
    key = data.get("key")

    if not username or not key:
        return jsonify({"error": "Missing credentials"}), 400

    if verify_user_key(username, key):
        # Remove from active sessions
        user_sessions.pop(username, None)
        # Remove from user_ips in memory to free up the username immediately
        user_ips.pop(username, None)

        # Remove from all rooms
        from models.state import room_users
        for room in room_users:
            room_users[room].pop(username, None)

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

    return jsonify({"error": "Invalid credentials"}), 403