Create app.py

app.py

@@ -0,0 +1,448 @@
+import os
+import json
+import time
+import uuid
+import base64
+from typing import Optional, Dict, Set
+import asyncio
+import aiosqlite
+import tempfile
+import io
+
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, Request, UploadFile, File
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import Response, JSONResponse, FileResponse
+import logging
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.backends import default_backend
+
+from storage_handler import store_file, retrieve_file, delete_file, store_file_stream
+
+# --------------------------- Configuration ---------------------------
+DATABASE_URL = os.environ.get("DATABASE_URL", "/data/infinitychat.db")
+MESSAGE_KEY_B64 = os.environ.get("SECRET_KEY", None)
+if MESSAGE_KEY_B64 is None:
+    # Auto-generate (set permanently in production!)
+    MESSAGE_KEY_B64 = base64.urlsafe_b64encode(os.urandom(32)).decode()
+    print("WARNING: SECRET_KEY not set – generated random key.")
+MESSAGE_KEY = base64.urlsafe_b64decode(MESSAGE_KEY_B64)
+assert len(MESSAGE_KEY) == 32
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("InfinityChat")
+
+# --------------------------- Database ---------------------------
+async def init_db():
+    db = await aiosqlite.connect(DATABASE_URL)
+    db.row_factory = aiosqlite.Row
+    await db.execute("PRAGMA journal_mode=WAL;")
+    await db.execute("PRAGMA foreign_keys=ON;")
+    await db.execute("""
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY,
+            username TEXT UNIQUE NOT NULL,
+            display_name TEXT NOT NULL DEFAULT '',
+            password_hash TEXT NOT NULL,
+            salt TEXT NOT NULL,
+            avatar_path TEXT,
+            token TEXT UNIQUE,
+            created_at INTEGER DEFAULT (strftime('%s','now'))
+        )
+    """)
+    await db.execute("""
+        CREATE TABLE IF NOT EXISTS messages (
+            id INTEGER PRIMARY KEY,
+            sender_id INTEGER NOT NULL,
+            encrypted_content TEXT NOT NULL,
+            timestamp_ms INTEGER NOT NULL,
+            reply_to_id INTEGER,
+            is_edited INTEGER DEFAULT 0,
+            is_deleted INTEGER DEFAULT 0,
+            file_path TEXT,
+            file_type TEXT,
+            file_name TEXT,
+            FOREIGN KEY(sender_id) REFERENCES users(id)
+        )
+    """)
+    await db.execute("""
+        CREATE TABLE IF NOT EXISTS read_receipts (
+            user_id INTEGER NOT NULL,
+            message_id INTEGER NOT NULL,
+            PRIMARY KEY (user_id, message_id)
+        )
+    """)
+    await db.execute("CREATE INDEX IF NOT EXISTS idx_msg_time ON messages(timestamp_ms)")
+    await db.commit()
+    await db.close()
+
+asyncio.get_event_loop().run_until_complete(init_db())
+
+# --------------------------- Helpers ---------------------------
+async def get_db():
+    db = await aiosqlite.connect(DATABASE_URL)
+    db.row_factory = aiosqlite.Row
+    await db.execute("PRAGMA journal_mode=WAL")
+    return db
+
+def encrypt_message(plain: str) -> str:
+    aesgcm = AESGCM(MESSAGE_KEY)
+    nonce = os.urandom(12)
+    ct = aesgcm.encrypt(nonce, plain.encode(), None)
+    return base64.urlsafe_b64encode(nonce + ct).decode()
+
+def decrypt_message(encrypted_b64: str) -> str:
+    raw = base64.urlsafe_b64decode(encrypted_b64)
+    nonce, ct = raw[:12], raw[12:]
+    aesgcm = AESGCM(MESSAGE_KEY)
+    return aesgcm.decrypt(nonce, ct, None).decode()
+
+def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
+    if salt is None:
+        salt = os.urandom(16).hex()
+    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), iterations=600000)
+    return base64.urlsafe_b64encode(kdf.derive(password.encode())).decode(), salt
+
+def verify_password(password: str, salt: str, stored_hash: str) -> bool:
+    key, _ = hash_password(password, salt)
+    return key == stored_hash
+
+def generate_token() -> str:
+    return base64.urlsafe_b64encode(os.urandom(32)).decode()
+
+async def get_user_by_token(token: str):
+    db = await get_db()
+    try:
+        async with db.execute("SELECT * FROM users WHERE token = ?", (token,)) as cur:
+            return await cur.fetchone()
+    finally:
+        await db.close()
+
+# --------------------------- FastAPI app ---------------------------
+app = FastAPI(title="InfinityChat")
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# --------------------------- HTTP endpoints ---------------------------
+@app.post("/signup")
+async def signup(username: str, password: str, display_name: str = ""):
+    db = await get_db()
+    if await (await db.execute("SELECT 1 FROM users WHERE username = ?", (username,))).fetchone():
+        await db.close()
+        raise HTTPException(400, "Username taken")
+    pwd_hash, salt = hash_password(password)
+    token = generate_token()
+    await db.execute("INSERT INTO users (username, display_name, password_hash, salt, token) VALUES (?,?,?,?,?)",
+                     (username, display_name, pwd_hash, salt, token))
+    await db.commit()
+    await db.close()
+    return {"token": token, "username": username, "display_name": display_name}
+
+@app.post("/login")
+async def login(username: str, password: str):
+    db = await get_db()
+    user = await (await db.execute("SELECT * FROM users WHERE username = ?", (username,))).fetchone()
+    if not user or not verify_password(password, user["salt"], user["password_hash"]):
+        await db.close()
+        raise HTTPException(400, "Invalid credentials")
+    token = generate_token()
+    await db.execute("UPDATE users SET token = ? WHERE id = ?", (token, user["id"]))
+    await db.commit()
+    await db.close()
+    return {"token": token, "username": user["username"], "display_name": user["display_name"], "avatar_path": user["avatar_path"]}
+
+@app.get("/me")
+async def me(token: str = Header(...)):
+    user = await get_user_by_token(token)
+    if not user:
+        raise HTTPException(401)
+    return {"id": user["id"], "username": user["username"], "display_name": user["display_name"], "avatar_path": user["avatar_path"]}
+
+@app.post("/upload_avatar")
+async def upload_avatar(file: UploadFile = File(...), token: str = Header(...)):
+    user = await get_user_by_token(token)
+    if not user:
+        raise HTTPException(401)
+    data = await file.read()
+    ext = os.path.splitext(file.filename)[1] if file.filename else ".jpg"
+    remote = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
+    store_file(remote, data)
+    db = await get_db()
+    await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote, user["id"]))
+    await db.commit()
+    await db.close()
+    return {"avatar_path": remote}
+
+# Chunked upload endpoint
+upload_sessions: Dict[str, dict] = {}  # upload_id -> {chunks: dict, filename, file_type, total_chunks, user_id}
+
+@app.post("/upload_chunk")
+async def upload_chunk(
+    file: UploadFile = File(...),
+    chunk_index: int = Query(...),
+    total_chunks: int = Query(...),
+    file_name: str = Query(...),
+    file_type: str = Query(...),
+    upload_id: str = Query(...),
+    token: str = Header(...)
+):
+    user = await get_user_by_token(token)
+    if not user:
+        raise HTTPException(401)
+    if upload_id not in upload_sessions:
+        upload_sessions[upload_id] = {
+            "chunks": {},
+            "filename": file_name,
+            "file_type": file_type,
+            "total_chunks": total_chunks,
+            "user_id": user["id"],
+            "received": 0
+        }
+    ses = upload_sessions[upload_id]
+    if chunk_index in ses["chunks"]:
+        return {"status": "duplicate"}
+    ses["chunks"][chunk_index] = await file.read()
+    ses["received"] += 1
+    if ses["received"] == total_chunks:
+        # Reassemble
+        full_data = b"".join(ses["chunks"][i] for i in sorted(ses["chunks"]))
+        remote = f"uploads/{user['username']}_{uuid.uuid4().hex}/{file_name}"
+        try:
+            # Encrypt and store via bucket
+            store_file(remote, full_data)
+        except Exception as e:
+            logger.error(f"Failed to store file: {e}")
+            del upload_sessions[upload_id]
+            raise HTTPException(500, "Storage error")
+        del upload_sessions[upload_id]
+        return {
+            "status": "complete",
+            "file_path": remote,
+            "file_name": file_name,
+            "file_type": file_type
+        }
+    return {"status": "chunk_ok"}
+
+# Downloads proxy (forces auth)
+@app.get("/download/{file_path:path}")
+async def download(file_path: str, token: str = Header(...)):
+    user = await get_user_by_token(token)
+    if not user:
+        raise HTTPException(401)
+    try:
+        data = retrieve_file(file_path)
+    except Exception as e:
+        raise HTTPException(404, "File not found or decryption failed")
+    import mimetypes
+    mt = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
+    return Response(content=data, media_type=mt)
+
+@app.get("/")
+async def root():
+    return FileResponse("static/index.html")
+
+# --------------------------- WebSocket ---------------------------
+class ConnectionManager:
+    def __init__(self):
+        self.active: Dict[str, Set[WebSocket]] = {}  # username -> set of websockets
+        self.user_read: Dict[int, int] = {}          # user_id -> last message id marked read
+
+    async def connect(self, ws: WebSocket, username: str):
+        await ws.accept()
+        self.active.setdefault(username, set()).add(ws)
+
+    def disconnect(self, ws: WebSocket, username: str):
+        if username in self.active:
+            self.active[username].discard(ws)
+            if not self.active[username]:
+                del self.active[username]
+
+    async def broadcast(self, msg: dict, exclude: WebSocket = None):
+        for uname, ws_set in self.active.items():
+            for w in list(ws_set):
+                if w == exclude:
+                    continue
+                try:
+                    await w.send_json(msg)
+                except Exception:
+                    self.disconnect(w, uname)
+
+    def online_users(self) -> list:
+        return list(self.active.keys())
+
+manager = ConnectionManager()
+
+@app.websocket("/ws")
+async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
+    if not token:
+        await ws.close(code=4001)
+        return
+    user = await get_user_by_token(token)
+    if not user:
+        await ws.close(code=4001, reason="Invalid token")
+        return
+
+    username = user["username"]
+    uid = user["id"]
+    await manager.connect(ws, username)
+
+    # Send online users immediately
+    online = manager.online_users()
+    await ws.send_json({"type": "online_users", "users": online})
+    await manager.broadcast({"type": "user_online", "username": username}, exclude=ws)
+
+    try:
+        while True:
+            data = await ws.receive_json()
+            mtype = data.get("type")
+
+            # ---- New message ----
+            if mtype == "chat_message":
+                content = data.get("content", "")
+                reply_to = data.get("reply_to_id")
+                client_id = data.get("client_id")
+                file_path = data.get("file_path")
+                file_type = data.get("file_type")
+                file_name = data.get("file_name")
+                ts = int(time.time() * 1000)
+
+                encrypted = encrypt_message(content)
+                db = await get_db()
+                cur = await db.execute(
+                    """INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name)
+                       VALUES (?, ?, ?, ?, ?, ?, ?)""",
+                    (uid, encrypted, ts, reply_to, file_path, file_type, file_name))
+                mid = cur.lastrowid
+                await db.commit()
+                await db.close()
+
+                msg_out = {
+                    "type": "new_message",
+                    "id": mid,
+                    "sender_id": uid,
+                    "username": username,
+                    "display_name": user["display_name"],
+                    "avatar_path": user["avatar_path"],
+                    "content": content,
+                    "timestamp_ms": ts,
+                    "reply_to_id": reply_to,
+                    "file_path": file_path,
+                    "file_type": file_type,
+                    "file_name": file_name,
+                    "is_edited": False,
+                    "is_deleted": False,
+                    "client_id": client_id,
+                    "status": "sent"
+                }
+                # Deliver to sender
+                await ws.send_json({**msg_out, "status": "delivered"})
+                # Broadcast to others
+                await manager.broadcast(msg_out, exclude=ws)
+
+            # ---- Load history (cursor‑based) ----
+            elif mtype == "load_messages":
+                before_id = data.get("before_id")
+                limit = min(data.get("limit", 50), 100)
+                db = await get_db()
+                q = "SELECT * FROM messages WHERE is_deleted = 0"
+                params = []
+                if before_id is not None:
+                    q += " AND id < ?"
+                    params.append(before_id)
+                q += " ORDER BY id DESC LIMIT ?"
+                params.append(limit)
+                rows = await (await db.execute(q, tuple(params))).fetchall()
+                rows = list(reversed(rows))
+                msgs = []
+                for row in rows:
+                    try:
+                        plain = decrypt_message(row["encrypted_content"])
+                    except Exception:
+                        plain = "[decryption error]"
+                    s = await db.execute("SELECT id, username, display_name, avatar_path FROM users WHERE id = ?", (row["sender_id"],))
+                    s = await s.fetchone()
+                    msgs.append({
+                        "id": row["id"],
+                        "sender_id": row["sender_id"],
+                        "username": s["username"] if s else "unknown",
+                        "display_name": s["display_name"] if s else "",
+                        "avatar_path": s["avatar_path"] if s else "",
+                        "content": plain,
+                        "timestamp_ms": row["timestamp_ms"],
+                        "reply_to_id": row["reply_to_id"],
+                        "file_path": row["file_path"],
+                        "file_type": row["file_type"],
+                        "file_name": row["file_name"],
+                        "is_edited": bool(row["is_edited"]),
+                        "is_deleted": bool(row["is_deleted"])
+                    })
+                await db.close()
+                await ws.send_json({"type": "messages_batch", "messages": msgs, "before_id": before_id})
+
+            # ---- Edit ----
+            elif mtype == "edit_message":
+                mid, new_text = data.get("message_id"), data.get("content")
+                db = await get_db()
+                row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0", (mid, uid))
+                if not await row.fetchone():
+                    await db.close()
+                    continue
+                new_enc = encrypt_message(new_text)
+                await db.execute("UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?", (new_enc, mid))
+                await db.commit()
+                await db.close()
+                await manager.broadcast({"type": "edit_message", "message_id": mid, "content": new_text})
+
+            # ---- Delete (permanent) ----
+            elif mtype == "delete_message":
+                mid = data.get("message_id")
+                db = await get_db()
+                row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid))
+                msg = await row.fetchone()
+                if not msg:
+                    await db.close()
+                    continue
+                if msg["file_path"]:
+                    try:
+                        delete_file(msg["file_path"])
+                    except Exception:
+                        pass
+                await db.execute("DELETE FROM messages WHERE id = ?", (mid,))
+                await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
+                await db.commit()
+                await db.close()
+                await manager.broadcast({"type": "delete_message", "message_id": mid})
+
+            # ---- Typing ----
+            elif mtype == "typing":
+                await manager.broadcast({"type": "typing_update", "username": username, "is_typing": data.get("is_typing", False)}, exclude=ws)
+
+            # ---- Read receipts ----
+            elif mtype == "mark_read":
+                up_to = data.get("up_to_message_id")
+                if not up_to:
+                    continue
+                db = await get_db()
+                # Which messages sent by others, <= up_to, not deleted, and not yet marked read by this user?
+                rows = await db.execute(
+                    """SELECT id, sender_id FROM messages
+                       WHERE id <= ? AND sender_id != ? AND is_deleted = 0
+                       EXCEPT SELECT message_id FROM read_receipts WHERE user_id = ?""",
+                    (up_to, uid, uid))
+                receipts = []
+                async for r in rows:
+                    await db.execute("INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?, ?)", (uid, r["id"]))
+                    receipts.append(r["id"])
+                await db.commit()
+                await db.close()
+                # Notify the original senders
+                for mid in receipts:
+                    await manager.broadcast({"type": "read_receipt", "message_id": mid, "reader_username": username})
+
+    except WebSocketDisconnect:
+        pass
+    finally:
+        manager.disconnect(ws, username)
+        if username not in manager.active:
+            await manager.broadcast({"type": "user_offline", "username": username})