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 import chatbox
from meshbay_common.adminop import OP_DIR_DELETE, 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_CHAT_KEYS,
    PURPOSE_INDEX,
    PURPOSE_ROSTER,
    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.leaseless_reads: list[tuple[str, int]] = []
        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 = {"chat_epoch": 1}   # the sealed configuration, opened
        self.chat_epochs = {1: secrets.token_bytes(32)}
        self.chat: list[dict] = []           # stored rows, oldest first
        # group_roster_resp's devices; this machine's own key is a root of its account.
        self.roster: list[dict] = [{
            "user_id": USER_ID, "username": "alice-test", "pk_ed25519": self.identity.pk_ed,
            "pk_x25519": self.identity.pk_x, "added_by_pk": "", "add_sig": "",
            "add_nonce": "", "add_ts": 0}]
        self.members: dict[str, Identity] = {}   # sender → the device it signs with
        self.unresponsive = False            # pings go unanswered, as a dead link's
        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 list_device(self, sender: str, device: Identity, *,
                    countersigned_by: Identity | None = None, proof: bool = True) -> None:
        """A roster row for `sender`'s device: a root, or countersigned by another
        of its devices — with the signature, or (`proof=False`) naming the
        countersigner and carrying nothing, as a lying node would."""
        row = {"user_id": "u-" + sender, "username": sender, "pk_ed25519": device.pk_ed,
               "pk_x25519": device.pk_x, "added_by_pk": "", "add_sig": "",
               "add_nonce": "", "add_ts": 0}
        if countersigned_by is not None:
            row["added_by_pk"] = countersigned_by.pk_ed
            if proof:
                nonce, ts = secrets.token_bytes(32), 1_700_000_000
                row.update(add_nonce=base64.b64encode(nonce).decode(), add_ts=ts,
                           add_sig=base64.b64encode(countersigned_by.sk_ed.sign(
                               device_add_transcript(
                                   node_pk_b64=NODE_PK, user_id="u-" + sender,
                                   pk_ed25519_b64=device.pk_ed,
                                   pk_x25519_b64=device.pk_x,
                                   nonce_node=nonce, ts=ts))).decode())
        self.roster.append(row)

    def say(self, text: str, *, sender: str = "bob-test", epoch: int = 1,
            forge: bool = False, device: Identity | None = None) -> Identity:
        """Another member's message, sealed and signed as a client does, from
        `device` (default: the sender's usual one, listed in the roster as a
        root the first time); `forge` signs it with a key other than the device
        it names."""
        if device is None:
            if sender not in self.members:
                self.members[sender] = Identity(Ed25519PrivateKey.generate(),
                                                X25519PrivateKey.generate())
                self.list_device(sender, self.members[sender])
            device = self.members[sender]
        other = device
        sealed = chatbox.seal(self.chat_epochs[epoch], GROUP_ID, epoch, other.pk_ed,
                              base64.b64decode(other.pk_ed), other.sk_ed,
                              {"text": text, "sender_name": sender})
        if forge:
            sealed["sig"] = Ed25519PrivateKey.generate().sign(b"something else")
        self._store(base64.b64decode(other.pk_ed), epoch, sealed, "u-" + sender)
        return other

    def _store(self, device: bytes, epoch: int, sealed: dict, sender_id: str) -> None:
        self.chat.append({"id": len(self.chat) + 1, "sender_id": sender_id,
                          "sender_name": "", "timestamp": 1_700_000_000 + len(self.chat),
                          "format": 1, "epoch": epoch,
                          "device": device, "nonce": sealed["nonce"], "sig": sealed["sig"],
                          "payload": "", "ct": sealed["ct"]})

    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 "tr" not in msg:
            # A read without a lease, as a player's: the node admits a dozen
            # files at once per connection.
            n.leaseless_reads.append((msg["file_id"], msg["chunk_index"]))
        elif 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)


    # ── Chat ─────────────────────────────────────────────────────────────────

    def _chat_keys_req(self, msg):
        n = self.node
        payload = {"epochs": [{"epoch": e, "key": k} for e, k in sorted(n.chat_epochs.items())],
                   "current": max(n.chat_epochs)}
        return {"type": "chat_keys_resp", "group_id": GROUP_ID,
                **seal(n.gek, PURPOSE_CHAT_KEYS, "chat_keys_resp", GROUP_ID, payload)}

    def _ping(self, msg):
        if self.node.unresponsive:
            return {"type": "error", "detail": "no answer"}
        return {"type": "pong", "token": msg.get("token")}

    def _group_roster_req(self, msg):
        n = self.node
        return {"type": "group_roster_resp", "group_id": GROUP_ID,
                **seal(n.gek, PURPOSE_ROSTER, "group_roster_resp", GROUP_ID,
                       {"devices": n.roster, "node_pk": NODE_PK})}

    def _chat_hist(self, msg):
        rows = self.node.chat
        if msg.get("before"):
            rows = [r for r in rows if r["id"] < msg["before"]]
        page = rows[-min(int(msg["limit"]), 200):]
        return {"type": "chat_hist_resp", "messages": page,
                "has_more": bool(page) and page[0]["id"] > rows[0]["id"]}

    def _chat_msg(self, msg):
        n = self.node
        if bytes(msg["device"]) != base64.b64decode(n.identity.pk_ed):
            return {"type": "error", "detail": "That is not the device on this connection"}
        if msg["epoch"] != max(n.chat_epochs) or msg["format"] != 1:
            return {"type": "error", "detail": "stale epoch or plaintext"}
        if not chatbox.verify(GROUP_ID, msg["epoch"], bytes(msg["device"]), msg["nonce"],
                              msg["ct"], msg["sig"]):
            return {"type": "error", "detail": "bad signature"}
        n._store(bytes(msg["device"]), msg["epoch"], msg, USER_ID)
        return {"type": "ack"}

    # ── 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 _dir_delete(self, msg):
        """Empty folders only, the operator's only, never a root — the node's rule."""
        n, path = self.node, msg["dir"].strip("/")
        if path not in n.dirs or "/" not in path:
            return {"type": "error", "detail": "Invalid directory"}
        if any(d.startswith(path + "/") for d in n.dirs) or \
                any(e["path"] == path or e["path"].startswith(path + "/")
                    for e in n.files.values()):
            return {"type": "error", "detail": "Directory is not empty"}
        if n.role != "operator":
            return {"type": "error", "detail": "No authorized key for deletion"}
        op_id, nonce = secrets.token_hex(8), secrets.token_bytes(32)
        n.admin_ops[op_id] = {"op": OP_DIR_DELETE, "subject": path, "nonce": nonce,
                              "ts": 1_700_000_000}
        return {"type": "admin_challenge", "op_id": op_id, "op": OP_DIR_DELETE,
                "subject": path, "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"}
        if op.get("op") == OP_DIR_DELETE:
            transcript = admin_transcript(op=OP_DIR_DELETE, node_pk_b64=NODE_PK,
                                          group_id=GROUP_ID, subject=op["subject"],
                                          nonce=op["nonce"], ts=op["ts"])
            if not (_signed_by_a_device(n, transcript, msg["signature"])
                    and n.role == "operator"):
                return {"type": "error", "detail": "Signature verification failed"}
            n.dirs.discard(op["subject"])
            return {"type": "dir_delete_ack", "dir": op["subject"]}
        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))