""" A node in memory, at the level of the MNP messages the CLI sends. It does the real cryptography — chunks encrypted under the per-chunk key, uploads unsealed under the group key, acks sealed back — and keeps the rules that shape the client: a lease before any chunk, chunks of an upload in order, the position an interrupted upload resumes from, no overwrite. Not a model of the node's code: a node the CLI must work against. """ import asyncio import hashlib import math import secrets from meshbay_cli.client.session import GroupSession, NodeRefusal from meshbay_common.groupbox import PURPOSE_INDEX, PURPOSE_UPLOAD, seal, unseal from meshbay_common.protocol import UPLOAD_PROBE_INDEX, file_upload_ack_wire from meshbay_common.webcrypto import chunk_key_aes, encrypt_chunk_aes GROUP_ID = "g1" CHUNK = 1024 * 1024 class FakeNode: def __init__(self, *, writable: bool = True): self.gek = secrets.token_bytes(32) self.roots = [{"name": "share", "available": True, "writable": writable}] self.dirs = {"share"} self.files: dict[str, dict] = {} # id → entry, with "data" self.leases: dict[str, str] = {} # tr → state self.queue_first = 0 # leases to answer "queued" before granting self.uploads: dict[tuple, dict] = {} # (dir, name) → {"next", "data"} self.fail_chunk: tuple[str, int] | None = None # (file_id, index) → refused self.refuse_upload_code = "" self.chunk_requests = 0 self.chunk_delay = 0.0 def add(self, path: str, name: str, data: bytes) -> dict: fid = hashlib.blake2b(data + path.encode() + name.encode(), digest_size=32).hexdigest() entry = {"id": fid, "name": name, "path": path, "size": len(data), "type": "document", "added_at": 1_700_000_000, "data": data} self.files[fid] = entry parts = path.split("/") for i in range(1, len(parts) + 1): self.dirs.add("/".join(parts[:i])) return entry def session(self) -> GroupSession: return GroupSession(None, FakeChannel(self), GROUP_ID, "node-1", "pk", self.gek, {}) def stored(self, path: str) -> bytes | None: parent, _, name = path.rpartition("/") for e in self.files.values(): if e["path"] == parent and e["name"] == name: return e["data"] return None class FakeChannel: def __init__(self, node: FakeNode): self.node = node self.closed = asyncio.Event() self._push: dict[str, list] = {} def on_push(self, t, h): self._push.setdefault(t, []).append(h) def off_push(self, t, h): if h in self._push.get(t, []): self._push[t].remove(h) def _emit(self, msg): for h in list(self._push.get(msg["type"], [])): h(msg) async def request(self, msg: dict, *expect: str, timeout: float = 60) -> dict: await asyncio.sleep(self.node.chunk_delay if msg["type"] == "file_req" else 0) reply = getattr(self, "_" + msg["type"])(msg) if reply.get("type") == "error": raise NodeRefusal(reply["detail"], reply.get("code", "")) return reply # ── Messages ───────────────────────────────────────────────────────────── def _index_sync(self, msg): n = self.node payload = {"version": 1, "roots": n.roots, "dirs": sorted(n.dirs), "entries": [{k: v for k, v in e.items() if k != "data"} for e in n.files.values()]} return {"type": "index_sync", "group_id": GROUP_ID, **seal(n.gek, PURPOSE_INDEX, "index_sync", GROUP_ID, payload)} def _transfer_open(self, msg): tr = msg["tr"] if self.node.queue_first: self.node.queue_first -= 1 self.node.leases[tr] = "queued" asyncio.get_running_loop().call_later(0.01, self._grant, tr) return {"type": "transfer_state", "tr": tr, "state": "queued", "ahead": 1} self.node.leases[tr] = "granted" return {"type": "transfer_state", "tr": tr, "state": "granted"} def _grant(self, tr): self.node.leases[tr] = "granted" self._emit({"type": "transfer_state", "tr": tr, "state": "granted"}) def _transfer_close(self, msg): self.node.leases.pop(msg["tr"], None) return {"type": "transfer_state", "tr": msg["tr"], "state": "closed"} def _file_req(self, msg): n = self.node n.chunk_requests += 1 if n.leases.get(msg.get("tr")) != "granted": return {"type": "error", "detail": "no lease", "code": "lease_not_granted"} if n.fail_chunk == (msg["file_id"], msg["chunk_index"]): return {"type": "error", "detail": "The drive stopped answering"} entry = n.files[msg["file_id"]] i = msg["chunk_index"] data = entry["data"][i * CHUNK:(i + 1) * CHUNK] nonce, ct = encrypt_chunk_aes(chunk_key_aes(n.gek, bytes.fromhex(entry["id"]), i), data) return {"type": "file_chunk", "chunk_index": i, "nonce": nonce, "ct": ct} def _dir_create(self, msg): root = self.node.roots[0] if not root["writable"]: return {"type": "error", "detail": "read-only", "code": "root_read_only"} path = f"{msg['dir']}/{msg['name']}" self.node.dirs.add(path) return {"type": "dir_create_ack", "dir": path} def _file_upload(self, msg): n = self.node if n.leases.get(msg.get("tr")) != "granted": return {"type": "error", "detail": "no lease", "code": "lease_not_granted"} p = unseal(n.gek, PURPOSE_UPLOAD, "file_upload", GROUP_ID, msg) name, folder, i = p["filename"], p["dir"], msg["chunk_index"] if n.refuse_upload_code: return {"type": "error", "detail": "refused", "code": n.refuse_upload_code} key = (folder, name) state = n.uploads.get(key) stored = name while n.stored(f"{folder}/{stored}") is not None: stem, dot, ext = name.partition(".") stored = f"{stem} (1){dot}{ext}" if i == UPLOAD_PROBE_INDEX: return file_upload_ack_wire(n.gek, GROUP_ID, upload_id=msg["upload_id"], chunk_index=i, filename=name, stored_as="", dir=folder, resume_from=state["next"] if state else 0) if state is None: state = n.uploads[key] = {"next": 0, "data": b""} if i != state["next"]: return {"type": "error", "detail": "Out of order", "code": "bad_chunk_index"} state["data"] += p["data"] state["next"] += 1 if state["next"] >= msg["total_chunks"]: del n.uploads[key] n.add(folder, stored, state["data"]) return file_upload_ack_wire(n.gek, GROUP_ID, upload_id=msg["upload_id"], chunk_index=i, filename=name, stored_as=stored, dir=folder) def chunks_of(size: int, chunk: int) -> int: return max(1, math.ceil(size / chunk))