from __future__ import annotations import hashlib import ipaddress import json import math import os import struct import time from collections.abc import Callable, Iterator from contextlib import suppress from dataclasses import dataclass from pathlib import Path from typing import IO, Literal, TypedDict import paho.mqtt.client as mqtt from paho.mqtt.enums import CallbackAPIVersion from paho.mqtt.properties import Properties from paho.mqtt.reasoncodes import ReasonCode from k1link.artifacts import utc_now_iso REPORT_TOPICS: tuple[str, ...] = ( "lixel/application/report/#", "RealtimePointcloud", "RealtimePath", "DeviceStatus", ) DEFAULT_MAX_MESSAGE_BYTES = 64 * 1024 * 1024 MAX_CONFIGURABLE_MESSAGE_BYTES = 256 * 1024 * 1024 MAX_TOPIC_BYTES = 65_535 CONNECT_TIMEOUT_SECONDS = 10.0 KEEPALIVE_SECONDS = 30 LOOP_INTERVAL_SECONDS = 0.25 # Eight-byte file signature followed by repeated >IQ, topic UTF-8 bytes, payload bytes. RAW_MAGIC = b"K1MQTT\x00\x01" FRAME_HEADER = struct.Struct(">IQ") _PRIVATE_NETWORKS = tuple( ipaddress.ip_network(cidr) for cidr in ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16") ) StopReason = Literal[ "duration_elapsed", "external_stop", "keyboard_interrupt", "message_too_large", "connection_failed", "connection_lost", "subscription_failed", "capture_error", ] class ArtifactPaths(TypedDict): raw: str metadata_jsonl: str summary: str class ArtifactHashes(TypedDict): raw_sha256: str metadata_jsonl_sha256: str class RawFormat(TypedDict): magic_hex: str frame_header_struct: str frame_layout: str class CaptureSummary(TypedDict): schema_version: int created_at_utc: str completed_at_utc: str sensitivity: str target_ipv4: str target_port: int mqtt_protocol: str subscription_qos: int clean_session: bool reconnect_enabled: bool publishing_enabled: bool subscriptions: list[str] requested_duration_seconds: float capture_elapsed_seconds: float operation_elapsed_seconds: float max_message_bytes: int connected: bool subscribed: bool stop_reason: StopReason error: str | None message_count: int rejected_message_count: int payload_bytes: int raw_bytes: int topic_counts: dict[str, int] raw_format: RawFormat artifacts: ArtifactPaths artifact_hashes: ArtifactHashes class CaptureError(RuntimeError): """A one-shot capture failed after preserving all artifacts written so far.""" def __init__(self, message: str, summary: CaptureSummary | None = None) -> None: super().__init__(message) self.summary = summary class MessageTooLargeError(CaptureError): """An MQTT message exceeded the configured evidence boundary.""" class CaptureFormatError(ValueError): """A raw capture is corrupt, truncated or outside configured reader bounds.""" @dataclass(frozen=True, slots=True) class CaptureFrame: sequence: int topic: str payload: bytes raw_frame_offset: int raw_payload_offset: int raw_frame_bytes: int @dataclass(frozen=True, slots=True) class CapturedMqttMessage: """A message flushed to the raw writer before it is exposed to live preview.""" sequence: int topic: str payload: bytes qos: int retain: bool dup: bool received_at_utc: str received_at_epoch_ns: int received_monotonic_ns: int @dataclass class _CaptureState: connected: bool = False subscribed: bool = False stopping: bool = False subscription_mid: int | None = None stop_reason: StopReason = "capture_error" error: str | None = None class _CaptureWriter: def __init__(self, out_dir: Path, max_message_bytes: int) -> None: self.out_dir = out_dir.expanduser().resolve() self.raw_path = self.out_dir / "mqtt.raw.k1mqtt" self.metadata_path = self.out_dir / "mqtt.metadata.jsonl" self.summary_path = self.out_dir / "mqtt.summary.json" self.max_message_bytes = max_message_bytes self.message_count = 0 self.rejected_message_count = 0 self.payload_bytes = 0 self.topic_counts: dict[str, int] = {} self._raw: IO[bytes] | None = None self._metadata: IO[str] | None = None def open(self) -> None: self.out_dir.mkdir(parents=True, exist_ok=True) artifact_paths = (self.raw_path, self.metadata_path, self.summary_path) existing = [path.name for path in artifact_paths if path.exists()] if existing: names = ", ".join(existing) raise FileExistsError(f"refusing to overwrite existing capture artifact(s): {names}") try: self._raw = _open_binary_exclusive(self.raw_path) self._raw.write(RAW_MAGIC) self._metadata = _open_text_exclusive(self.metadata_path) except BaseException: with suppress(OSError): self.close() raise def record(self, message: mqtt.MQTTMessage) -> CapturedMqttMessage: raw = self._require_raw() metadata = self._require_metadata() topic = message.topic topic_bytes = topic.encode("utf-8") payload = message.payload received_at_utc = utc_now_iso() received_at_epoch_ns = time.time_ns() received_monotonic_ns = time.monotonic_ns() if not 1 <= len(topic_bytes) <= MAX_TOPIC_BYTES: raise ValueError( f"incoming MQTT topic is {len(topic_bytes)} bytes; expected 1..{MAX_TOPIC_BYTES}" ) if len(payload) > self.max_message_bytes: self.rejected_message_count += 1 record = { "schema_version": 1, "record_type": "rejected_message", "sequence": self.message_count + self.rejected_message_count, "received_at_utc": received_at_utc, "received_monotonic_ns": received_monotonic_ns, "topic": topic, "payload_bytes": len(payload), "max_message_bytes": self.max_message_bytes, "reason": "message_too_large", } self._write_metadata(metadata, record) raise MessageTooLargeError( f"message on {topic!r} is {len(payload)} bytes; " f"limit is {self.max_message_bytes} bytes" ) offset = raw.tell() header = FRAME_HEADER.pack(len(topic_bytes), len(payload)) raw.write(header) raw.write(topic_bytes) raw.write(payload) raw.flush() self.message_count += 1 self.payload_bytes += len(payload) self.topic_counts[topic] = self.topic_counts.get(topic, 0) + 1 frame_bytes = len(header) + len(topic_bytes) + len(payload) record = { "schema_version": 1, "record_type": "message", "sequence": self.message_count, "received_at_utc": received_at_utc, "received_at_epoch_ns": received_at_epoch_ns, "received_monotonic_ns": received_monotonic_ns, "topic": topic, "qos": message.qos, "retain": message.retain, "dup": message.dup, "payload_bytes": len(payload), "payload_sha256": hashlib.sha256(payload).hexdigest(), "raw_frame_offset": offset, "raw_payload_offset": offset + len(header) + len(topic_bytes), "raw_frame_bytes": frame_bytes, } self._write_metadata(metadata, record) return CapturedMqttMessage( sequence=self.message_count, topic=topic, payload=payload, qos=message.qos, retain=message.retain, dup=message.dup, received_at_utc=received_at_utc, received_at_epoch_ns=received_at_epoch_ns, received_monotonic_ns=received_monotonic_ns, ) def close(self) -> None: first_error: OSError | None = None # Make raw frames durable before making their JSONL references durable. for stream in (self._raw, self._metadata): if stream is None or stream.closed: continue try: stream.flush() os.fsync(stream.fileno()) except OSError as exc: if first_error is None: first_error = exc finally: try: stream.close() except OSError as exc: if first_error is None: first_error = exc if first_error is not None: raise first_error @property def raw_bytes(self) -> int: if self.raw_path.exists(): return self.raw_path.stat().st_size return 0 def _require_raw(self) -> IO[bytes]: if self._raw is None or self._raw.closed: raise RuntimeError("capture writer is not open") return self._raw def _require_metadata(self) -> IO[str]: if self._metadata is None or self._metadata.closed: raise RuntimeError("capture writer is not open") return self._metadata @staticmethod def _write_metadata(stream: IO[str], record: dict[str, object]) -> None: stream.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n") stream.flush() def validate_private_ipv4(value: str) -> str: """Require a literal RFC1918 address so capture cannot target arbitrary hosts.""" try: address = ipaddress.ip_address(value) except ValueError as exc: raise ValueError("host must be a literal private IPv4 address") from exc if not isinstance(address, ipaddress.IPv4Address) or not any( address in network for network in _PRIVATE_NETWORKS ): raise ValueError("host must be an RFC1918 private IPv4 address") return str(address) def iter_capture_frames( path: Path, *, max_payload_bytes: int = MAX_CONFIGURABLE_MESSAGE_BYTES, max_topic_bytes: int = MAX_TOPIC_BYTES, ) -> Iterator[CaptureFrame]: """Yield validated frames from a K1 MQTT raw capture without decoding payloads.""" if not 1 <= max_payload_bytes <= MAX_CONFIGURABLE_MESSAGE_BYTES: raise ValueError( f"max_payload_bytes must be between 1 and {MAX_CONFIGURABLE_MESSAGE_BYTES}" ) if not 1 <= max_topic_bytes <= MAX_TOPIC_BYTES: raise ValueError(f"max_topic_bytes must be between 1 and {MAX_TOPIC_BYTES}") with path.expanduser().open("rb") as stream: magic = stream.read(len(RAW_MAGIC)) if magic != RAW_MAGIC: if len(magic) < len(RAW_MAGIC): raise CaptureFormatError("raw capture is truncated before the complete magic") raise CaptureFormatError("raw capture magic/version is not supported") sequence = 0 while True: frame_offset = stream.tell() header = stream.read(FRAME_HEADER.size) if not header: return if len(header) != FRAME_HEADER.size: raise CaptureFormatError( f"frame at offset {frame_offset} has a truncated length header" ) topic_length, payload_length = FRAME_HEADER.unpack(header) if not 1 <= topic_length <= max_topic_bytes: raise CaptureFormatError( f"frame at offset {frame_offset} topic length {topic_length} " f"is outside 1..{max_topic_bytes}" ) if payload_length > max_payload_bytes: raise CaptureFormatError( f"frame at offset {frame_offset} payload length {payload_length} " f"exceeds {max_payload_bytes}" ) topic_raw = _read_exact( stream, topic_length, description=f"topic at frame offset {frame_offset}", ) try: topic = topic_raw.decode("utf-8") except UnicodeDecodeError as exc: raise CaptureFormatError( f"frame at offset {frame_offset} topic is not valid UTF-8" ) from exc payload_offset = stream.tell() payload = _read_exact( stream, payload_length, description=f"payload at frame offset {frame_offset}", ) sequence += 1 yield CaptureFrame( sequence=sequence, topic=topic, payload=payload, raw_frame_offset=frame_offset, raw_payload_offset=payload_offset, raw_frame_bytes=FRAME_HEADER.size + topic_length + payload_length, ) def capture_mqtt( host: str, out_dir: Path, *, port: int = 1883, duration_seconds: float = 60.0, max_message_bytes: int = DEFAULT_MAX_MESSAGE_BYTES, on_ready: Callable[[], None] | None = None, on_message_recorded: Callable[[CapturedMqttMessage], None] | None = None, should_stop: Callable[[], bool] | None = None, _client_factory: Callable[[], mqtt.Client] | None = None, ) -> CaptureSummary: """Capture the fixed K1 report subscriptions once, without publishing or reconnecting.""" target_ipv4 = validate_private_ipv4(host) if not 1 <= port <= 65535: raise ValueError("port must be between 1 and 65535") if not math.isfinite(duration_seconds) or duration_seconds <= 0: raise ValueError("duration_seconds must be finite and greater than zero") if not 1 <= max_message_bytes <= MAX_CONFIGURABLE_MESSAGE_BYTES: raise ValueError( f"max_message_bytes must be between 1 and {MAX_CONFIGURABLE_MESSAGE_BYTES}" ) client = ( _client_factory() if _client_factory is not None else mqtt.Client( callback_api_version=CallbackAPIVersion.VERSION2, clean_session=True, protocol=mqtt.MQTTv311, reconnect_on_failure=False, ) ) writer = _CaptureWriter(out_dir, max_message_bytes) writer.open() state = _CaptureState() created_at_utc = utc_now_iso() operation_started = time.monotonic() capture_started: float | None = None failure: CaptureError | None = None def fail(reason: StopReason, message: str) -> None: if state.error is None: state.stop_reason = reason state.error = message def on_connect( callback_client: mqtt.Client, _userdata: object, _flags: mqtt.ConnectFlags, reason_code: ReasonCode, _properties: Properties | None, ) -> None: if reason_code.is_failure: fail("connection_failed", f"broker rejected connection: {reason_code}") return state.connected = True try: result, mid = callback_client.subscribe([(topic, 0) for topic in REPORT_TOPICS]) except (OSError, RuntimeError, ValueError) as exc: fail( "subscription_failed", f"subscribe failed: {type(exc).__name__}: {exc}", ) return if result != mqtt.MQTT_ERR_SUCCESS or mid is None: fail("subscription_failed", f"subscribe failed: {mqtt.error_string(result)}") return state.subscription_mid = mid def on_subscribe( _callback_client: mqtt.Client, _userdata: object, mid: int, reason_codes: list[ReasonCode], _properties: Properties | None, ) -> None: if mid != state.subscription_mid: fail("subscription_failed", f"unexpected SUBACK message id: {mid}") return if len(reason_codes) != len(REPORT_TOPICS) or any( reason_code.is_failure for reason_code in reason_codes ): fail("subscription_failed", "broker rejected one or more fixed subscriptions") return state.subscribed = True def on_message( _callback_client: mqtt.Client, _userdata: object, message: mqtt.MQTTMessage, ) -> None: if state.error is not None: return try: recorded = writer.record(message) except MessageTooLargeError as exc: fail("message_too_large", str(exc)) except (OSError, RuntimeError, ValueError) as exc: fail("capture_error", f"artifact write failed: {type(exc).__name__}: {exc}") return if on_message_recorded is not None: try: on_message_recorded(recorded) except (OSError, RuntimeError, ValueError) as exc: fail("capture_error", f"preview callback failed: {type(exc).__name__}: {exc}") def on_disconnect( _callback_client: mqtt.Client, _userdata: object, _flags: mqtt.DisconnectFlags, reason_code: ReasonCode, _properties: Properties | None, ) -> None: if not state.stopping: fail("connection_lost", f"broker connection ended: {reason_code}") client.on_connect = on_connect client.on_subscribe = on_subscribe client.on_message = on_message client.on_disconnect = on_disconnect connect_attempted = False try: connect_attempted = True connect_result = client.connect( target_ipv4, port=port, keepalive=KEEPALIVE_SECONDS, ) if connect_result != mqtt.MQTT_ERR_SUCCESS: fail("connection_failed", f"connect failed: {mqtt.error_string(connect_result)}") while state.error is None: now = time.monotonic() if should_stop is not None and should_stop(): state.stop_reason = "external_stop" break if state.subscribed and capture_started is None: capture_started = now if on_ready is not None: on_ready() if capture_started is not None and now - capture_started >= duration_seconds: state.stop_reason = "duration_elapsed" break if capture_started is None and now - operation_started >= CONNECT_TIMEOUT_SECONDS: fail("connection_failed", "timed out waiting for CONNACK/SUBACK") break loop_result = client.loop(timeout=LOOP_INTERVAL_SECONDS) if loop_result != mqtt.MQTT_ERR_SUCCESS and state.error is None: fail( "connection_lost", f"MQTT network loop failed: {mqtt.error_string(loop_result)}", ) except KeyboardInterrupt: state.stop_reason = "keyboard_interrupt" except (OSError, RuntimeError, ValueError) as exc: fail("connection_failed", f"MQTT capture failed: {type(exc).__name__}: {exc}") finally: state.stopping = True if connect_attempted: try: client.disconnect() except (OSError, RuntimeError, ValueError) as exc: if state.error is None: fail("capture_error", f"disconnect failed: {type(exc).__name__}: {exc}") try: writer.close() except OSError as exc: if state.error is None: fail("capture_error", f"artifact close failed: {type(exc).__name__}: {exc}") operation_completed = time.monotonic() capture_elapsed = 0.0 if capture_started is None else operation_completed - capture_started summary = _build_summary( writer=writer, target_ipv4=target_ipv4, port=port, duration_seconds=duration_seconds, capture_elapsed=capture_elapsed, operation_elapsed=operation_completed - operation_started, max_message_bytes=max_message_bytes, created_at_utc=created_at_utc, state=state, ) try: _write_summary_exclusive(writer.summary_path, summary) except OSError as exc: raise CaptureError(f"could not write capture summary: {type(exc).__name__}: {exc}") from exc if state.error is not None: failure = CaptureError(state.error, summary) if failure is not None: raise failure return summary def _build_summary( *, writer: _CaptureWriter, target_ipv4: str, port: int, duration_seconds: float, capture_elapsed: float, operation_elapsed: float, max_message_bytes: int, created_at_utc: str, state: _CaptureState, ) -> CaptureSummary: return { "schema_version": 1, "created_at_utc": created_at_utc, "completed_at_utc": utc_now_iso(), "sensitivity": "contains raw K1 MQTT payloads and local addressing; do not commit", "target_ipv4": target_ipv4, "target_port": port, "mqtt_protocol": "3.1.1", "subscription_qos": 0, "clean_session": True, "reconnect_enabled": False, "publishing_enabled": False, "subscriptions": list(REPORT_TOPICS), "requested_duration_seconds": duration_seconds, "capture_elapsed_seconds": round(capture_elapsed, 6), "operation_elapsed_seconds": round(operation_elapsed, 6), "max_message_bytes": max_message_bytes, "connected": state.connected, "subscribed": state.subscribed, "stop_reason": state.stop_reason, "error": state.error, "message_count": writer.message_count, "rejected_message_count": writer.rejected_message_count, "payload_bytes": writer.payload_bytes, "raw_bytes": writer.raw_bytes, "topic_counts": dict(sorted(writer.topic_counts.items())), "raw_format": { "magic_hex": RAW_MAGIC.hex(), "frame_header_struct": FRAME_HEADER.format, "frame_layout": "topic_length:uint32, payload_length:uint64, topic_utf8, payload", }, "artifacts": { "raw": writer.raw_path.name, "metadata_jsonl": writer.metadata_path.name, "summary": writer.summary_path.name, }, "artifact_hashes": { "raw_sha256": _sha256_file(writer.raw_path), "metadata_jsonl_sha256": _sha256_file(writer.metadata_path), }, } def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _read_exact(stream: IO[bytes], length: int, *, description: str) -> bytes: value = stream.read(length) if len(value) != length: raise CaptureFormatError( f"{description} is truncated: expected {length} bytes, got {len(value)}" ) return value def _open_binary_exclusive(path: Path) -> IO[bytes]: descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) try: return os.fdopen(descriptor, "wb") except BaseException: os.close(descriptor) raise def _open_text_exclusive(path: Path) -> IO[str]: descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) try: return os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") except BaseException: os.close(descriptor) raise def _write_summary_exclusive(path: Path, summary: CaptureSummary) -> None: serialized = json.dumps(summary, ensure_ascii=False, indent=2) + "\n" with _open_text_exclusive(path) as stream: stream.write(serialized) stream.flush() os.fsync(stream.fileno())