Update app.py

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from typing import Optional, List
+from typing import Optional, List, Dict
 import httpx
 import json
 import logging
@@ -10,6 +10,7 @@ import asyncio
 import hashlib
 from datetime import datetime, timezone
 import re
+import fastapi
 
 app = FastAPI()
 logging.basicConfig(level=logging.INFO)
@@ -31,13 +32,24 @@ if not os.path.isdir(PERSISTENT_ROOT):
     CHAT_DIR = os.path.join(".", "data", "chat")
 os.makedirs(CHAT_DIR, exist_ok=True)
 
+USERS_DIR = os.path.join(PERSISTENT_ROOT, "users")
+os.makedirs(USERS_DIR, exist_ok=True)
+USERS_FILE = os.path.join(USERS_DIR, "users.jsonl")
+SESSIONS_FILE = os.path.join(USERS_DIR, "sessions.json")
+ADMIN_KEY = os.environ.get("ADMIN_KEY", "")  # set in HF Space secrets
+
+# Async file locks
+app.state.chat_locks = {}
+app.state.users_lock = asyncio.Lock()
+app.state.sessions_lock = asyncio.Lock()
+app.state.sessions: Dict[str, str] = {}
+
+
 def _chat_file_for(video_id: str) -> str:
     # Stable filename to avoid path traversal
     h = hashlib.sha256(video_id.encode("utf-8")).hexdigest()[:32]
     return os.path.join(CHAT_DIR, f"{h}.jsonl")
 
-# Async file locks
-app.state.chat_locks = {}
 
 def _lock_for(path: str) -> asyncio.Lock:
     lock = app.state.chat_locks.get(path)
@@ -46,44 +58,126 @@ def _lock_for(path: str) -> asyncio.Lock:
         app.state.chat_locks[path] = lock
     return lock
 
+
 def _now_iso() -> str:
     return datetime.now(timezone.utc).isoformat()
 
+
 def _valid_author(s: Optional[str]) -> str:
     s = (s or "anon").strip()
     s = re.sub(r"\s+", " ", s)
     return s[:32] or "anon"
 
+
 def _valid_text(s: str) -> str:
     s = (s or "").rstrip("\n")
     return s[:2000]
 
+
 async def _append_jsonl(path: str, record: dict) -> None:
     line = json.dumps(record, ensure_ascii=False) + "\n"
+
     def _write():
         with open(path, "a", encoding="utf-8") as f:
             f.write(line)
+
     await asyncio.to_thread(_write)
 
+
 async def _read_jsonl(path: str) -> List[dict]:
     if not os.path.exists(path):
         return []
+
     def _read():
         with open(path, "r", encoding="utf-8") as f:
             return [json.loads(x) for x in f if x.strip()]
+
     return await asyncio.to_thread(_read)
 
+
+# ---------------------------
+# 👤 Users & Auth helpers
+# ---------------------------
+
+def _safe_handle(s: str) -> str:
+    s = (s or "").strip()
+    s = re.sub(r"[^a-zA-Z0-9_.~-]", "_", s)
+    return s[:24]
+
+
+def _hash_pin(pin: str) -> str:
+    return hashlib.sha256(("tlks|" + pin).encode("utf-8")).hexdigest()
+
+
+def _rand_token() -> str:
+    return hashlib.sha256(os.urandom(32)).hexdigest()
+
+
+async def _append_user(record: dict) -> None:
+    line = json.dumps(record, ensure_ascii=False) + "\n"
+
+    def _write():
+        with open(USERS_FILE, "a", encoding="utf-8") as f:
+            f.write(line)
+
+    async with app.state.users_lock:
+        await asyncio.to_thread(_write)
+
+
+async def _read_users() -> List[dict]:
+    if not os.path.exists(USERS_FILE):
+        return []
+
+    def _read():
+        with open(USERS_FILE, "r", encoding="utf-8") as f:
+            return [json.loads(x) for x in f if x.strip()]
+
+    async with app.state.users_lock:
+        return await asyncio.to_thread(_read)
+
+
+async def _load_sessions():
+    if os.path.exists(SESSIONS_FILE):
+        def _read():
+            with open(SESSIONS_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+        app.state.sessions = await asyncio.to_thread(_read)
+    else:
+        app.state.sessions = {}
+
+
+async def _save_sessions():
+    def _write(data):
+        with open(SESSIONS_FILE, "w", encoding="utf-8") as f:
+            json.dump(data, f)
+
+    async with app.state.sessions_lock:
+        await asyncio.to_thread(_write, app.state.sessions)
+
+
+@app.on_event("startup")
+async def _startup_load_sessions():
+    await _load_sessions()
+
+
+def _require_admin(request: Request):
+    provided = request.headers.get("x-admin-key", "")
+    if not ADMIN_KEY or provided != ADMIN_KEY:
+        raise fastapi.HTTPException(status_code=401, detail="Admin key required")
+
+
 # ---------------------------
-# 💬 Chat API
+# 💬 Chat API (public read, authed write)
 # ---------------------------
 class NewMessage(BaseModel):
-    author: Optional[str] = Field(default="anon", max_length=64)
+    author: Optional[str] = Field(default="anon", max_length=64)  # ignored on server write
     text: str = Field(..., min_length=1, max_length=5000)
 
+
 @app.get("/chat/{video_id}")
 async def get_messages(video_id: str, limit: int = 50, since: Optional[str] = None):
     """
-    Fetch messages for a video.
+    Fetch messages for a room/video.
     - limit: max messages (default 50)
     - since: ISO8601 timestamp; return only messages newer than this
     """
@@ -107,16 +201,167 @@ async def get_messages(video_id: str, limit: int = 50, since: Optional[str] = No
 
     return {"video_id": video_id, "count": len(items), "messages": items}
 
+
+# ---------------------------
+# 👤 Users & Auth API
+# ---------------------------
+class NewUser(BaseModel):
+    email: str
+    first_name: str
+    last_name: str
+    handle: str = Field(..., min_length=2, max_length=24)
+    klass: str = Field(..., max_length=64)
+    profile_image: Optional[str] = ""
+    description: Optional[str] = ""
+    pin: str = Field(..., min_length=3, max_length=32)
+
+
+class UpdateUser(BaseModel):
+    first_name: Optional[str] = None
+    last_name: Optional[str] = None
+    klass: Optional[str] = None
+    profile_image: Optional[str] = None
+    description: Optional[str] = None
+    pin: Optional[str] = None
+    disabled: Optional[bool] = None
+
+
+class LoginReq(BaseModel):
+    handle: str
+    pin: str
+
+
+@app.post("/admin/users")
+async def admin_create_user(user: NewUser, request: Request):
+    _require_admin(request)
+    users = await _read_users()
+
+    handle = _safe_handle(user.handle.lstrip("@"))
+    if any(u.get("handle") == handle for u in users):
+        return {"ok": False, "error": "Handle already exists"}
+
+    rec = {
+        "email": user.email.strip(),
+        "first_name": user.first_name.strip(),
+        "last_name": user.last_name.strip(),
+        "handle": handle,
+        "klass": user.klass.strip(),
+        "profile_image": (user.profile_image or "").strip(),
+        "description": (user.description or "").strip(),
+        "pin_hash": _hash_pin(user.pin),
+        "disabled": False,
+        "created_at": _now_iso(),
+    }
+    await _append_user(rec)
+    pub = {k: v for k, v in rec.items() if k != "pin_hash"}
+    return {"ok": True, "user": pub}
+
+
+@app.get("/admin/users")
+async def admin_list_users(request: Request, q: Optional[str] = None, limit: int = 200):
+    _require_admin(request)
+    users = await _read_users()
+    if q:
+        ql = q.lower()
+        users = [u for u in users if ql in u.get("email", "").lower() or ql in u.get("handle", "").lower()]
+    users = users[: max(1, min(limit, 1000))]
+    for u in users:
+        u.pop("pin_hash", None)
+    return {"count": len(users), "users": users}
+
+
+@app.put("/admin/users/{handle}")
+async def admin_update_user(handle: str, patch: UpdateUser, request: Request):
+    _require_admin(request)
+    handle = _safe_handle(handle)
+    users = await _read_users()
+    changed = False
+    for u in users:
+        if u.get("handle") == handle:
+            if patch.first_name is not None:
+                u["first_name"] = patch.first_name.strip()
+            if patch.last_name is not None:
+                u["last_name"] = patch.last_name.strip()
+            if patch.klass is not None:
+                u["klass"] = patch.klass.strip()
+            if patch.profile_image is not None:
+                u["profile_image"] = patch.profile_image.strip()
+            if patch.description is not None:
+                u["description"] = patch.description.strip()
+            if patch.disabled is not None:
+                u["disabled"] = bool(patch.disabled)
+            if patch.pin is not None:
+                u["pin_hash"] = _hash_pin(patch.pin)
+            changed = True
+            break
+    if not changed:
+        raise fastapi.HTTPException(status_code=404, detail="User not found")
+
+    # rewrite file
+    async with app.state.users_lock:
+        def _write_all():
+            with open(USERS_FILE, "w", encoding="utf-8") as f:
+                for rec in users:
+                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        await asyncio.to_thread(_write_all)
+    return {"ok": True}
+
+
+@app.post("/auth/login")
+async def login(req: LoginReq):
+    users = await _read_users()
+    handle = _safe_handle(req.handle.lstrip("@"))
+    u = next((x for x in users if x.get("handle") == handle), None)
+    if not u or u.get("disabled"):
+        return {"ok": False, "error": "Invalid user"}
+    if u.get("pin_hash") != _hash_pin(req.pin):
+        return {"ok": False, "error": "Invalid PIN"}
+
+    token = _rand_token()
+    app.state.sessions[token] = handle
+    await _save_sessions()
+    return {"ok": True, "token": token}
+
+
+async def _require_user(request: Request) -> dict:
+    auth = request.headers.get("authorization", "")
+    if not auth.lower().startswith("bearer "):
+        raise fastapi.HTTPException(status_code=401, detail="Missing bearer token")
+    token = auth.split(" ", 1)[1].strip()
+    handle = app.state.sessions.get(token)
+    if not handle:
+        raise fastapi.HTTPException(status_code=401, detail="Invalid session")
+
+    users = await _read_users()
+    user = next((u for u in users if u.get("handle") == handle), None)
+    if not user or user.get("disabled"):
+        raise fastapi.HTTPException(status_code=401, detail="User disabled")
+    return user
+
+
+@app.get("/me")
+async def me(request: Request):
+    user = await _require_user(request)
+    pub = {k: v for k, v in user.items() if k not in ("pin_hash",)}
+    return {"ok": True, "user": pub}
+
+
 @app.post("/chat/{video_id}")
 async def post_message(video_id: str, msg: NewMessage, request: Request):
     """
-    Append a message to a video's chat.
-    Body: { "author": "Ross", "text": "hello" }
+    Append a message to a room's chat.
+    Body: { "text": "hello" }
     """
     path = _chat_file_for(video_id)
     lock = _lock_for(path)
 
-    author = _valid_author(msg.author)
+    # Require a valid session and take author from the user profile
+    user = await _require_user(request)
+    author = _valid_author(f"{user.get('first_name','')} {user.get('last_name','')}".strip() or user["handle"])
+    handle = user["handle"]
+    klass = user.get("klass", "")
+    profile_image = user.get("profile_image", "")
+
     text = _valid_text(msg.text)
     if not text:
         return {"ok": False, "error": "Empty message"}
@@ -128,11 +373,14 @@ async def post_message(video_id: str, msg: NewMessage, request: Request):
     record = {
         "id": mid,
         "video_id": video_id,
-        "author": author,  # ✅ FIXED: uses frontend-provided name
+        "author": author,
+        "handle": handle,
+        "klass": klass,
+        "profile_image": profile_image,
         "text": text,
         "created_at": created,
         "ip": ip,
-        "ua": request.headers.get("user-agent", "")[:200]
+        "ua": request.headers.get("user-agent", "")[:200],
     }
 
     async with lock:
@@ -140,22 +388,26 @@ async def post_message(video_id: str, msg: NewMessage, request: Request):
 
     return {"ok": True, "message": record}
 
+
 # ---------------------------
 # (Original iCloud album endpoints)
 # ---------------------------
 BASE_62_MAP = {c: i for i, c in enumerate("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")}
 
+
 async def get_client() -> httpx.AsyncClient:
     if not hasattr(app.state, "client"):
         app.state.client = httpx.AsyncClient(timeout=15.0)
     return app.state.client
 
+
 def base62_to_int(token: str) -> int:
     result = 0
     for ch in token:
         result = result * 62 + BASE_62_MAP[ch]
     return result
 
+
 async def get_base_url(token: str) -> str:
     first = token[0]
     if first == "A":
@@ -164,12 +416,14 @@ async def get_base_url(token: str) -> str:
         n = base62_to_int(token[1:3])
     return f"https://p{n:02d}-sharedstreams.icloud.com/{token}/sharedstreams/"
 
+
 ICLOUD_HEADERS = {
     "Origin": "https://www.icloud.com",
-    "Content-Type": "text/plain"
+    "Content-Type": "text/plain",
 }
 ICLOUD_PAYLOAD = '{"streamCtag":null}'
 
+
 async def get_redirected_base_url(base_url: str, token: str) -> str:
     client = await get_client()
     resp = await client.post(
@@ -191,21 +445,25 @@ async def get_redirected_base_url(base_url: str, token: str) -> str:
     else:
         resp.raise_for_status()
 
+
 async def post_json(path: str, base_url: str, payload: str) -> dict:
     client = await get_client()
     resp = await client.post(f"{base_url}{path}", headers=ICLOUD_HEADERS, data=payload)
     resp.raise_for_status()
     return resp.json()
 
+
 async def get_metadata(base_url: str) -> list:
     data = await post_json("webstream", base_url, ICLOUD_PAYLOAD)
     return data.get("photos", [])
 
+
 async def get_asset_urls(base_url: str, guids: list) -> dict:
     payload = json.dumps({"photoGuids": guids})
     data = await post_json("webasseturls", base_url, payload)
     return data.get("items", {})
 
+
 @app.get("/album/{token}")
 async def get_album(token: str):
     try:
@@ -225,7 +483,7 @@ async def get_album(token: str):
             best = max(
                 (d for k, d in derivatives.items() if k.lower() != "posterframe"),
                 key=lambda d: int(d.get("fileSize") or 0),
-                default=None
+                default=None,
             )
             if not best:
                 continue
@@ -246,7 +504,7 @@ async def get_album(token: str):
             videos.append({
                 "caption": photo.get("caption", ""),
                 "url": video_url,
-                "poster": poster or ""
+                "poster": poster or "",
             })
 
         return {"videos": videos}
@@ -255,6 +513,7 @@ async def get_album(token: str):
         logging.exception("Error in get_album")
         return {"error": str(e)}
 
+
 @app.get("/album/{token}/raw")
 async def get_album_raw(token: str):
     try:
@@ -266,4 +525,4 @@ async def get_album_raw(token: str):
         return {"metadata": metadata, "asset_urls": asset_map}
     except Exception as e:
         logging.exception("Error in get_album_raw")
-        return {"error": str(e)}
+        return {"error": str(e)}