684 lines
23 KiB
Python
684 lines
23 KiB
Python
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())
|