Update app.py

app.py

@@ -19,7 +19,10 @@ 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, file_exists
+from storage_handler import (
+    store_file, retrieve_file, delete_file, file_exists,
+    download_database, upload_database, start_db_sync
+)
 
 # ------------------------------------------------------------------------
 # Configuration
@@ -36,9 +39,17 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("InfinityChat")
 
 # ------------------------------------------------------------------------
-# Database initialization (once at startup)
+# Database initialization
 # ------------------------------------------------------------------------
 async def init_database():
+    os.makedirs(os.path.dirname(DATABASE_URL), exist_ok=True)
+
+    db_existed = download_database(DATABASE_URL)
+    if db_existed:
+        logger.info("✅ Restored database from bucket")
+    else:
+        logger.info("🆕 Starting with fresh database")
+
     async with aiosqlite.connect(DATABASE_URL) as db:
         db.row_factory = aiosqlite.Row
         await db.execute("PRAGMA journal_mode=WAL;")
@@ -90,6 +101,9 @@ async def init_database():
         await db.execute("CREATE INDEX IF NOT EXISTS idx_messages_not_deleted ON messages(is_deleted, id)")
         await db.execute("CREATE INDEX IF NOT EXISTS idx_receipts_user_msg ON read_receipts(user_id, message_id)")
         await db.commit()
+
+    upload_database(DATABASE_URL)
+    start_db_sync(DATABASE_URL)
     logger.info("✅ Database initialized successfully")
 
 # ------------------------------------------------------------------------
@@ -99,9 +113,17 @@ async def init_database():
 async def lifespan(app: FastAPI):
     await init_database()
     yield
+    logger.info("🔄 Final database sync on shutdown...")
+    upload_database(DATABASE_URL)
 
 app = FastAPI(title="InfinityChat", version="1.0.0", lifespan=lifespan)
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"]
+)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # ------------------------------------------------------------------------
@@ -114,6 +136,11 @@ async def get_db() -> aiosqlite.Connection:
     await conn.execute("PRAGMA foreign_keys=ON")
     return conn
 
+def schedule_db_sync():
+    import threading
+    t = threading.Thread(target=upload_database, args=(DATABASE_URL,), daemon=True)
+    t.start()
+
 # ------------------------------------------------------------------------
 # Encryption helpers (AES-256-GCM)
 # ------------------------------------------------------------------------
@@ -138,7 +165,13 @@ def decrypt_message(encrypted_b64: str) -> str:
 def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
     if salt is None:
         salt = os.urandom(32).hex()
-    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), iterations=600000, backend=default_backend())
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt.encode(),
+        iterations=600000,
+        backend=default_backend()
+    )
     key = base64.urlsafe_b64encode(kdf.derive(password.encode())).decode()
     return key, salt
 
@@ -157,10 +190,16 @@ async def authenticate_user(token: str) -> Optional[dict]:
         return None
     db = await get_db()
     try:
-        cursor = await db.execute("SELECT id, username, display_name, avatar_path, status FROM users WHERE token = ?", (token,))
+        cursor = await db.execute(
+            "SELECT id, username, display_name, avatar_path, status FROM users WHERE token = ?",
+            (token,)
+        )
         user = await cursor.fetchone()
         if user:
-            await db.execute("UPDATE users SET last_seen = ? WHERE id = ?", (int(time.time()), user['id']))
+            await db.execute(
+                "UPDATE users SET last_seen = ? WHERE id = ?",
+                (int(time.time()), user['id'])
+            )
             await db.commit()
             return dict(user)
         return None
@@ -175,9 +214,13 @@ async def health_check():
     return {"status": "healthy", "timestamp": int(time.time()), "version": "1.0.0"}
 
 @app.post("/api/auth/signup")
-async def signup(username: str = Query(..., min_length=3, max_length=30), password: str = Query(..., min_length=6), display_name: str = Query(None, max_length=50)):
-    if not username.isalnum() and '_' not in username:
-        raise HTTPException(400, "Invalid username")
+async def signup(
+    username: str = Query(..., min_length=3, max_length=30),
+    password: str = Query(..., min_length=6),
+    display_name: str = Query(None, max_length=50)
+):
+    if not username.replace('_', '').isalnum():
+        raise HTTPException(400, "Invalid username - only letters, numbers and underscores allowed")
     db = await get_db()
     try:
         async with db.execute("SELECT id FROM users WHERE username = ?", (username.lower(),)) as cursor:
@@ -187,11 +230,21 @@ async def signup(username: str = Query(..., min_length=3, max_length=30), passwo
         token = generate_token()
         display = display_name or username
         await db.execute(
-            "INSERT INTO users (username, display_name, password_hash, salt, token, status) VALUES (?,?,?,?,?, 'online')",
+            "INSERT INTO users (username, display_name, password_hash, salt, token, status) VALUES (?,?,?,?,?,'online')",
             (username.lower(), display, pwd_hash, salt, token)
         )
         await db.commit()
-        return {"token": token, "user": {"id": None, "username": username, "display_name": display, "avatar_path": None, "status": "online"}}
+        schedule_db_sync()
+        return {
+            "token": token,
+            "user": {
+                "id": None,
+                "username": username.lower(),
+                "display_name": display,
+                "avatar_path": None,
+                "status": "online"
+            }
+        }
     finally:
         await db.close()
 
@@ -204,8 +257,12 @@ async def login(username: str = Query(...), password: str = Query(...)):
         if not user or not verify_password(password, user['salt'], user['password_hash']):
             raise HTTPException(401, "Invalid credentials")
         token = generate_token()
-        await db.execute("UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?", (token, int(time.time()), user['id']))
+        await db.execute(
+            "UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?",
+            (token, int(time.time()), user['id'])
+        )
         await db.commit()
+        schedule_db_sync()
         return {
             "token": token,
             "user": {
@@ -222,15 +279,20 @@ async def login(username: str = Query(...), password: str = Query(...)):
 @app.get("/api/auth/verify")
 async def verify_token(token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user: raise HTTPException(401)
+    if not user:
+        raise HTTPException(401)
     return {"user": user}
 
 @app.post("/api/auth/logout")
 async def logout(token: str = Header(..., alias="X-Auth-Token")):
     db = await get_db()
     try:
-        await db.execute("UPDATE users SET token = NULL, status = 'offline', last_seen = ? WHERE token = ?", (int(time.time()), token))
+        await db.execute(
+            "UPDATE users SET token = NULL, status = 'offline', last_seen = ? WHERE token = ?",
+            (int(time.time()), token)
+        )
         await db.commit()
+        schedule_db_sync()
         return {"status": "logged_out"}
     finally:
         await db.close()
@@ -241,36 +303,53 @@ async def logout(token: str = Header(..., alias="X-Auth-Token")):
 @app.get("/api/profile")
 async def get_profile(token: str = Header(..., alias="X-Auth-Token")):
     user = await authenticate_user(token)
-    if not user: raise HTTPException(401)
+    if not user:
+        raise HTTPException(401)
     return {"user": user}
 
 @app.patch("/api/profile")
-async def update_profile(display_name: str = Query(...), token: str = Header(..., alias="X-Auth-Token")):
+async def update_profile(
+    display_name: str = Query(...),
+    token: str = Header(..., alias="X-Auth-Token")
+):
     user = await authenticate_user(token)
-    if not user: raise HTTPException(401)
+    if not user:
+        raise HTTPException(401)
     db = await get_db()
     try:
-        await db.execute("UPDATE users SET display_name = ? WHERE id = ?", (display_name, user['id']))
+        await db.execute(
+            "UPDATE users SET display_name = ? WHERE id = ?",
+            (display_name, user['id'])
+        )
         await db.commit()
+        schedule_db_sync()
         user['display_name'] = display_name
         return {"user": user}
     finally:
         await db.close()
 
 @app.post("/api/profile/avatar")
-async def upload_avatar(file: UploadFile = File(...), token: str = Header(..., alias="X-Auth-Token")):
+async def upload_avatar(
+    file: UploadFile = File(...),
+    token: str = Header(..., alias="X-Auth-Token")
+):
     user = await authenticate_user(token)
-    if not user: raise HTTPException(401)
+    if not user:
+        raise HTTPException(401)
     data = await file.read()
-    if len(data) > 5*1024*1024:
-        raise HTTPException(400, "Image too large")
+    if len(data) > 5 * 1024 * 1024:
+        raise HTTPException(400, "Image too large (max 5MB)")
     ext = os.path.splitext(file.filename)[1] if file.filename else '.jpg'
     remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
     store_file(remote_path, data)
     db = await get_db()
     try:
-        await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote_path, user['id']))
+        await db.execute(
+            "UPDATE users SET avatar_path = ? WHERE id = ?",
+            (remote_path, user['id'])
+        )
         await db.commit()
+        schedule_db_sync()
     finally:
         await db.close()
     return {"avatar_path": remote_path}
@@ -292,42 +371,78 @@ async def upload_chunk(
     token: str = Header(..., alias="X-Auth-Token")
 ):
     user = await authenticate_user(token)
-    if not user: raise HTTPException(401)
+    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,
-            "file_size": file_size, "total_chunks": total_chunks, "received": 0,
-            "user_id": user['id'], "created_at": time.time()
+            "chunks": {},
+            "filename": file_name,
+            "file_type": file_type,
+            "file_size": file_size,
+            "total_chunks": total_chunks,
+            "received": 0,
+            "user_id": user['id'],
+            "created_at": time.time()
         }
+
     session = upload_sessions[upload_id]
-    if session['user_id'] != user['id']: raise HTTPException(403)
+    if session['user_id'] != user['id']:
+        raise HTTPException(403)
+
     chunk_data = await file.read()
     session['chunks'][chunk_index] = chunk_data
     session['received'] += 1
+
     if session['received'] == total_chunks:
         full_data = b"".join(session['chunks'][i] for i in sorted(session['chunks']))
         remote_path = f"uploads/{user['username']}/{uuid.uuid4().hex}/{file_name}"
         store_file(remote_path, full_data)
         del upload_sessions[upload_id]
-        return {"status": "complete", "file_path": remote_path, "file_name": file_name, "file_type": file_type, "file_size": file_size}
-    return {"status": "in_progress", "received": session['received'], "total": total_chunks}
+        return {
+            "status": "complete",
+            "file_path": remote_path,
+            "file_name": file_name,
+            "file_type": file_type,
+            "file_size": file_size
+        }
+
+    return {
+        "status": "in_progress",
+        "received": session['received'],
+        "total": total_chunks
+    }
 
 # ------------------------------------------------------------------------
-# File Download
+# File Download - accepts token as header OR ?t= query param
+# so <img src> and window.open() work without fetch headers
 # ------------------------------------------------------------------------
 @app.get("/api/download/{file_path:path}")
-async def download_file(file_path: str, token: str = Header(None, alias="X-Auth-Token"), t: str = Query(None)):
+async def download_file(
+    file_path: str,
+    token: str = Header(None, alias="X-Auth-Token"),
+    t: str = Query(None)
+):
     auth_token = token or t
     user = await authenticate_user(auth_token)
-    if not user: raise HTTPException(401)
-    if '..' in file_path: raise HTTPException(400)
+    if not user:
+        raise HTTPException(401)
+    if '..' in file_path:
+        raise HTTPException(400, "Invalid path")
     try:
         data = retrieve_file(file_path)
     except FileNotFoundError:
-        raise HTTPException(404)
+        raise HTTPException(404, "File not found")
+    except Exception as e:
+        logger.error(f"Download error for {file_path}: {e}")
+        raise HTTPException(500, "Failed to retrieve file")
     import mimetypes
     mt = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-    return Response(content=data, media_type=mt)
+    return Response(
+        content=data,
+        media_type=mt,
+        headers={"Cache-Control": "private, max-age=3600"}
+    )
 
 # ------------------------------------------------------------------------
 # WebSocket Manager
@@ -350,7 +465,8 @@ class ConnectionManager:
 
     async def broadcast(self, msg: dict, exclude_user_id: Optional[int] = None):
         for uid, info in list(self.active.items()):
-            if uid == exclude_user_id: continue
+            if uid == exclude_user_id:
+                continue
             for ws in list(info["connections"]):
                 try:
                     await ws.send_json(msg)
@@ -366,7 +482,10 @@ class ConnectionManager:
                     self.disconnect(ws, user_id)
 
     def get_online_users(self) -> List[Dict]:
-        return [{"id": uid, "username": info["username"]} for uid, info in self.active.items()]
+        return [
+            {"id": uid, "username": info["username"]}
+            for uid, info in self.active.items()
+        ]
 
     def is_online(self, user_id: int) -> bool:
         return user_id in self.active
@@ -390,7 +509,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
 
     await manager.connect(ws, uid, username)
 
-    # Send initial data
     await ws.send_json({
         "type": "connection_established",
         "user_id": uid,
@@ -400,7 +518,6 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
         "online_users": manager.get_online_users()
     })
 
-    # Broadcast new online status
     await manager.broadcast({
         "type": "user_status",
         "user_id": uid,
@@ -426,10 +543,18 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 client_id = data.get("client_id")
 
                 if not content and not file_path:
-                    await ws.send_json({"type": "error", "code": "EMPTY", "message": "Message cannot be empty"})
+                    await ws.send_json({
+                        "type": "error",
+                        "code": "EMPTY",
+                        "message": "Message cannot be empty"
+                    })
                     continue
                 if len(content) > 10000:
-                    await ws.send_json({"type": "error", "code": "TOO_LONG", "message": "Message too long"})
+                    await ws.send_json({
+                        "type": "error",
+                        "code": "TOO_LONG",
+                        "message": "Message too long"
+                    })
                     continue
 
                 encrypted = encrypt_message(content)
@@ -438,11 +563,16 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 db = await get_db()
                 try:
                     cursor = await db.execute(
-                        "INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name, file_size) VALUES (?,?,?,?,?,?,?,?)",
-                        (uid, encrypted, ts, reply_to_id, file_path, file_type, file_name, file_size)
+                        """INSERT INTO messages
+                           (sender_id, encrypted_content, timestamp_ms, reply_to_id,
+                            file_path, file_type, file_name, file_size)
+                           VALUES (?,?,?,?,?,?,?,?)""",
+                        (uid, encrypted, ts, reply_to_id,
+                         file_path, file_type, file_name, file_size)
                     )
                     mid = cursor.lastrowid
                     await db.commit()
+                    schedule_db_sync()
                 finally:
                     await db.close()
 
@@ -464,14 +594,12 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     "status": "sent"
                 }
 
-                # Confirm to sender with delivered status
                 await ws.send_json({
                     "type": "new_message",
                     "message": {**msg_obj, "status": "delivered"},
                     "client_id": client_id
                 })
 
-                # Broadcast to others
                 await manager.broadcast({
                     "type": "new_message",
                     "message": msg_obj
@@ -487,24 +615,26 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                         rows = await db.execute(
                             """SELECT m.*, u.username, u.display_name, u.avatar_path
                                FROM messages m JOIN users u ON m.sender_id = u.id
-                               WHERE m.is_deleted = 0 AND m.id < ? ORDER BY m.id DESC LIMIT ?""",
+                               WHERE m.is_deleted = 0 AND m.id < ?
+                               ORDER BY m.id DESC LIMIT ?""",
                             (cursor_id, limit)
                         )
                     else:
                         rows = await db.execute(
                             """SELECT m.*, u.username, u.display_name, u.avatar_path
                                FROM messages m JOIN users u ON m.sender_id = u.id
-                               WHERE m.is_deleted = 0 ORDER BY m.id DESC LIMIT ?""",
+                               WHERE m.is_deleted = 0
+                               ORDER BY m.id DESC LIMIT ?""",
                             (limit,)
                         )
+
                     msgs = []
                     for row in reversed(await rows.fetchall()):
                         try:
                             decrypted = decrypt_message(row['encrypted_content'])
-                        except:
+                        except Exception:
                             decrypted = "[Decryption Error]"
 
-                        # Determine read status
                         read_cursor = await db.execute(
                             "SELECT COUNT(*) as cnt FROM read_receipts WHERE message_id = ?",
                             (row['id'],)
@@ -552,7 +682,11 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                         (mid, uid)
                     ) as cursor:
                         if not await cursor.fetchone():
-                            await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
+                            await ws.send_json({
+                                "type": "error",
+                                "code": "NOT_FOUND",
+                                "message": "Message not found"
+                            })
                             continue
                     new_enc = encrypt_message(new_content)
                     await db.execute(
@@ -560,6 +694,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                         (new_enc, mid)
                     )
                     await db.commit()
+                    schedule_db_sync()
                     await manager.broadcast({
                         "type": "message_edited",
                         "message_id": mid,
@@ -582,16 +717,21 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     ) as cursor:
                         msg = await cursor.fetchone()
                     if not msg:
-                        await ws.send_json({"type": "error", "code": "NOT_FOUND", "message": "Message not found"})
+                        await ws.send_json({
+                            "type": "error",
+                            "code": "NOT_FOUND",
+                            "message": "Message not found"
+                        })
                         continue
                     if msg['file_path']:
                         try:
                             delete_file(msg['file_path'])
-                        except:
+                        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()
+                    schedule_db_sync()
                     await manager.broadcast({
                         "type": "message_deleted",
                         "message_id": mid
@@ -620,7 +760,9 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     rows = await db.execute(
                         """SELECT m.id, m.sender_id FROM messages m
                            WHERE m.id <= ? AND m.sender_id != ? AND m.is_deleted = 0
-                           AND m.id NOT IN (SELECT message_id FROM read_receipts WHERE user_id = ?)""",
+                           AND m.id NOT IN (
+                               SELECT message_id FROM read_receipts WHERE user_id = ?
+                           )""",
                         (up_to, uid, uid)
                     )
                     new_reads = []
@@ -631,6 +773,8 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                         )
                         new_reads.append(r)
                     await db.commit()
+                    if new_reads:
+                        schedule_db_sync()
                     for r in new_reads:
                         await manager.send_to_user(r['sender_id'], {
                             "type": "message_read",
@@ -660,6 +804,7 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 (int(time.time()), uid)
             )
             await db.commit()
+            schedule_db_sync()
         finally:
             await db.close()
         await manager.broadcast({