150 lines
4.4 KiB
Python
150 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import struct
|
|
import sys
|
|
import zlib
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
VENDOR_ROOT = SCRIPT_DIR / "_vendor"
|
|
if str(VENDOR_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(VENDOR_ROOT))
|
|
|
|
from qrcode.constants import ERROR_CORRECT_M
|
|
from qrcode.main import QRCode
|
|
|
|
|
|
DEFAULT_OUTPUT_DIR = SCRIPT_DIR.parent / "output" / "qr"
|
|
SERIAL_PREFIXES = ("TalkingQ-", "TQ_")
|
|
SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate a binding QR PNG for the banban device bind flow."
|
|
)
|
|
parser.add_argument("device_id", help="Device ID written into the QR payload.")
|
|
parser.add_argument(
|
|
"serial_number",
|
|
help=f"Device serial number. It must start with one of {SERIAL_PREFIXES!r}.",
|
|
)
|
|
parser.add_argument(
|
|
"-o",
|
|
"--output",
|
|
help="Output PNG path. Defaults to talkingq-url/output/qr/<device_id>.png",
|
|
)
|
|
parser.add_argument(
|
|
"--box-size",
|
|
type=int,
|
|
default=10,
|
|
help="Pixel size of one QR module. Default: 10",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def validate_inputs(device_id: str, serial_number: str, box_size: int) -> tuple[str, str]:
|
|
normalized_device_id = device_id.strip()
|
|
normalized_serial_number = serial_number.strip()
|
|
|
|
if not normalized_device_id:
|
|
raise ValueError("device_id cannot be empty")
|
|
if not normalized_serial_number:
|
|
raise ValueError("serial_number cannot be empty")
|
|
if not normalized_serial_number.startswith(SERIAL_PREFIXES):
|
|
raise ValueError(f"serial_number must start with one of {SERIAL_PREFIXES}")
|
|
if box_size <= 0:
|
|
raise ValueError("box_size must be greater than 0")
|
|
|
|
return normalized_device_id, normalized_serial_number
|
|
|
|
|
|
def build_payload(device_id: str, serial_number: str) -> str:
|
|
return json.dumps(
|
|
{"device_id": device_id, "serial_number": serial_number},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
|
|
def build_qr_matrix(payload: str) -> list[list[bool]]:
|
|
qr = QRCode(error_correction=ERROR_CORRECT_M, border=4, box_size=10)
|
|
qr.add_data(payload)
|
|
qr.make(fit=True)
|
|
return qr.get_matrix()
|
|
|
|
|
|
def write_png(path: Path, matrix: list[list[bool]], box_size: int) -> None:
|
|
width = len(matrix[0]) * box_size
|
|
height = len(matrix) * box_size
|
|
raw_rows = bytearray()
|
|
|
|
for row in matrix:
|
|
expanded_row = bytearray()
|
|
for cell in row:
|
|
pixel = 0 if cell else 255
|
|
expanded_row.extend([pixel] * box_size)
|
|
row_bytes = bytes(expanded_row)
|
|
for _ in range(box_size):
|
|
raw_rows.append(0)
|
|
raw_rows.extend(row_bytes)
|
|
|
|
ihdr = struct.pack("!IIBBBBB", width, height, 8, 0, 0, 0, 0)
|
|
compressed = zlib.compress(bytes(raw_rows), level=9)
|
|
|
|
with path.open("wb") as fp:
|
|
fp.write(b"\x89PNG\r\n\x1a\n")
|
|
write_png_chunk(fp, b"IHDR", ihdr)
|
|
write_png_chunk(fp, b"IDAT", compressed)
|
|
write_png_chunk(fp, b"IEND", b"")
|
|
|
|
|
|
def write_png_chunk(fp, chunk_type: bytes, data: bytes) -> None:
|
|
fp.write(struct.pack("!I", len(data)))
|
|
fp.write(chunk_type)
|
|
fp.write(data)
|
|
crc = zlib.crc32(chunk_type)
|
|
crc = zlib.crc32(data, crc)
|
|
fp.write(struct.pack("!I", crc & 0xFFFFFFFF))
|
|
|
|
|
|
def resolve_output_path(device_id: str, output_arg: str | None) -> Path:
|
|
if output_arg:
|
|
return Path(output_arg).expanduser().resolve()
|
|
|
|
safe_name = SAFE_NAME_RE.sub("_", device_id).strip("._") or "bind_qr"
|
|
return (DEFAULT_OUTPUT_DIR / f"{safe_name}.png").resolve()
|
|
|
|
|
|
def write_payload_copy(txt_path: Path, payload: str) -> None:
|
|
txt_path.write_text(payload, encoding="utf-8")
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
try:
|
|
device_id, serial_number = validate_inputs(args.device_id, args.serial_number, args.box_size)
|
|
except ValueError as exc:
|
|
print(f"Error: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
payload = build_payload(device_id, serial_number)
|
|
output_path = resolve_output_path(device_id, args.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
matrix = build_qr_matrix(payload)
|
|
write_png(output_path, matrix, args.box_size)
|
|
write_payload_copy(output_path.with_suffix(".txt"), payload)
|
|
|
|
print(output_path)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|