aboutsummaryrefslogtreecommitdiffstats
path: root/poc/spike5_node.py
blob: 6104e00be1bcb204c29363f3950fdf4c064b91d2 (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
"""
MeshBay Spike 5 — Encrypted Transfer NODE (runs on Fedora, behind NAT)

Acts as the "file server":
  - Creates a 5 MB test file if not present
  - Generates or loads GEK from node_state.json
  - Connects TCP OUT to meshbay.org:19003 (outgoing = always works through NAT)
  - Receives chunk request
  - Reads chunk from disk, derives per-chunk key via HKDF, encrypts, signs
  - Sends encrypted chunk with metadata
  - Reports timing

Note: node initiates the TCP connection (reversed for POC — in production,
the QUIC client connects to the node using the hole-punching from Spike 4).
"""

import asyncio, json, base64, struct, os, time
from pathlib import Path
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes, serialization
import blake3

MESHBAY_HOST = "meshbay.org"
MESHBAY_PORT = 19003
CHUNK_SIZE   = 1024 * 1024   # 1 MB
TEST_FILE    = Path("testfile.bin")
STATE_FILE   = Path("node_state.json")
PASS = "✓"; FAIL = "✗"

# ── Wire helpers ───────────────────────────────────────────────────────────────

async def send_msg(writer, obj: dict):
    data = json.dumps(obj).encode()
    writer.write(struct.pack('>I', len(data)) + data)
    await writer.drain()

async def recv_msg(reader) -> dict:
    raw_len = await reader.readexactly(4)
    length  = struct.unpack('>I', raw_len)[0]
    data    = await reader.readexactly(length)
    return json.loads(data)

# ── Chunk encryption ───────────────────────────────────────────────────────────

def make_encrypted_chunk(
    sk_node: Ed25519PrivateKey,
    gek_raw: bytes,
    file_path: Path,
    file_hash: bytes,
    chunk_index: int
) -> dict:
    """Read chunk from disk, compress (skipped), encrypt, sign. Return wire dict."""

    # Read plaintext chunk
    with open(file_path, 'rb') as f:
        f.seek(chunk_index * CHUNK_SIZE)
        plaintext = f.read(CHUNK_SIZE)

    pt_hash = blake3.blake3(plaintext).digest()

    # Derive per-chunk key (deterministic from GEK + file identity + chunk index)
    chunk_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None,
        info=b"file:" + file_hash + b":chunk:" + chunk_index.to_bytes(4, 'big')
    ).derive(gek_raw)

    # Encrypt with ChaCha20-Poly1305
    nonce      = os.urandom(12)
    ciphertext = ChaCha20Poly1305(chunk_key).encrypt(nonce, plaintext, None)
    ct_hash    = blake3.blake3(ciphertext).digest()

    # Sign: chunk_index || nonce || ct_hash
    sig_payload = chunk_index.to_bytes(4, 'big') + nonce + ct_hash
    signature   = sk_node.sign(sig_payload)

    pk_bytes = sk_node.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)

    return {
        "chunk_index":    chunk_index,
        "plaintext_size": len(plaintext),
        "nonce_b64":      base64.b64encode(nonce).decode(),
        "ct_b64":         base64.b64encode(ciphertext).decode(),
        "ct_hash_b64":    base64.b64encode(ct_hash).decode(),
        "pt_hash_b64":    base64.b64encode(pt_hash).decode(),
        "sig_b64":        base64.b64encode(signature).decode(),
        "pk_node_b64":    base64.b64encode(pk_bytes).decode(),
        "file_hash_b64":  base64.b64encode(file_hash).decode(),
        "gek_b64":        base64.b64encode(gek_raw).decode(),   # POC ONLY
    }

# ── Main ──────────────────────────────────────────────────────────────────────

async def main():
    print("\n=== MeshBay Spike 5 — Encrypted Transfer NODE ===\n")

    # Load state
    state = json.loads(STATE_FILE.read_text()) if STATE_FILE.exists() else {}

    # Load node Ed25519 key
    sk_node = Ed25519PrivateKey.from_private_bytes(
        base64.b64decode(state["sk_ed25519_b64"]))
    pk_bytes = sk_node.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    print(f"  Node PK: {base64.b64encode(pk_bytes).decode()[:24]}...")

    # GEK — load or generate
    if "gek_b64" not in state:
        gek_raw = ChaCha20Poly1305.generate_key()
        state["gek_b64"] = base64.b64encode(gek_raw).decode()
        STATE_FILE.write_text(json.dumps(state, indent=2))
        print(f"  {PASS} GEK generated and saved")
    else:
        gek_raw = base64.b64decode(state["gek_b64"])
        print(f"  {PASS} GEK loaded from state")

    # Create test file if needed (5 MB random)
    if not TEST_FILE.exists():
        print(f"  Creating {TEST_FILE} (5 MB)...")
        with open(TEST_FILE, 'wb') as f:
            f.write(os.urandom(5 * 1024 * 1024))
    file_size = TEST_FILE.stat().st_size
    print(f"  {PASS} Test file: {TEST_FILE} ({file_size/1024/1024:.1f} MB)")

    # Pre-compute file hash (used in HKDF info string — identifies the file)
    print(f"  Computing file hash...")
    t0 = time.perf_counter()
    file_hash = blake3.blake3(TEST_FILE.read_bytes()).digest()
    print(f"  {PASS} File hash: {file_hash.hex()[:16]}... ({(time.perf_counter()-t0)*1000:.0f}ms)")

    # Connect to meshbay.org
    print(f"\n  Connecting to {MESHBAY_HOST}:{MESHBAY_PORT} ...")
    reader, writer = await asyncio.open_connection(MESHBAY_HOST, MESHBAY_PORT)
    print(f"  {PASS} TCP connection established")

    # Receive chunk request
    req = await recv_msg(reader)
    assert req["type"] == "chunk_req", f"Unexpected message type: {req['type']}"
    chunk_index = req["chunk_index"]
    file_id     = req["file_id"]
    print(f"  {PASS} Request received: file={file_id} chunk={chunk_index}")

    # Encrypt chunk on-the-fly
    print(f"\n  Encrypting chunk {chunk_index} ({CHUNK_SIZE/1024:.0f} KB)...")
    t_enc = time.perf_counter()
    chunk_data = make_encrypted_chunk(sk_node, gek_raw, TEST_FILE, file_hash, chunk_index)
    enc_ms = (time.perf_counter() - t_enc) * 1000

    ct_size = len(base64.b64decode(chunk_data["ct_b64"]))
    print(f"  {PASS} Encrypted: {CHUNK_SIZE/1024:.0f} KB → {ct_size/1024:.1f} KB ciphertext")
    print(f"  {PASS} Encrypt+sign time: {enc_ms:.1f} ms")

    # Send encrypted chunk
    print(f"\n  Sending to {MESHBAY_HOST}:{MESHBAY_PORT}...")
    t_send = time.perf_counter()
    await send_msg(writer, chunk_data)
    writer.close()
    await writer.wait_closed()
    send_ms = (time.perf_counter() - t_send) * 1000

    total_ms = enc_ms + send_ms
    payload_kb = (4 + len(json.dumps(chunk_data).encode())) / 1024
    throughput = (payload_kb / 1024) / (send_ms / 1000)

    print(f"  {PASS} Sent {payload_kb:.1f} KB in {send_ms:.0f}ms ({throughput:.1f} MB/s)")
    print(f"\n  {'='*50}")
    print(f"  Timings (node side):")
    print(f"    Encrypt + sign : {enc_ms:.1f} ms")
    print(f"    TCP send       : {send_ms:.0f} ms")
    print(f"    Total          : {total_ms:.0f} ms")
    print(f"\n  {PASS} Spike 5 NODE — chunk served successfully.")

asyncio.run(main())