New storage scheme

cmd_chat/client/client.py

@@ -6,6 +6,8 @@ from typing import Optional
 import srp
 import requests
 from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from cryptography.hazmat.primitives import hashes
 import websockets
 from rich.console import Console
 from rich.panel import Panel
@@ -23,6 +25,7 @@ class Client:
         self.password = (password or "").encode()
         self.user_id: Optional[str] = None
         self.fernet: Optional[Fernet] = None
+        self.room_fernet: Optional[Fernet] = None
 
         self.console = Console()
         self.messages: list[dict] = []
@@ -67,6 +70,16 @@ class Client:
             self.user_id = init_data["user_id"]
             B = base64.b64decode(init_data["B"])
             salt = base64.b64decode(init_data["salt"])
+            room_salt = base64.b64decode(init_data["room_salt"])
+
+            hkdf = HKDF(
+                algorithm=hashes.SHA256(),
+                length=32,
+                salt=room_salt,
+                info=b"cmd-chat-room-key",
+            )
+            room_key = hkdf.derive(self.password)
+            self.room_fernet = Fernet(base64.urlsafe_b64encode(room_key))
 
             M = usr.process_challenge(salt, B)
 
@@ -96,6 +109,15 @@ class Client:
 
         self.success(f"SRP authenticated (session: {self.user_id[:8]}...)")
 
+    def decrypt_message(self, msg: dict) -> dict:
+        if "text" in msg and msg["text"]:
+            try:
+                decrypted = self.room_fernet.decrypt(msg["text"].encode()).decode()
+                msg["text"] = decrypted
+            except Exception:
+                msg["text"] = "[decrypt failed]"
+        return msg
+
     def render_messages(self) -> None:
         self.console.clear()
 
@@ -131,12 +153,15 @@ class Client:
                 msg_type = data.get("type", "")
 
                 if msg_type == "init":
-                    self.messages = data.get("messages", [])
+                    messages = [
+                        self.decrypt_message(m) for m in data.get("messages", [])
+                    ]
+                    self.messages = messages
                     self.users = data.get("users", [])
                     self.connected = True
                     self.render_messages()
                 elif msg_type == "message":
-                    msg_data = data.get("data", {})
+                    msg_data = self.decrypt_message(data.get("data", {}))
                     self.messages.append(msg_data)
                     self.render_messages()
                 elif msg_type == "user_left":
@@ -156,7 +181,8 @@ class Client:
                     self.running = False
                     break
                 if text.strip():
-                    await ws.send(text)
+                    encrypted = self.room_fernet.encrypt(text.encode()).decode()
+                    await ws.send(encrypted)
             except (EOFError, KeyboardInterrupt):
                 self.running = False
                 break

cmd_chat/client/config.py

@@ -1 +0,0 @@
-MESSAGES_TO_SHOW = 5

cmd_chat/server/factory.py

@@ -3,7 +3,7 @@ from contextlib import suppress
 from cryptography.fernet import Fernet
 from sanic import Sanic
 from sanic_ext import Extend
-
+import os
 from .managers import ConnectionManager
 from .stores import MessageStore, UserSessionStore
 from .srp_auth import SRPAuthManager
@@ -19,7 +19,7 @@ def create_app(password: str = "", name: str = "cmd-chat-server") -> Sanic:
     app.ctx.session_store = UserSessionStore()
     app.ctx.connection_manager = ConnectionManager()
     app.ctx.srp_manager = SRPAuthManager(password)
-    app.ctx.fernet_key = Fernet.generate_key()
+    app.ctx.room_salt = os.urandom(16)
     app.ctx.cleanup_task = None
 
     register_lifecycle(app)

cmd_chat/server/stores.py

@@ -9,7 +9,6 @@ class MessageStore:
     def add(self, message: Message) -> None:
         self._messages.append(message)
 
-
     def get_all(self) -> list[Message]:
         return self._messages.copy()
 
@@ -17,7 +16,6 @@ class MessageStore:
         count = len(self._messages)
         self._messages.clear()
 
-
     def count(self) -> int:
         return len(self._messages)
 
@@ -29,7 +27,6 @@ class UserSessionStore:
     def add(self, session: UserSession) -> None:
         self._sessions[session.user_id] = session
 
-
     def get(self, user_id: str) -> Optional[UserSession]:
         return self._sessions.get(user_id)
 
@@ -41,7 +38,6 @@ class UserSessionStore:
         if user_id in self._sessions:
             del self._sessions[user_id]
 
-
     def cleanup_stale(self, timeout_seconds: int = 3600) -> int:
         stale_ids = [
             uid for uid, s in self._sessions.items() if s.is_stale(timeout_seconds)

cmd_chat/server/views.py

@@ -35,6 +35,7 @@ async def srp_init(request: Request, app: Sanic) -> HTTPResponse:
                 "user_id": user_id,
                 "B": base64.b64encode(B).decode(),
                 "salt": base64.b64encode(salt).decode(),
+                "room_salt": base64.b64encode(app.ctx.room_salt).decode(),
             }
         )
 
@@ -66,7 +67,6 @@ async def srp_verify(request: Request, app: Sanic) -> HTTPResponse:
         )
         app.ctx.session_store.add(session)
 
-
         return response.json(
             {
                 "H_AMK": base64.b64encode(H_AMK).decode(),