Files
CaptainJack/tools/capture_export.py
2026-05-29 00:29:14 -05:00

266 lines
7.5 KiB
Python
Executable File

#!/usr/bin/env python3
"""Download CaptainJack capture files from the watch over USB serial."""
from __future__ import annotations
import argparse
import binascii
import datetime as dt
import sys
import time
from pathlib import Path
try:
import serial
from serial.tools import list_ports
except ImportError: # pragma: no cover - dependency error path
serial = None
list_ports = None
BAUD = 115200
MARKER = "@WS "
CAPTURE_FILES = ("captures.pcapng", "captures.22000")
class ProtocolError(RuntimeError):
pass
def require_serial() -> None:
if serial is None:
raise SystemExit("pyserial is required: python3 -m pip install pyserial")
def open_serial(port: str, timeout: float = 3.0):
require_serial()
ser = serial.Serial()
ser.port = port
ser.baudrate = BAUD
ser.timeout = timeout
ser.write_timeout = timeout
ser.dtr = False
ser.rts = False
ser.open()
ser.dtr = False
ser.rts = False
time.sleep(0.15)
ser.reset_input_buffer()
return ser
def protocol_line(ser, timeout: float = 8.0) -> str:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
raw = ser.readline()
if not raw:
continue
text = raw.decode("utf-8", "replace").strip()
pos = text.find(MARKER)
if pos >= 0:
return text[pos:]
raise TimeoutError("timed out waiting for @WS response")
def send(ser, command: str) -> None:
ser.write((command.strip() + "\n").encode("ascii"))
ser.flush()
def ping(ser) -> bool:
for _ in range(3):
send(ser, "PING")
try:
line = protocol_line(ser, timeout=4.0)
except TimeoutError:
continue
if line == "@WS READY":
send(ser, "PING")
continue
if line == "@WS PONG":
return True
return False
def list_captures(ser) -> list[tuple[str, int]]:
send(ser, "LIST")
files: list[tuple[str, int]] = []
saw_begin = False
while True:
line = protocol_line(ser)
if line == "@WS LIST BEGIN":
saw_begin = True
continue
if line == "@WS LIST END":
return files
if line.startswith("@WS ERR "):
raise ProtocolError(line)
if line.startswith("@WS FILE "):
if not saw_begin:
raise ProtocolError("FILE before LIST BEGIN")
parts = line.split(" ", 3)
if len(parts) != 4:
raise ProtocolError(f"bad FILE line: {line}")
files.append((parts[2], int(parts[3])))
def download_one(ser, name: str, out_dir: Path) -> Path:
send(ser, f"GET {name}")
first = protocol_line(ser)
if first.startswith("@WS ERR "):
raise ProtocolError(first)
parts = first.split(" ", 3)
if len(parts) != 4 or parts[:2] != ["@WS", "BEGIN"]:
raise ProtocolError(f"bad BEGIN line: {first}")
remote_name = parts[2]
expected_size = int(parts[3])
data = bytearray()
while True:
line = protocol_line(ser, timeout=12.0)
if line.startswith("@WS DATA "):
try:
data.extend(binascii.unhexlify(line[9:]))
except binascii.Error as exc:
raise ProtocolError(f"bad DATA hex for {remote_name}") from exc
continue
if line == f"@WS END {remote_name}":
break
if line.startswith("@WS ERR "):
raise ProtocolError(line)
raise ProtocolError(f"unexpected line while downloading {remote_name}: {line}")
if len(data) != expected_size:
raise ProtocolError(f"{remote_name}: expected {expected_size} bytes, got {len(data)}")
out_path = out_dir / remote_name
out_path.write_bytes(data)
return out_path
def clear_watch(ser) -> str:
send(ser, "CLEAR")
line = protocol_line(ser)
if line.startswith("@WS OK "):
return line[7:]
raise ProtocolError(line)
def yes_no(prompt: str) -> bool:
answer = input(f"{prompt} [y/N] ").strip().lower()
return answer in {"y", "yes"}
def default_output_dir() -> Path:
stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S")
return Path("captures") / stamp
def cmd_ports(_args) -> int:
require_serial()
ports = list(list_ports.comports())
if not ports:
print("No serial ports found.")
return 1
for port in ports:
desc = f" - {port.description}" if port.description else ""
print(f"{port.device}{desc}")
return 0
def cmd_ping(args) -> int:
with open_serial(args.port) as ser:
if ping(ser):
print("Watch responded: PONG")
return 0
print("No PONG from watch.")
return 1
def cmd_list(args) -> int:
with open_serial(args.port) as ser:
files = list_captures(ser)
if not files:
print("No captures stored.")
return 0
for name, size in files:
print(f"{name}\t{size} bytes")
return 0
def cmd_download(args) -> int:
out_dir = args.output
out_dir.mkdir(parents=True, exist_ok=True)
with open_serial(args.port) as ser:
if not ping(ser):
raise SystemExit("watch did not respond to PING")
files = list_captures(ser)
if not files:
print("No captures stored.")
return 0
downloaded: list[Path] = []
for name, size in files:
if size <= 0:
continue
path = download_one(ser, name, out_dir)
downloaded.append(path)
print(f"Downloaded {name} -> {path} ({path.stat().st_size} bytes)")
if downloaded and not args.no_clean:
if yes_no("Clear downloaded captures from the watch now?"):
result = clear_watch(ser)
print(f"Watch cleanup: {result}")
else:
print("Left captures on watch.")
return 0
def cmd_clear(args) -> int:
if not args.yes and not yes_no("Delete capture files and captured-BSSID cache from the watch?"):
print("Cancelled.")
return 0
with open_serial(args.port) as ser:
result = clear_watch(ser)
print(f"Watch cleanup: {result}")
return 0
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
sub = parser.add_subparsers(dest="command", required=True)
sub.add_parser("ports", help="list available serial ports").set_defaults(func=cmd_ports)
for name, help_text in (
("ping", "check watch protocol response"),
("list", "list capture files stored on the watch"),
("download", "download capture files and optionally clear them"),
("clear", "delete capture files from the watch"),
):
p = sub.add_parser(name, help=help_text)
p.add_argument("--port", required=True, help="serial port, e.g. /dev/cu.usbserial-...")
if name == "download":
p.add_argument("-o", "--output", type=Path, default=default_output_dir())
p.add_argument("--no-clean", action="store_true", help="do not ask to clear after download")
if name == "clear":
p.add_argument("-y", "--yes", action="store_true", help="clear without prompting")
p.set_defaults(func=globals()[f"cmd_{name}"])
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
try:
return args.func(args)
except (TimeoutError, ProtocolError, serial.SerialException if serial else Exception) as exc:
print(f"error: {exc}", file=sys.stderr)
return 1
if __name__ == "__main__":
raise SystemExit(main())