Files
audio_project_hub/scripts/simulate_node.py
T
2026-08-14 08:24:33 -06:00

310 lines
10 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Simulate one audio Node for the ESP32 Hub over the node UDP protocol.
The Hub (see src/network/node_protocol.h) runs:
* control plane on UDP 4210 (PROBE / START / STOP / POLL)
* data plane on UDP 4211 (DATA packets from Nodes)
This script behaves like a Node with 4 microphones:
* answers PROBE with a PROBE_RESP (so the Hub registers it)
* streams 4 interleaved 48 kHz / 16-bit sine tones (one per mic) as
DATA packets while recording, paced to real time so the WAV plays back
at the correct pitch and duration
* sends HEARTBEATs so the Hub keeps it marked online
Usage:
# PC connected to the Hub's AP (192.168.4.x), or any host on the subnet
python3 scripts/simulate_node.py
# multiple virtual nodes (distinct tones):
python3 scripts/simulate_node.py --node-id 2 --frequencies 330,415,495,660
The four default tones are an A-major chord, one per mic, so each channel is
clearly distinguishable when the 40-channel WAV is played back.
Requires only the Python standard library.
"""
import argparse
import math
import select
import signal
import socket
import struct
import sys
import time
# --- Protocol constants (mirror src/network/network_config.h + node_protocol.h)
PROTOCOL_VERSION = 1
CONTROL_PORT = 4210
DATA_PORT = 4211
MSG_PROBE = 0x01
MSG_PROBE_RESP = 0x02
MSG_START = 0x03
MSG_STOP = 0x04
MSG_HEARTBEAT = 0x05
MSG_POLL = 0x06
MSG_POLL_ACK = 0x07
MSG_DATA = 0x10
STATUS_IDLE = 0
STATUS_RECORDING = 1
MICS_PER_NODE = 1
SAMPLE_RATE = 16000
BITS_PER_SAMPLE = 16
FRAME_BYTES = MICS_PER_NODE * 2 # 8 bytes per frame (4 ch x 2 B)
# Frames per DATA packet: (header 8 B) + n * 8 B must stay below 1472.
FRAMES_PER_PACKET = 100
AMPLITUDE = 8000.0 # ~0.25 x full scale
HEARTBEAT_MS = 2000
def u16(x):
return struct.pack('<H', x & 0xFFFF)
def u32(x):
return struct.pack('<I', x & 0xFFFFFFFF)
def build_probe_resp(node_id, mac, firmware, mic_count, sample_rate,
bits_per_sample, buffer_bytes, status):
fw = firmware.encode('ascii')[:15].ljust(15, b'\x00') + b'\x00'
return (bytes([MSG_PROBE_RESP, PROTOCOL_VERSION, node_id])
+ mac
+ fw
+ bytes([mic_count])
+ u32(sample_rate)
+ u16(bits_per_sample)
+ u32(buffer_bytes)
+ bytes([status]))
def build_heartbeat(node_id, uptime_ms, buffer_fill, status):
return (bytes([MSG_HEARTBEAT, PROTOCOL_VERSION, node_id])
+ u32(uptime_ms)
+ u32(buffer_fill)
+ bytes([status]))
def build_data(node_id, mic_count, frame_count, pcm):
return (bytes([MSG_DATA, PROTOCOL_VERSION, node_id, mic_count])
+ u32(frame_count)
+ pcm)
class SineGenerator:
"""Four independent phase accumulators producing interleaved 16-bit PCM."""
def __init__(self, frequencies, amplitude):
self.frequencies = frequencies
self.amplitude = amplitude
self.phase = [0.0] * len(frequencies)
self._two_pi = 2.0 * math.pi
def frames(self, count):
"""Generate `count` interleaved frames (count * FRAME_BYTES bytes)."""
pcm = bytearray(count * FRAME_BYTES)
step = [self._two_pi * f / SAMPLE_RATE for f in self.frequencies]
for f in range(count):
for ch in range(len(self.frequencies)):
self.phase[ch] += step[ch]
if self.phase[ch] >= self._two_pi:
self.phase[ch] -= self._two_pi
sample = int(self.amplitude * math.sin(self.phase[ch]))
off = (f * MICS_PER_NODE + ch) * 2
pcm[off] = sample & 0xFF
pcm[off + 1] = (sample >> 8) & 0xFF
return bytes(pcm)
class SimulatedNode:
def __init__(self, args):
self.node_id = args.node_id
self.mac = bytes([0x02, 0x00, 0x00, 0x00, 0x00, 0x10 + (args.node_id & 0xEF)])
self.firmware = "py-node-v1"
self.buffer_bytes = 64 * 1024
self.hub_ip = args.hub_ip
self.recording = False
self.status = STATUS_IDLE
self.uptime_ms = 0
self.started_at = 0.0
self.produced = 0
self.last_heartbeat = 0.0
self.running = True
self.sine = SineGenerator(args.frequencies, args.amplitude)
self.stop_after = args.duration
self.ctrl = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.ctrl.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
try:
self.ctrl.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
except OSError:
pass
self.ctrl.bind(('', CONTROL_PORT))
self.data = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
self.data.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
except OSError:
pass
# --- helpers ----------------------------------------------------------
def send_probe_resp(self):
pkt = build_probe_resp(self.node_id, self.mac, self.firmware,
MICS_PER_NODE, SAMPLE_RATE, BITS_PER_SAMPLE,
self.buffer_bytes, self.status)
self.ctrl.sendto(pkt, (self.hub_ip, CONTROL_PORT))
print(f"[node {self.node_id}] PROBE_RESP sent to {self.hub_ip}")
def send_heartbeat(self):
pkt = build_heartbeat(self.node_id, self.uptime_ms,
self.buffered_frames() * FRAME_BYTES, self.status)
self.ctrl.sendto(pkt, (self.hub_ip, CONTROL_PORT))
def buffered_frames(self):
if not self.recording:
return 0
return max(0, int(time.monotonic() - self.started_at) * SAMPLE_RATE
- self.produced)
def stream_pacing(self):
"""Send the frames that real time says are due, in DATA packets."""
if not self.recording or self.hub_ip is None:
return
target = int((time.monotonic() - self.started_at) * SAMPLE_RATE)
while self.produced < target:
n = min(target - self.produced, FRAMES_PER_PACKET)
pcm = self.sine.frames(n)
pkt = build_data(self.node_id, MICS_PER_NODE, n, pcm)
self.data.sendto(pkt, (self.hub_ip, DATA_PORT))
self.produced += n
# --- control packet handling ------------------------------------------
def handle(self, packet, src):
if len(packet) < 2:
return
msg, ver = packet[0], packet[1]
if ver != PROTOCOL_VERSION:
return
if src[0] != self.hub_ip:
self.hub_ip = src[0] # learn the Hub from any control packet
print(f"[node {self.node_id}] learned Hub at {self.hub_ip}")
if msg == MSG_PROBE:
self.send_probe_resp()
elif msg == MSG_START:
self.recording = True
self.status = STATUS_RECORDING
self.started_at = time.monotonic()
self.produced = 0
print(f"[node {self.node_id}] START - streaming 4 tones at "
f"{SAMPLE_RATE} Hz, {MICS_PER_NODE} ch")
elif msg == MSG_STOP:
if self.recording:
print(f"[node {self.node_id}] STOP - sent "
f"{self.produced} frames")
self.recording = False
self.status = STATUS_IDLE
elif msg == MSG_POLL:
# We stream continuously; nothing extra to send.
pass
# --- main loop --------------------------------------------------------
def run(self):
print(f"[node {self.node_id}] listening on UDP {CONTROL_PORT}...")
print(f"[node {self.node_id}] 4 tones: {self.sine.frequencies} Hz")
if self.hub_ip:
print(f"[node {self.node_id}] Hub set to {self.hub_ip}")
else:
print("[node %d] waiting for a PROBE to learn the Hub IP "
"(PC must be on the Hub's AP network)" % self.node_id)
while self.running:
ready, _, _ = select.select([self.ctrl], [], [], 0.002)
if ready:
try:
packet, src = self.ctrl.recvfrom(2048)
self.handle(packet, src)
except OSError:
pass
if self.recording:
self.stream_pacing()
if self.stop_after is not None:
if time.monotonic() - self.started_at >= self.stop_after:
print(f"[node {self.node_id}] --duration reached, "
"waiting for STOP")
self.recording = False
self.status = STATUS_IDLE
now_ms = int(time.monotonic() * 1000)
self.uptime_ms = now_ms
if self.hub_ip and now_ms - self.last_heartbeat >= HEARTBEAT_MS:
self.last_heartbeat = now_ms
self.send_heartbeat()
self.ctrl.close()
self.data.close()
def parse_args(argv):
parser = argparse.ArgumentParser(
description='Simulate a 4-mic audio Node for the ESP32 Hub.')
parser.add_argument('--node-id', type=int, default=1,
help='Node ID (default 1). Overrides the matching '
'built-in simulator node.')
parser.add_argument('--hub-ip', default=None,
help='Hub IP, e.g. 192.168.4.1 (auto-learned from '
'the first control packet if omitted).')
parser.add_argument('--frequencies', default='440,554.37,659.25,880',
help='Comma-separated sine frequencies per mic '
'(default an A-major chord).')
parser.add_argument('--amplitude', type=float, default=AMPLITUDE,
help='Sine amplitude 0..32767 (default 8000).')
parser.add_argument('--duration', type=float, default=None,
help='Auto-stop after N seconds of recording '
'(default: until the Hub sends STOP).')
return parser.parse_args(argv)
def main(argv=None):
args = parse_args(argv)
args.frequencies = [float(f) for f in args.frequencies.split(',')]
if len(args.frequencies) < MICS_PER_NODE:
print('error: provide at least %d frequencies' % MICS_PER_NODE,
file=sys.stderr)
return 1
node = SimulatedNode(args)
def on_sigint(_sig, _frame):
node.running = False
print('\n[node %d] shutting down' % args.node_id)
signal.signal(signal.SIGINT, on_sigint)
node.run()
return 0
if __name__ == '__main__':
sys.exit(main())