Files
Artemis-Plugins/plugins/artemis-device-test/skill/scripts/vendor/device_upgrade.py
T

381 lines
14 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""设备 V2 基础信息预检与 EPFP/4001 固件升级工具。"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import re
import select
import socket
import stat
import struct
import sys
import time
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple
INFO_PORT = 8942
UPGRADE_PORT = 8943
INFO_COMMAND = 30101
UPGRADE_COMMAND = 4001
MAX_JSON_SIZE = 1024 * 1024
MAX_STATUS_SIZE = 1024 * 1024
MAX_ROM_SIZE = 0xFFFFFFFF
STATUS_SIZE = 21
class UpgradeError(RuntimeError):
"""表示预检、协议或升级过程失败。"""
def recv_exact(sock: socket.socket, size: int) -> bytes:
"""读取指定字节数;对端提前关闭时立即失败。"""
data = bytearray()
while len(data) < size:
chunk = sock.recv(size - len(data))
if not chunk:
raise UpgradeError("设备在响应完整之前关闭了连接")
data.extend(chunk)
return bytes(data)
def query_device_info(host: str, timeout: float, logger: EventLogger) -> Dict[str, Any]:
"""通过 EPTP/30101 查询并严格校验设备基础信息。"""
session = int(time.time() * 1000) & 0xFFFFFFFF
request_object = {"cmd": INFO_COMMAND, "session": session}
request = json.dumps(request_object, separators=(",", ":")).encode("utf-8")
logger.emit("info_request", host=host, port=INFO_PORT, request=request_object)
with socket.create_connection((host, INFO_PORT), timeout=timeout) as sock:
sock.settimeout(timeout)
sock.sendall(b"EPTP" + struct.pack("!I", len(request)) + request)
header = recv_exact(sock, 8)
if header[:4] != b"EPTP":
raise UpgradeError("基础信息响应同步头不是 EPTP:%r" % (header[:4],))
payload_size = struct.unpack("!I", header[4:])[0]
if payload_size == 0 or payload_size > MAX_JSON_SIZE:
raise UpgradeError("基础信息响应长度不合理:%d" % payload_size)
payload = recv_exact(sock, payload_size)
try:
response = json.loads(payload.decode("utf-8"))
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise UpgradeError("基础信息响应不是有效 UTF-8 JSON:%s" % exc) from exc
if not isinstance(response, dict):
raise UpgradeError("基础信息响应 JSON 顶层不是对象")
if response.get("cmd") != INFO_COMMAND:
raise UpgradeError("基础信息响应 cmd 不匹配:%r" % response.get("cmd"))
if response.get("session") != session:
raise UpgradeError(
"基础信息响应 session 不匹配:期望 %d,实际 %r"
% (session, response.get("session"))
)
if response.get("code") != 0:
raise UpgradeError("基础信息响应 code 非 0:%r" % response.get("code"))
logger.emit("device_info", response=response)
return response
def response_data(response: Dict[str, Any]) -> Dict[str, Any]:
"""取得响应 data 对象,格式错误时拒绝继续。"""
data = response.get("data")
if not isinstance(data, dict):
raise UpgradeError("基础信息响应缺少 data 对象")
return data
def collect_models(data: Dict[str, Any]) -> List[str]:
"""收集 inner、outer 和 hardware 中可用的型号。"""
models: List[str] = []
for section_name in ("inner", "outer", "hardware"):
section = data.get(section_name)
if isinstance(section, dict):
value = section.get("type")
if isinstance(value, str) and value and value not in models:
models.append(value)
return models
def collect_versions(data: Dict[str, Any]) -> Dict[str, Any]:
"""收集各固件分区和硬件版本,便于升级前后比对。"""
versions: Dict[str, Any] = {}
for section_name in ("inner", "outer", "hardware"):
section = data.get(section_name)
if isinstance(section, dict):
versions[section_name] = section.get("version")
return versions
def validate_identity(
response: Dict[str, Any], expected_sn: str, expected_model: Optional[str]
) -> Tuple[Dict[str, Any], List[str]]:
"""对完整 SN 做强校验,并按需校验型号。"""
data = response_data(response)
actual_sn = data.get("sn")
if actual_sn != expected_sn:
raise UpgradeError("设备 SN 不匹配:期望 %s,实际 %r" % (expected_sn, actual_sn))
models = collect_models(data)
if expected_model and expected_model not in models:
raise UpgradeError(
"设备型号不匹配:期望 %s,设备报告 %s"
% (expected_model, models if models else "无可用型号")
)
return data, models
def infer_rom_version(path: Path) -> Optional[str]:
"""仅从路径中识别明确的四段数字版本号,不猜测其他命名。"""
candidates: List[str] = []
pattern = re.compile(r"(?<!\d)(\d+\.\d+\.\d+\.\d+)(?!\d)")
for part in path.parts:
candidates.extend(pattern.findall(part))
return candidates[-1] if candidates else None
def validate_rom(path: Path) -> Dict[str, Any]:
"""确认 ROM 为普通非空文件,并计算 SHA-256。"""
absolute = path.expanduser().resolve(strict=True)
metadata = os.stat(str(absolute))
if not stat.S_ISREG(metadata.st_mode):
raise UpgradeError("ROM 不是普通文件:%s" % absolute)
if metadata.st_size <= 0:
raise UpgradeError("ROM 文件为空:%s" % absolute)
if metadata.st_size > MAX_ROM_SIZE:
raise UpgradeError("ROM 大于协议 uint32 长度上限:%d" % metadata.st_size)
digest = hashlib.sha256()
with absolute.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return {
"path": str(absolute),
"size": metadata.st_size,
"modified": time.strftime(
"%Y-%m-%d %H:%M:%S", time.localtime(metadata.st_mtime)
),
"sha256": digest.hexdigest(),
"target_version_from_path": infer_rom_version(absolute),
}
def parse_status(payload: bytes) -> Dict[str, int]:
"""解析 21 字节升级状态业务数据。"""
if len(payload) < STATUS_SIZE:
raise UpgradeError("升级状态不足 21 字节:%d" % len(payload))
status, session, download_total, downloaded, flash_total, flashed = struct.unpack(
"!BIIIII", payload[:STATUS_SIZE]
)
if status not in (0, 1, 2, 3):
raise UpgradeError("未知升级状态值:%d" % status)
return {
"status": status,
"session": session,
"download_total": download_total,
"downloaded": downloaded,
"flash_total": flash_total,
"flashed": flashed,
}
def extract_status_payloads(buffer: bytearray) -> Iterable[Tuple[str, bytes]]:
"""从 TCP 字节流中提取完整 EPFP 帧或兼容设备实测的裸状态。"""
while buffer:
# 同步头被拆包时等待后续数据,避免把半个 EPFP 当裸状态。
if len(buffer) < 4 and b"EPFP".startswith(bytes(buffer)):
return
if buffer[:4] == b"EPFP":
if len(buffer) < 12:
return
command, payload_size = struct.unpack("!II", buffer[4:12])
if command != UPGRADE_COMMAND:
raise UpgradeError("升级状态帧 cmd 不匹配:%d" % command)
if payload_size < STATUS_SIZE or payload_size > MAX_STATUS_SIZE:
raise UpgradeError("升级状态帧长度不合理:%d" % payload_size)
if len(buffer) < 12 + payload_size:
return
payload = bytes(buffer[12 : 12 + payload_size])
del buffer[: 12 + payload_size]
yield "framed", payload
continue
if buffer[0] in (0, 1, 2, 3):
# 合法裸状态内部也可能包含字节串 EPFP,不能在其业务数据中误重同步。
if len(buffer) < STATUS_SIZE:
return
payload = bytes(buffer[:STATUS_SIZE])
del buffer[:STATUS_SIZE]
yield "raw", payload
continue
sync_position = buffer.find(b"EPFP")
if sync_position > 0:
# 首字节不是合法 status,且后方存在同步头时,丢弃无效前缀。
del buffer[:sync_position]
continue
if len(buffer) < STATUS_SIZE:
return
# 保留这 21 字节交给状态解析器报告明确的非法 status。
payload = bytes(buffer[:STATUS_SIZE])
del buffer[:STATUS_SIZE]
yield "raw", payload
def transfer_rom(
host: str,
rom_path: Path,
rom_size: int,
connect_timeout: float,
stall_timeout: float,
total_timeout: float,
chunk_size: int,
logger: EventLogger,
) -> Dict[str, Any]:
"""发送 ROM,并在同一循环中处理回包、停滞和终态。"""
result: Dict[str, Any] = {"status": None, "bytes_sent": 0}
receive_buffer = bytearray()
pending = bytearray(b"EPFP" + struct.pack("!II", UPGRADE_COMMAND, rom_size))
started = time.monotonic()
last_progress = started
previous_progress: Optional[Tuple[int, int]] = None
upload_complete = False
logger.emit(
"upgrade_connect", host=host, port=UPGRADE_PORT, rom_size=rom_size
)
with socket.create_connection((host, UPGRADE_PORT), timeout=connect_timeout) as sock:
sock.setblocking(False)
with rom_path.open("rb") as rom:
while True:
now = time.monotonic()
if now - started >= total_timeout:
raise UpgradeError("升级超过总超时 %.1f 秒" % total_timeout)
if now - last_progress >= stall_timeout:
raise UpgradeError("设备下载或烧录进度连续 %.1f 秒无变化" % stall_timeout)
readable, writable, exceptional = select.select(
[sock], [] if upload_complete else [sock], [sock], 0.5
)
if exceptional:
raise UpgradeError("升级 TCP 连接出现异常状态")
if readable:
chunk = sock.recv(4096)
if not chunk:
raise UpgradeError("设备在 status=2 之前关闭升级连接")
logger.emit("upgrade_rx", size=len(chunk), hex=chunk.hex())
receive_buffer.extend(chunk)
for frame_type, payload in extract_status_payloads(receive_buffer):
status = parse_status(payload)
logger.emit("device_status", frame=frame_type, **status)
progress = (status["downloaded"], status["flashed"])
if previous_progress is None or progress != previous_progress:
previous_progress = progress
last_progress = time.monotonic()
result.update(status)
if status["status"] == 3:
raise UpgradeError("设备返回升级失败状态 status=3")
if status["status"] == 2:
incomplete_counters = {}
if (
status["download_total"] > 0
and status["downloaded"] < status["download_total"]
):
incomplete_counters["download"] = {
"total": status["download_total"],
"current": status["downloaded"],
}
if (
status["flash_total"] > 0
and status["flashed"] < status["flash_total"]
):
incomplete_counters["flash"] = {
"total": status["flash_total"],
"current": status["flashed"],
}
if incomplete_counters:
logger.emit(
"protocol_warning",
message="status=2,但设备进度计数尚未达到其自身总量",
counters=incomplete_counters,
)
logger.emit("protocol_success", status=2)
return result
if writable:
if not pending:
next_chunk = rom.read(chunk_size)
if next_chunk:
pending.extend(next_chunk)
else:
upload_complete = True
logger.emit(
"client_upload_complete",
bytes_sent=result["bytes_sent"],
rom_size=rom_size,
)
if pending:
sent = sock.send(pending)
if sent <= 0:
raise UpgradeError("升级 TCP 连接无法继续发送")
# 前 12 字节是 EPFP 头,不计入 ROM 上传进度。
previous_wire_bytes = result.get("wire_bytes_sent", 0)
wire_bytes = previous_wire_bytes + sent
result["wire_bytes_sent"] = wire_bytes
result["bytes_sent"] = max(0, min(rom_size, wire_bytes - 12))
del pending[:sent]