summaryrefslogtreecommitdiffstats
path: root/packages/meshbay-node/src/meshbay_node/indexer/group_index.py
blob: 4edeeab2f240c5a326295114b5515bc2446ef17a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
"""
Mesh Group Index — encrypted file listing for a group.

Wire format (private group):
  msgpack({entries, version, group_id}) → zstd compress → GEK ChaCha20 encrypt → sign

Wire format (public group):
  msgpack({entries, version, group_id}) → sign (no encryption)

Delta format:
  {base_version, version, additions: [...], deletions: [id, ...]}
"""

import base64
import logging
import os
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Iterator

import blake3
import msgpack
import zstandard as zstd
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from meshbay_common.crypto import (
    chunk_key as derive_chunk_key,
    encrypt_chunk,
    decrypt_chunk,
    sign_chunk,
    verify_chunk_signature,
    pk_to_b64,
    generate_gek,
)
from meshbay_common.protocol import IndexEntry, IndexDelta

log = logging.getLogger(__name__)

ZSTD_LEVEL   = 3    # fast compression
INDEX_CHUNK  = 0    # the index itself is treated as chunk 0 of a virtual "index file"


@dataclass
class GroupIndex:
    """
    Encrypted, signed Mesh Group Index for one group.

    Usage:
        idx = GroupIndex(group_id="...", sk_node=sk, gek=gek_bytes)
        idx.add_entry(entry)
        wire_bytes = idx.serialize()          # for sending to members
        recovered  = GroupIndex.deserialize(wire_bytes, sk_node=sk, gek=gek_bytes)
    """

    group_id: str
    sk_node:  Ed25519PrivateKey
    gek:      bytes | None = None   # None → public group (no encryption)
    version:  int = 1
    _entries: dict = field(default_factory=dict, repr=False)  # id → IndexEntry

    # ── Entry management ──────────────────────────────────────────────────────

    def add_entry(self, entry: IndexEntry) -> None:
        self._entries[entry.id] = entry

    def remove_entry(self, file_id: str) -> bool:
        return self._entries.pop(file_id, None) is not None

    def get_entry(self, file_id: str) -> IndexEntry | None:
        return self._entries.get(file_id)

    @property
    def entries(self) -> list[IndexEntry]:
        return list(self._entries.values())

    @property
    def count(self) -> int:
        return len(self._entries)

    # ── Serialisation ─────────────────────────────────────────────────────────

    def serialize(self) -> bytes:
        """
        Produce a wire-ready byte string:
          msgpack → zstd → [GEK encrypt if private] → sign → length-prefixed envelope
        """
        payload = msgpack.packb({
            "group_id": self.group_id,
            "version":  self.version,
            "entries":  [asdict(e) for e in self.entries],
        }, use_bin_type=True)

        compressed = zstd.compress(payload, level=ZSTD_LEVEL)

        if self.gek is not None:
            # Private group: encrypt with GEK-derived key
            idx_hash   = blake3.blake3(compressed).digest()
            ckey       = derive_chunk_key(self.gek, idx_hash, INDEX_CHUNK)
            nonce, ct  = encrypt_chunk(ckey, compressed)
            sig        = sign_chunk(self.sk_node, INDEX_CHUNK, nonce, blake3.blake3(ct).digest())
            envelope = msgpack.packb({
                "type":       "index",
                "encrypted":  True,
                "version":    self.version,
                "group_id":   self.group_id,
                "idx_hash_b64": base64.b64encode(idx_hash).decode(),
                "nonce_b64":  base64.b64encode(nonce).decode(),
                "ct_b64":     base64.b64encode(ct).decode(),
                "sig_b64":    base64.b64encode(sig).decode(),
                "pk_node_b64": pk_to_b64(self.sk_node.public_key()),
            }, use_bin_type=True)
        else:
            # Public group: just sign the compressed payload
            payload_hash = blake3.blake3(compressed).digest()
            sig = sign_chunk(self.sk_node, INDEX_CHUNK,
                             bytes(12),        # zero nonce for plaintext
                             payload_hash)
            envelope = msgpack.packb({
                "type":       "index",
                "encrypted":  False,
                "version":    self.version,
                "group_id":   self.group_id,
                "data_b64":   base64.b64encode(compressed).decode(),
                "hash_b64":   base64.b64encode(payload_hash).decode(),
                "sig_b64":    base64.b64encode(sig).decode(),
                "pk_node_b64": pk_to_b64(self.sk_node.public_key()),
            }, use_bin_type=True)

        return envelope

    @classmethod
    def deserialize(
        cls,
        data: bytes,
        sk_node: Ed25519PrivateKey,
        gek: bytes | None = None,
    ) -> "GroupIndex":
        """Deserialize, verify signature, and decrypt (if private)."""
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from meshbay_common.crypto import verify_chunk_signature

        envelope = msgpack.unpackb(data, raw=False)

        pk_node_raw = base64.b64decode(envelope["pk_node_b64"])
        pk_node     = Ed25519PublicKey.from_public_bytes(pk_node_raw)
        sig         = base64.b64decode(envelope["sig_b64"])

        if envelope["encrypted"]:
            if gek is None:
                raise ValueError("GEK required to decrypt private group index")
            ct        = base64.b64decode(envelope["ct_b64"])
            nonce     = base64.b64decode(envelope["nonce_b64"])
            ct_hash   = blake3.blake3(ct).digest()
            verify_chunk_signature(pk_node, INDEX_CHUNK, nonce, ct_hash, sig)

            idx_hash  = base64.b64decode(envelope["idx_hash_b64"])
            ckey      = derive_chunk_key(gek, idx_hash, INDEX_CHUNK)
            compressed = decrypt_chunk(ckey, nonce, ct)
        else:
            compressed  = base64.b64decode(envelope["data_b64"])
            payload_hash = base64.b64decode(envelope["hash_b64"])
            verify_chunk_signature(pk_node, INDEX_CHUNK, bytes(12), payload_hash, sig)

        payload = msgpack.unpackb(zstd.decompress(compressed), raw=False)
        idx = cls(
            group_id=payload["group_id"],
            sk_node=sk_node,
            gek=gek,
            version=payload["version"],
        )
        for e in payload["entries"]:
            idx.add_entry(IndexEntry(**e))
        return idx

    # ── Delta ─────────────────────────────────────────────────────────────────

    def diff(self, previous: "GroupIndex") -> IndexDelta:
        """Compute what changed since a previous version of this index."""
        prev_ids = set(previous._entries)
        curr_ids = set(self._entries)

        additions = [self._entries[i] for i in curr_ids - prev_ids]
        deletions = list(prev_ids - curr_ids)

        return IndexDelta(
            base_version=previous.version,
            version=self.version,
            additions=additions,
            deletions=deletions,
        )