app.py

import json
import time
import re
import os
import hashlib
import secrets
import base64
import sqlite3
import asyncio
from collections import defaultdict
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
DB_PATH = "/data/chat.db"

def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn

def init_db():
    with get_db() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                email TEXT UNIQUE NOT NULL,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                salt TEXT NOT NULL,
                created_at TEXT DEFAULT (datetime('now'))
            );
            CREATE TABLE IF NOT EXISTS sessions (
                token TEXT PRIMARY KEY,
                user_id INTEGER NOT NULL,
                created_at TEXT DEFAULT (datetime('now')),
                FOREIGN KEY (user_id) REFERENCES users(id)
            );
        """)

@asynccontextmanager
async def lifespan(app):
    init_db()
    yield

app = FastAPI(lifespan=lifespan)
app.mount("/static", StaticFiles(directory="static"), name="static")

def hash_password(password: str, salt: str = None):
    if salt is None:
        salt = secrets.token_urlsafe(32)
    key = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 100000, dklen=32)
    return base64.b64encode(key).decode(), salt

def verify_password(password: str, salt: str, stored_hash: str) -> bool:
    computed, _ = hash_password(password, salt)
    return computed == stored_hash

attempts: dict[str, list[float]] = defaultdict(list)
blocked: dict[str, float] = {}

def check_rate_limit(ip: str) -> bool:
    now = time.time()
    if block_until := blocked.get(ip):
        if block_until > now:
            return False
        del blocked[ip]
    attempts[ip] = [t for t in attempts[ip] if t > now - 240]
    if len(attempts[ip]) >= 5:
        blocked[ip] = now + 240
        attempts[ip] = []
        return False
    return True

def record_failed(ip: str):
    now = time.time()
    attempts[ip].append(now)
    if len(attempts[ip]) >= 5:
        blocked[ip] = now + 240
        attempts[ip] = []

def create_session(user_id: int) -> str:
    token = secrets.token_urlsafe(32)
    with get_db() as conn:
        conn.execute("INSERT INTO sessions (token, user_id) VALUES (?, ?)", (token, user_id))
    return token

def get_user_from_session(token: str):
    with get_db() as conn:
        row = conn.execute(
            "SELECT u.id, u.username, u.email FROM users u "
            "JOIN sessions s ON u.id = s.user_id WHERE s.token = ?",
            (token,),
        ).fetchone()
    return dict(row) if row else None

def delete_session(token: str):
    with get_db() as conn:
        conn.execute("DELETE FROM sessions WHERE token = ?", (token,))

messages: list[dict] = []
connections: dict[int, list[WebSocket]] = {}

USERNAME_RE = re.compile(r"^[a-zA-Z0-9_\-]{3,20}$")

def get_client_ip(request: Request) -> str:
    fwd = request.headers.get("x-forwarded-for")
    return fwd.split(",")[0].strip() if fwd else (request.client.host if request.client else "unknown")

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.post("/api/signup")
async def signup(request: Request):
    try:
        data = await request.json()
    except Exception:
        raise HTTPException(400, "Invalid JSON")

    email = data.get("email", "").strip().lower()
    username = data.get("username", "").strip()
    password = data.get("password", "")

    if not email or "@" not in email or "." not in email.split("@")[-1]:
        return JSONResponse({"error": "Invalid email address."}, 400)
    if not USERNAME_RE.match(username):
        return JSONResponse({"error": "Username must be 3-20 chars (letters, numbers, _, -)."}, 400)
    if len(password) < 8:
        return JSONResponse({"error": "Password must be at least 8 characters."}, 400)

    pw_hash, salt = hash_password(password)

    try:
        with get_db() as conn:
            conn.execute(
                "INSERT INTO users (email, username, password_hash, salt) VALUES (?, ?, ?, ?)",
                (email, username, pw_hash, salt),
            )
    except sqlite3.IntegrityError:
        return JSONResponse({"error": "Email or username already taken."}, 409)

    with get_db() as conn:
        user = conn.execute("SELECT id FROM users WHERE username = ?", (username,)).fetchone()
    token = create_session(user["id"])

    resp = JSONResponse({"success": True, "username": username, "token": token})
    resp.set_cookie("session_token", token, httponly=True, samesite="strict", max_age=86400 * 30)
    return resp

@app.post("/api/login")
async def login(request: Request):
    ip = get_client_ip(request)
    if not check_rate_limit(ip):
        return JSONResponse({"error": "Too many attempts. Wait 4 minutes."}, 429)

    try:
        data = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON."}, 400)

    login_id = data.get("login", "").strip()
    password = data.get("password", "")

    if not login_id or not password:
        record_failed(ip)
        return JSONResponse({"error": "Both fields are required."}, 400)

    with get_db() as conn:
        user = conn.execute(
            "SELECT id, username, email, password_hash, salt FROM users "
            "WHERE username = ? OR email = ?",
            (login_id, login_id.lower()),
        ).fetchone()

    if not user:
        record_failed(ip)
        return JSONResponse({"error": "Invalid credentials."}, 401)

    if not verify_password(password, user["salt"], user["password_hash"]):
        record_failed(ip)
        return JSONResponse({"error": "Invalid credentials."}, 401)

    token = create_session(user["id"])
    resp = JSONResponse({"success": True, "username": user["username"], "token": token})
    resp.set_cookie("session_token", token, httponly=True, samesite="strict", max_age=86400 * 30)
    return resp

@app.post("/api/logout")
async def logout(request: Request):
    token = request.cookies.get("session_token", "")
    if token:
        delete_session(token)
    resp = JSONResponse({"success": True})
    resp.delete_cookie("session_token")
    return resp

WEBHOOK_PASSWORD = os.environ.get("WEBHOOK_PASSWORD", "admin123")

@app.post("/api/webhook")
async def webhook(request: Request):
    try:
        data = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON"}, 400)

    if data.get("password") != WEBHOOK_PASSWORD:
        return JSONResponse({"error": "Unauthorized"}, 401)

    message = data.get("message", "").strip()
    if not message:
        return JSONResponse({"error": "Message is required"}, 400)

    msg_data = {
        "user": "System",
        "msg": message,
        "server_time": int(time.time() * 1000),
        "is_system": True,
    }
    messages.append(msg_data)

    for conns in connections.values():
        for conn in conns:
            try:
                await conn.send_json(msg_data)
            except Exception:
                pass

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

@app.get("/api/session")
async def check_session(request: Request):
    token = request.cookies.get("session_token", "")
    if not token:
        return JSONResponse({"authenticated": False})
    user = get_user_from_session(token)
    if not user:
        return JSONResponse({"authenticated": False})
    return JSONResponse({"authenticated": True, "username": user["username"], "token": token})

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    token = websocket.query_params.get("token", "")
    if not token:
        await websocket.close(code=4001)
        return

    user = get_user_from_session(token)
    if not user:
        await websocket.close(code=4001)
        return

    await websocket.accept()
    user_id = user["id"]
    username = user["username"]

    connections.setdefault(user_id, []).append(websocket)

    try:
        for msg in messages:
            await websocket.send_json(msg)

        while True:
            try:
                raw = await asyncio.wait_for(websocket.receive_text(), timeout=300)
            except asyncio.TimeoutError:
                try:
                    await websocket.close(code=4002, reason="idle_timeout")
                except Exception:
                    pass
                break

            msg_data = json.loads(raw)
            msg_data["user_id"] = user_id
            msg_data["user"] = username
            msg_data["server_time"] = int(time.time() * 1000)

            messages.append(msg_data)

            for uid, conns in connections.items():
                for conn in conns:
                    try:
                        await conn.send_json(msg_data)
                    except Exception:
                        pass

    except WebSocketDisconnect:
        pass
    except Exception:
        pass
    finally:
        conns = connections.get(user_id, [])
        if websocket in conns:
            conns.remove(websocket)
        if not conns:
            connections.pop(user_id, None)