path: root/packages/meshbay-cli/tests/fake_node.py

"""
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 base64
import hashlib
import math
import secrets

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from fake_hub import USER_ID
from meshbay_cli.client.keys import Identity
from meshbay_cli.client.session import GroupSession, NodeRefusal
from meshbay_common.adminop import OP_FILE_DELETE, admin_transcript
from meshbay_common.device import (
    device_add_transcript,
    device_code_hash,
    device_revoke_transcript,
)
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"
NODE_PK = base64.b64encode(b"\x05" * 32).decode()
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
        # This machine's key here, the account's other devices, and the requests
        # waiting for an approval — each `{pk_ed25519, pk_x25519, code_hash}`.
        self.identity = Identity(Ed25519PrivateKey.generate(), X25519PrivateKey.generate())
        self.nonce = secrets.token_bytes(32)
        self.devices = [{"pk_ed25519": self.identity.pk_ed, "pk_x25519": self.identity.pk_x,
                         "label": "", "pinned_at": "2026-10-01T10:00:00+00:00",
                         "pinned_via": "code"}]
        self.requests: list[dict] = []
        self.role = ""                       # "operator": this key was paired here
        self.ack: dict = {}                  # the sealed configuration, opened
        self.admin_ops: dict[str, dict] = {}

    def add(self, path: str, name: str, data: bytes, uploader_id: str = "") -> 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,
                 "uploader_id": uploader_id}
        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", NODE_PK, self.gek,
                            dict(self.ack), self.identity, self.nonce, self.role)

    def request_device(self, code: str) -> Identity:
        """Another machine files a request with `code`, as it would before its proof."""
        other = Identity(Ed25519PrivateKey.generate(), X25519PrivateKey.generate())
        self.requests.append({"pk_ed25519": other.pk_ed, "pk_x25519": other.pk_x,
                              "created_at": "2026-10-14T10:00:00+00:00",
                              "code_hash": device_code_hash(code.replace("-", ""),
                                                            other.pk_ed, other.pk_x)})
        return other

    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)


    # ── Deleting: a challenge, signed with this machine's key ──────────────

    def _file_delete(self, msg):
        n = self.node
        if msg["file_id"] not in n.files:
            return {"type": "error", "detail": "File not found"}
        op_id, nonce = secrets.token_hex(8), secrets.token_bytes(32)
        n.admin_ops[op_id] = {"subject": msg["file_id"], "nonce": nonce, "ts": 1_700_000_000}
        return {"type": "admin_challenge", "op_id": op_id, "op": OP_FILE_DELETE,
                "subject": msg["file_id"], "nonce": base64.b64encode(nonce).decode(),
                "ts": 1_700_000_000, "node_pk": NODE_PK, "group_id": GROUP_ID}

    def _admin_response(self, msg):
        n = self.node
        op = n.admin_ops.pop(msg["op_id"], None)
        if op is None:
            return {"type": "error", "detail": "No pending admin operation"}
        entry = n.files.get(op["subject"])
        transcript = admin_transcript(op=OP_FILE_DELETE, node_pk_b64=NODE_PK,
                                      group_id=GROUP_ID, subject=op["subject"],
                                      nonce=op["nonce"], ts=op["ts"])
        signed = _signed_by_a_device(n, transcript, msg["signature"])
        allowed = n.role == "operator" or (entry and entry["uploader_id"] == USER_ID)
        if not (signed and allowed):
            return {"type": "error", "detail": "Signature verification failed"}
        del n.files[op["subject"]]
        return {"type": "file_delete_ack", "file_id": op["subject"]}

    # ── Devices ──────────────────────────────────────────────────────────────

    def _device_list(self, msg):
        n = self.node
        return {"type": "device_list_result", "pending": len(n.requests),
                "devices": [{**d, "is_this_one": d["pk_ed25519"] == n.identity.pk_ed}
                            for d in n.devices]}

    def _device_lookup(self, msg):
        return {"type": "device_lookup_result", "requests": list(self.node.requests)}

    def _device_add(self, msg):
        n = self.node
        transcript = device_add_transcript(
            node_pk_b64=NODE_PK, user_id=USER_ID, pk_ed25519_b64=msg["pk_ed25519"],
            pk_x25519_b64=msg["pk_x25519"], nonce_node=n.nonce, ts=msg["ts"])
        if not _signed_by_a_device(n, transcript, msg["sig"]):
            return {"type": "error", "detail": "Not signed by a device already paired here"}
        req = next((r for r in n.requests if r["code_hash"] == msg["code_hash"]
                    and r["pk_ed25519"] == msg["pk_ed25519"]), None)
        if req is None:
            return {"type": "error", "detail": "No pending request"}
        n.requests.remove(req)
        n.devices.append({"pk_ed25519": req["pk_ed25519"], "pk_x25519": req["pk_x25519"],
                          "label": "", "pinned_at": "2026-10-14T10:01:00+00:00",
                          "pinned_via": "device"})
        return {"type": "device_add_ack", "pk_ed25519": req["pk_ed25519"]}

    def _device_revoke(self, msg):
        n = self.node
        transcript = device_revoke_transcript(
            node_pk_b64=NODE_PK, user_id=USER_ID, pk_ed25519_b64=msg["pk_ed25519"],
            nonce_node=n.nonce, ts=msg["ts"])
        if not _signed_by_a_device(n, transcript, msg["sig"]):
            return {"type": "error", "detail": "Not signed by a device already paired here"}
        if len(n.devices) <= 1:
            return {"type": "error", "detail": "This is your only device here"}
        n.devices = [d for d in n.devices if d["pk_ed25519"] != msg["pk_ed25519"]]
        return {"type": "device_add_ack", "revoked": msg["pk_ed25519"]}


def _signed_by_a_device(node: FakeNode, transcript: bytes, sig: str) -> bool:
    for d in node.devices:
        try:
            Ed25519PublicKey.from_public_bytes(base64.b64decode(d["pk_ed25519"])).verify(
                base64.b64decode(sig), transcript)
            return True
        except Exception:
            continue
    return False


def chunks_of(size: int, chunk: int) -> int:
    return max(1, math.ceil(size / chunk))