websocket_client.py (18925B)
1 """ 2 Obsidian Sync WebSocket client. 3 4 Reversed from Obsidian 1.11.7 app.js (N9 class). 5 6 Protocol: 7 1. Connect to wss://<host> (host comes from vault.host in API response) 8 2. Send init message: 9 {"op":"init", "token":<token>, "id":<vault_id>, "keyhash":<keyhash>, 10 "version":<last_uid>, "initial":<bool>, "device":<device_name>, 11 "encryption_version":<int>} 12 3. Receive response: {"res":"ok", "version":<uid>, "userId":<int>, "perFileMax":<int>} 13 OR {"res":"err", ...} / {"status":"err", ...} 14 4. After successful init, server sends "ready" message: 15 {"op":"ready", "version":<uid>} 16 5. Server may push file updates: {"op":"push", "path":<enc_path>, "hash":<hash>, ...} 17 18 Operations (sent as JSON): 19 pull: {"op":"pull", "uid":<int>} 20 -> response: {"size":<int>, "pieces":<int>, "deleted":<bool>, ...} 21 -> then <pieces> binary ArrayBuffer messages with the file data 22 23 push: {"op":"push", "path":<enc>, "relatedpath":<enc|null>, "extension":<str>, 24 "hash":<hash>, "ctime":<int>, "mtime":<int>, "folder":<bool>, 25 "deleted":<bool>, "size":<int>, "pieces":<int>} 26 -> response: {"res":"ok"} or chunk acknowledgements 27 -> then send <pieces> binary chunks of 2MB each, waiting for {"res":"ok"} after each 28 29 deleted: {"op":"deleted", "suppressrenames":true} -> {items:[{uid, path, ...}]} 30 history: {"op":"history", "path":<enc>, "last":<uid>} -> {items:[...]} 31 restore: {"op":"restore", "uid":<int>} -> {res:"ok"} 32 usernames:{"op":"usernames"} -> {usernames:[...]} 33 purge: {"op":"purge"} -> {res:"ok"} 34 size: {"op":"size"} -> {size:<int>, limit:<int>, vault_size:<int>} 35 ping: {"op":"ping"} -> {"op":"pong"} 36 37 File data in pull/push is encrypted (encrypt/decrypt via encryption provider). 38 File paths are encoded deterministically (deterministicEncodeStr/deterministicDecodeStr). 39 """ 40 41 from __future__ import annotations 42 43 import asyncio 44 import json 45 import hashlib 46 import os 47 from typing import Optional, Callable, Any 48 from websockets.asyncio.client import connect as ws_connect 49 from websockets.protocol import State as WsState 50 from websockets.exceptions import ConnectionClosed 51 52 from .crypto import EncryptionV0, EncryptionV2, create_encryption_provider 53 54 CHUNK_SIZE = 2 * 1024 * 1024 # 2MB chunks per push 55 HEARTBEAT_INTERVAL = 20 # seconds 56 HEARTBEAT_TIMEOUT = 120 # disconnect if no message for this long 57 REQUEST_TIMEOUT = 60 # seconds 58 59 60 def _get_ws_url(host: str) -> str: 61 """Convert host to WebSocket URL. Matches getHost() in app.js.""" 62 if host.startswith("127.0.0.1") or host.startswith("localhost"): 63 return f"ws://{host}" 64 return f"wss://{host}" 65 66 67 def _validate_host(url: str) -> bool: 68 """Validate that the host ends with .obsidian.md or is localhost. 69 Matches the security check in connect() in app.js. 70 """ 71 from urllib.parse import urlparse 72 parsed = urlparse(url) 73 hostname = parsed.hostname or "" 74 return hostname.endswith(".obsidian.md") or hostname == "127.0.0.1" 75 76 77 class SyncError(Exception): 78 pass 79 80 81 class SyncConnectionError(SyncError): 82 pass 83 84 85 class SyncAuthError(SyncError): 86 pass 87 88 89 class ObsidianSyncClient: 90 """ 91 WebSocket client for Obsidian Sync. 92 93 Manages connection, authentication, and file sync operations. 94 """ 95 96 def __init__( 97 self, 98 host: str, 99 token: str, 100 vault_id: str, 101 encryption_provider: EncryptionV0 | EncryptionV2, 102 device_name: str = "obsidian-sync-client", 103 last_version: int = 0, 104 initial: bool = True, 105 ): 106 self.host = host 107 self.token = token 108 self.vault_id = vault_id 109 self.enc = encryption_provider 110 self.device_name = device_name 111 self.last_version = last_version 112 self.initial = initial 113 114 self._ws = None # websockets connection 115 self._response_queue: asyncio.Queue = asyncio.Queue() 116 self._data_queue: asyncio.Queue = asyncio.Queue() 117 self._push_callbacks: list[Callable] = [] 118 self._ready_callbacks: list[Callable] = [] 119 self._running = False 120 self._receive_task: Optional[asyncio.Task] = None 121 self.user_id: int = -1 122 self.per_file_max: int = 200 * 1024 * 1024 # default 200MB 123 self.last_message_ts: float = 0 124 125 # ------------------------------------------------------------------ 126 # Connection management 127 # ------------------------------------------------------------------ 128 129 async def connect(self) -> None: 130 """Connect and authenticate with the sync server.""" 131 url = _get_ws_url(self.host) 132 if not _validate_host(url): 133 raise SyncConnectionError(f"Invalid sync host: {self.host}") 134 135 try: 136 self._ws = await ws_connect( 137 url, 138 ping_interval=None, # We handle heartbeats manually 139 close_timeout=10, 140 max_size=256 * 1024 * 1024, # 256MB max message 141 ) 142 except Exception as e: 143 raise SyncConnectionError(f"Unable to connect to server: {e}") from e 144 145 # Send init message 146 init_msg = { 147 "op": "init", 148 "token": self.token, 149 "id": self.vault_id, 150 "keyhash": self.enc.key_hash, 151 "version": self.last_version, 152 "initial": self.initial, 153 "device": self.device_name, 154 "encryption_version": self.enc.encryption_version, 155 } 156 await self._ws.send(json.dumps(init_msg)) 157 158 # Wait for auth response 159 try: 160 raw = await asyncio.wait_for(self._ws.recv(), timeout=REQUEST_TIMEOUT) 161 except asyncio.TimeoutError: 162 raise SyncConnectionError("Timeout waiting for auth response") 163 164 if not isinstance(raw, str): 165 raise SyncConnectionError("Server returned binary instead of JSON auth response") 166 167 try: 168 msg = json.loads(raw) 169 except json.JSONDecodeError as e: 170 raise SyncConnectionError(f"Server JSON failed to parse: {raw}") from e 171 172 if msg.get("op") == "pong": 173 # Unexpected pong, try again 174 raw = await asyncio.wait_for(self._ws.recv(), timeout=REQUEST_TIMEOUT) 175 msg = json.loads(raw) 176 177 if msg.get("status") == "err" or msg.get("res") == "err": 178 err = msg.get("msg", "Authentication failed") 179 if "expired" in err: 180 raise SyncAuthError(f"Subscription expired: {err}") 181 raise SyncAuthError(f"Failed to authenticate: {err}") 182 183 if msg.get("res") != "ok": 184 raise SyncConnectionError(f"Did not respond to login request: {raw}") 185 186 # Parse server capabilities 187 if "perFileMax" in msg: 188 val = msg["perFileMax"] 189 if isinstance(val, int) and val >= 0: 190 self.per_file_max = val 191 192 if msg.get("userId"): 193 self.user_id = msg["userId"] 194 195 # Start background receive loop 196 self._running = True 197 self._receive_task = asyncio.create_task(self._receive_loop()) 198 199 async def disconnect(self) -> None: 200 """Disconnect from the server.""" 201 self._running = False 202 if self._receive_task: 203 self._receive_task.cancel() 204 try: 205 await self._receive_task 206 except asyncio.CancelledError: 207 pass 208 self._receive_task = None 209 if self._ws: 210 await self._ws.close() 211 self._ws = None 212 # Drain queues with disconnect errors 213 error = SyncConnectionError("Disconnected") 214 try: 215 while True: 216 future = self._response_queue.get_nowait() 217 if hasattr(future, 'set_exception'): 218 future.set_exception(error) 219 except asyncio.QueueEmpty: 220 pass 221 222 def is_connected(self) -> bool: 223 return self._ws is not None and self._ws.state == WsState.OPEN 224 225 # ------------------------------------------------------------------ 226 # Internal message handling 227 # ------------------------------------------------------------------ 228 229 async def _receive_loop(self) -> None: 230 """Background task to receive and dispatch messages.""" 231 import time 232 if self._ws is None: 233 return 234 try: 235 async for raw in self._ws: 236 self.last_message_ts = time.time() 237 if isinstance(raw, str): 238 try: 239 msg = json.loads(raw) 240 except json.JSONDecodeError: 241 continue 242 243 op = msg.get("op") 244 if op == "pong": 245 continue 246 elif op == "ready": 247 version = msg.get("version", 0) 248 for cb in self._ready_callbacks: 249 asyncio.create_task(cb(version)) 250 elif op == "push": 251 # Server is pushing a file update to us 252 asyncio.create_task(self._handle_server_push(msg)) 253 else: 254 # Response to a request we made 255 await self._response_queue.put(msg) 256 else: 257 # Binary data 258 await self._data_queue.put(raw) 259 except ConnectionClosed: 260 pass 261 except Exception: 262 pass 263 finally: 264 self._running = False 265 266 async def _handle_server_push(self, msg: dict) -> None: 267 """Handle a server-pushed file update.""" 268 # Decode encrypted path 269 try: 270 path = await asyncio.get_event_loop().run_in_executor( 271 None, self.enc.deterministic_decode_str, msg.get("path", "") 272 ) 273 msg["path"] = path 274 275 if msg.get("hash"): 276 decoded_hash = await asyncio.get_event_loop().run_in_executor( 277 None, self.enc.deterministic_decode_str, msg["hash"] 278 ) 279 msg["hash"] = decoded_hash 280 except Exception: 281 pass # Best-effort decoding 282 283 for cb in self._push_callbacks: 284 try: 285 await cb(msg) 286 except Exception: 287 pass 288 289 async def _get_response(self, timeout: float = REQUEST_TIMEOUT) -> dict: 290 try: 291 return await asyncio.wait_for(self._response_queue.get(), timeout=timeout) 292 except asyncio.TimeoutError: 293 await self.disconnect() 294 raise SyncError("Request timed out") 295 296 async def _get_data(self, timeout: float = REQUEST_TIMEOUT) -> bytes: 297 try: 298 return await asyncio.wait_for(self._data_queue.get(), timeout=timeout) 299 except asyncio.TimeoutError: 300 await self.disconnect() 301 raise SyncError("Data receive timed out") 302 303 async def _send(self, msg: dict) -> None: 304 if not self._ws: 305 raise SyncConnectionError("Not connected") 306 await self._ws.send(json.dumps(msg)) 307 308 async def _send_binary(self, data: bytes) -> None: 309 if not self._ws: 310 raise SyncConnectionError("Not connected") 311 await self._ws.send(data) 312 313 async def _send_ping(self) -> None: 314 await self._send({"op": "ping"}) 315 316 # ------------------------------------------------------------------ 317 # Public API 318 # ------------------------------------------------------------------ 319 320 def on_push(self, callback: Callable) -> None: 321 """Register callback for server-pushed file updates.""" 322 self._push_callbacks.append(callback) 323 324 def on_ready(self, callback: Callable) -> None: 325 """Register callback for when server sends 'ready' (after initial sync version exchange).""" 326 self._ready_callbacks.append(callback) 327 328 async def pull(self, uid: int) -> Optional[bytes]: 329 """Download a file version by UID. Returns decrypted content or None if deleted.""" 330 await self._send({"op": "pull", "uid": uid}) 331 response = await self._get_response() 332 333 if response.get("deleted"): 334 return None 335 336 size = response.get("size", 0) 337 pieces = response.get("pieces", 0) 338 339 if size == 0 or pieces == 0: 340 return b"" 341 342 # Receive all binary pieces 343 buffer = bytearray(size) 344 offset = 0 345 for _ in range(pieces): 346 chunk = await self._get_data() 347 n = len(chunk) 348 buffer[offset:offset + n] = chunk 349 offset += n 350 351 encrypted = bytes(buffer) 352 if len(encrypted) == 0: 353 return b"" 354 355 # Decrypt 356 loop = asyncio.get_event_loop() 357 return await loop.run_in_executor(None, self.enc.decrypt, encrypted) 358 359 async def push( 360 self, 361 path: str, 362 data: Optional[bytes], 363 related_path: Optional[str] = None, 364 folder: bool = False, 365 deleted: bool = False, 366 ctime: int = 0, 367 mtime: int = 0, 368 file_hash: str = "", 369 ) -> None: 370 """Upload or delete a file. 371 372 - For folders or deletions, data should be None. 373 - file_hash should be hex(SHA-256(raw_data)) before encryption. 374 """ 375 loop = asyncio.get_event_loop() 376 377 # Encode path deterministically 378 enc_path = await loop.run_in_executor( 379 None, self.enc.deterministic_encode_str, path 380 ) 381 382 enc_related = None 383 if related_path: 384 enc_related = await loop.run_in_executor( 385 None, self.enc.deterministic_encode_str, related_path 386 ) 387 388 # Get extension 389 if not folder: 390 ext = path.rsplit(".", 1)[-1] if "." in path.split("/")[-1] else "" 391 else: 392 ext = "" 393 394 if folder or deleted or data is None: 395 # Folder or delete - no content 396 await self._send({ 397 "op": "push", 398 "path": enc_path, 399 "relatedpath": enc_related, 400 "extension": ext, 401 "hash": "", 402 "ctime": 0, 403 "mtime": 0, 404 "folder": folder, 405 "deleted": deleted, 406 }) 407 await self._get_response() 408 return 409 410 # Encrypt content 411 encrypted = await loop.run_in_executor(None, self.enc.encrypt, data) 412 413 # Encode hash 414 enc_hash = "" 415 if file_hash: 416 enc_hash = await loop.run_in_executor( 417 None, self.enc.deterministic_encode_str, file_hash 418 ) 419 420 size = len(encrypted) 421 pieces = max(1, (size + CHUNK_SIZE - 1) // CHUNK_SIZE) if size > 0 else 1 422 423 await self._send({ 424 "op": "push", 425 "path": enc_path, 426 "relatedpath": enc_related, 427 "extension": ext, 428 "hash": enc_hash, 429 "ctime": ctime, 430 "mtime": mtime, 431 "folder": False, 432 "deleted": False, 433 "size": size, 434 "pieces": pieces, 435 }) 436 437 response = await self._get_response() 438 439 # If server says ok already, we're done 440 if response.get("res") == "ok": 441 return 442 443 # Send binary chunks 444 for i in range(pieces): 445 start = i * CHUNK_SIZE 446 end = min(start + CHUNK_SIZE, size) 447 chunk = encrypted[start:end] 448 await self._send_binary(bytes(chunk)) 449 await self._get_response() # Wait for ack per chunk 450 451 async def list_deleted(self) -> list[dict]: 452 """List deleted/renamed files on the server.""" 453 await self._send({"op": "deleted", "suppressrenames": True}) 454 response = await self._get_response() 455 items = response.get("items", []) 456 loop = asyncio.get_event_loop() 457 for item in items: 458 try: 459 item["path"] = await loop.run_in_executor( 460 None, self.enc.deterministic_decode_str, item["path"] 461 ) 462 except Exception: 463 pass 464 return items 465 466 async def list_history(self, path: str, last: int = 0) -> dict: 467 """List version history for a file.""" 468 loop = asyncio.get_event_loop() 469 enc_path = await loop.run_in_executor( 470 None, self.enc.deterministic_encode_str, path 471 ) 472 await self._send({"op": "history", "path": enc_path, "last": last}) 473 response = await self._get_response() 474 475 for item in response.get("items", []): 476 try: 477 item["path"] = await loop.run_in_executor( 478 None, self.enc.deterministic_decode_str, item.get("path", "") 479 ) 480 except Exception: 481 pass 482 if item.get("relatedpath"): 483 try: 484 item["relatedpath"] = await loop.run_in_executor( 485 None, self.enc.deterministic_decode_str, item["relatedpath"] 486 ) 487 except Exception: 488 pass 489 return response 490 491 async def restore(self, uid: int) -> None: 492 """Restore a file to a previous version by UID.""" 493 await self._send({"op": "restore", "uid": uid}) 494 await self._get_response() 495 496 async def get_usernames(self) -> dict: 497 """Get usernames of vault collaborators.""" 498 await self._send({"op": "usernames"}) 499 return await self._get_response() 500 501 async def purge(self) -> None: 502 """Purge version history.""" 503 await self._send({"op": "purge"}) 504 await self._get_response() 505 506 async def get_size(self) -> dict: 507 """Get storage usage. Returns {size, limit, vault_size}.""" 508 await self._send({"op": "size"}) 509 response = await self._get_response() 510 return { 511 "size": response.get("size", -1), 512 "limit": response.get("limit", -1), 513 "vault_size": response.get("vault_size", -1), 514 } 515 516 async def heartbeat_loop(self) -> None: 517 """Run periodic heartbeats (call as background task).""" 518 import time 519 while self._running and self.is_connected(): 520 await asyncio.sleep(HEARTBEAT_INTERVAL) 521 if not self.is_connected(): 522 break 523 elapsed = time.time() - self.last_message_ts 524 if elapsed > HEARTBEAT_TIMEOUT: 525 await self.disconnect() 526 break 527 elif elapsed > 10: 528 try: 529 await self._send_ping() 530 except Exception: 531 break