154 lines
4.9 KiB
Python
154 lines
4.9 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
try:
|
||
from cryptography.exceptions import InvalidTag
|
||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||
except Exception as exc: # noqa: BLE001
|
||
AESGCM = None # type: ignore[assignment]
|
||
InvalidTag = Exception # type: ignore[assignment]
|
||
_CRYPTO_IMPORT_ERROR = exc
|
||
else:
|
||
_CRYPTO_IMPORT_ERROR = None
|
||
|
||
|
||
MAGIC = b"TQFWENC1"
|
||
NONCE_SIZE = 12
|
||
SHA256_SIZE = 32
|
||
MIN_PAYLOAD_SIZE = len(MAGIC) + NONCE_SIZE + SHA256_SIZE + 16
|
||
DEFAULT_FIRMWARE_KEY_RELATIVE_PATH = Path("keys") / "firmware" / "fw_key.bin"
|
||
|
||
|
||
class FirmwareEnvelopeError(Exception):
|
||
pass
|
||
|
||
|
||
def _require_crypto() -> None:
|
||
if _CRYPTO_IMPORT_ERROR is not None or AESGCM is None:
|
||
raise FirmwareEnvelopeError("cryptography 不可用，无法处理 .enc 固件。")
|
||
|
||
|
||
def _parse_key_material(raw: bytes, source: str) -> bytes:
|
||
if len(raw) == 32:
|
||
return raw
|
||
|
||
try:
|
||
text = raw.decode("utf-8").strip()
|
||
except UnicodeDecodeError as exc:
|
||
raise FirmwareEnvelopeError(f"固件密钥格式无效（{source}）：既不是 32 字节原始密钥，也不是文本编码。") from exc
|
||
|
||
if not text:
|
||
raise FirmwareEnvelopeError(f"固件密钥为空（{source}）。")
|
||
|
||
lowered = text.lower()
|
||
if lowered.startswith("hex:"):
|
||
text = text[4:].strip()
|
||
if len(text) == 64:
|
||
try:
|
||
key = bytes.fromhex(text)
|
||
except ValueError as exc:
|
||
raise FirmwareEnvelopeError(f"固件密钥 HEX 无效（{source}）。") from exc
|
||
if len(key) == 32:
|
||
return key
|
||
|
||
if lowered.startswith("base64:"):
|
||
text = text[7:].strip()
|
||
try:
|
||
key = base64.b64decode(text.encode("utf-8"), validate=True)
|
||
except Exception: # noqa: BLE001
|
||
key = b""
|
||
if len(key) == 32:
|
||
return key
|
||
|
||
raise FirmwareEnvelopeError(
|
||
f"固件密钥长度无效（{source}）：需要 32 字节，或 64 位 HEX，或 32 字节 Base64。"
|
||
)
|
||
|
||
|
||
def _iter_runtime_roots() -> Iterable[Path]:
|
||
exe_path = Path(sys.executable).resolve()
|
||
yield exe_path.parent
|
||
for parent in exe_path.parents:
|
||
if parent.suffix.lower() == ".app":
|
||
yield parent.parent
|
||
break
|
||
|
||
|
||
def _candidate_key_paths() -> Iterable[Path]:
|
||
if getattr(sys, "frozen", False):
|
||
for root in _iter_runtime_roots():
|
||
yield root / DEFAULT_FIRMWARE_KEY_RELATIVE_PATH
|
||
|
||
for parent in Path(__file__).resolve().parents:
|
||
yield parent / DEFAULT_FIRMWARE_KEY_RELATIVE_PATH
|
||
|
||
|
||
def _load_key_from_fixed_path(purpose: str) -> bytes:
|
||
for path in _candidate_key_paths():
|
||
if not path.exists():
|
||
continue
|
||
try:
|
||
raw = path.read_bytes()
|
||
except OSError as exc:
|
||
raise FirmwareEnvelopeError(f"{purpose}密钥文件读取失败：{path}") from exc
|
||
return _parse_key_material(raw, str(path))
|
||
|
||
raise FirmwareEnvelopeError(
|
||
f"缺少{purpose}密钥文件。请在运行目录提供 {DEFAULT_FIRMWARE_KEY_RELATIVE_PATH}。"
|
||
)
|
||
|
||
|
||
def load_firmware_encrypt_key() -> bytes:
|
||
return _load_key_from_fixed_path(purpose="固件加密")
|
||
|
||
|
||
def load_firmware_decrypt_key() -> bytes:
|
||
return _load_key_from_fixed_path(purpose="固件解密")
|
||
|
||
|
||
def encrypt_firmware_blob(plain: bytes, key: bytes) -> bytes:
|
||
_require_crypto()
|
||
if len(key) != 32:
|
||
raise FirmwareEnvelopeError("固件加密密钥长度无效，必须为 32 字节。")
|
||
nonce = os.urandom(NONCE_SIZE)
|
||
digest = hashlib.sha256(plain).digest()
|
||
aes = AESGCM(key)
|
||
ciphertext = aes.encrypt(nonce, plain, MAGIC)
|
||
return MAGIC + nonce + digest + ciphertext
|
||
|
||
|
||
def decrypt_firmware_blob(payload: bytes, key: bytes) -> bytes:
|
||
_require_crypto()
|
||
if len(key) != 32:
|
||
raise FirmwareEnvelopeError("固件解密密钥长度无效，必须为 32 字节。")
|
||
if len(payload) < MIN_PAYLOAD_SIZE:
|
||
raise FirmwareEnvelopeError("固件 .enc 文件格式无效：长度不足。")
|
||
if payload[: len(MAGIC)] != MAGIC:
|
||
raise FirmwareEnvelopeError("固件 .enc 文件格式无效：magic 不匹配。")
|
||
|
||
cursor = len(MAGIC)
|
||
nonce = payload[cursor : cursor + NONCE_SIZE]
|
||
cursor += NONCE_SIZE
|
||
expected_digest = payload[cursor : cursor + SHA256_SIZE]
|
||
cursor += SHA256_SIZE
|
||
ciphertext = payload[cursor:]
|
||
|
||
aes = AESGCM(key)
|
||
try:
|
||
plain = aes.decrypt(nonce, ciphertext, MAGIC)
|
||
except InvalidTag as exc:
|
||
raise FirmwareEnvelopeError("固件 .enc 解密失败：密钥错误或密文损坏。") from exc
|
||
except Exception as exc: # noqa: BLE001
|
||
raise FirmwareEnvelopeError("固件 .enc 解密失败。") from exc
|
||
|
||
digest = hashlib.sha256(plain).digest()
|
||
if digest != expected_digest:
|
||
raise FirmwareEnvelopeError("固件 .enc 解密后哈希不一致。")
|
||
return plain
|