Files
AI_Printer/tools/esptool-factory/factory_common/signing.py

183 lines
6.9 KiB
Python

from __future__ import annotations
import base64
import hashlib
import json
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
from factory_common.signing_pubkey import PUBLIC_KEY_PEM
try:
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import dsa, ec, ed25519, ed448, padding, rsa
except Exception as exc: # noqa: BLE001
InvalidSignature = None # type: ignore[assignment]
hashes = None # type: ignore[assignment]
serialization = None # type: ignore[assignment]
dsa = ec = ed25519 = ed448 = padding = rsa = None # type: ignore[assignment]
_CRYPTO_IMPORT_ERROR = exc
else:
_CRYPTO_IMPORT_ERROR = None
class SignatureError(Exception):
pass
def _require_crypto() -> None:
if _CRYPTO_IMPORT_ERROR is not None or serialization is None:
raise SignatureError("cryptography is required for signature operations")
def sig_path_for(path: Path) -> Path:
return Path(str(path) + ".sig")
def _load_public_key(pem: bytes) -> object:
_require_crypto()
if not pem:
raise SignatureError("public key is empty")
try:
return serialization.load_pem_public_key(pem)
except Exception as exc: # noqa: BLE001
raise SignatureError("invalid public key format") from exc
def _load_private_key(path: Path, password: Optional[str]) -> object:
_require_crypto()
try:
data = path.read_bytes()
except OSError as exc:
raise SignatureError(f"unable to read private key: {path}") from exc
password_bytes = password.encode("utf-8") if password else None
try:
return serialization.load_pem_private_key(data, password=password_bytes)
except (TypeError, ValueError) as exc:
raise SignatureError("invalid private key or password") from exc
def _sha256_b64(data: bytes) -> str:
digest = hashlib.sha256(data).digest()
return base64.b64encode(digest).decode("ascii")
def _sign_bytes(data: bytes, private_key: object) -> Tuple[bytes, str]:
_require_crypto()
if ed25519 and isinstance(private_key, ed25519.Ed25519PrivateKey):
return private_key.sign(data), "ed25519"
if ed448 and isinstance(private_key, ed448.Ed448PrivateKey):
return private_key.sign(data), "ed448"
if ec and isinstance(private_key, ec.EllipticCurvePrivateKey):
return private_key.sign(data, ec.ECDSA(hashes.SHA256())), "ecdsa-sha256"
if rsa and isinstance(private_key, rsa.RSAPrivateKey):
signature = private_key.sign(
data,
padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
hashes.SHA256(),
)
return signature, "rsa-pss-sha256"
if dsa and isinstance(private_key, dsa.DSAPrivateKey):
return private_key.sign(data, hashes.SHA256()), "dsa-sha256"
raise SignatureError("unsupported private key type")
def _verify_bytes(data: bytes, signature: bytes, public_key: object, alg_hint: Optional[str]) -> None:
_require_crypto()
try:
if ed25519 and isinstance(public_key, ed25519.Ed25519PublicKey):
public_key.verify(signature, data)
return
if ed448 and isinstance(public_key, ed448.Ed448PublicKey):
public_key.verify(signature, data)
return
if ec and isinstance(public_key, ec.EllipticCurvePublicKey):
public_key.verify(signature, data, ec.ECDSA(hashes.SHA256()))
return
if rsa and isinstance(public_key, rsa.RSAPublicKey):
public_key.verify(
signature,
data,
padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
hashes.SHA256(),
)
return
if dsa and isinstance(public_key, dsa.DSAPublicKey):
public_key.verify(signature, data, hashes.SHA256())
return
except InvalidSignature as exc:
raise SignatureError("signature verification failed") from exc
except Exception as exc: # noqa: BLE001
raise SignatureError("signature verification failed") from exc
raise SignatureError("unsupported public key type")
def _parse_signature(sig_path: Path) -> Tuple[bytes, Dict[str, Any]]:
try:
raw = sig_path.read_bytes()
except OSError as exc:
raise SignatureError(f"unable to read signature file: {sig_path}") from exc
text = raw.decode("utf-8", errors="strict").strip()
try:
payload = json.loads(text)
except json.JSONDecodeError:
payload = None
if isinstance(payload, dict):
sig_b64 = payload.get("sig_b64")
if not isinstance(sig_b64, str):
raise SignatureError("signature file missing sig_b64")
try:
signature = base64.b64decode(sig_b64.encode("ascii"))
except Exception as exc: # noqa: BLE001
raise SignatureError("invalid base64 signature") from exc
return signature, payload
try:
signature = base64.b64decode(text.encode("ascii"))
except Exception as exc: # noqa: BLE001
raise SignatureError("invalid signature format") from exc
return signature, {}
def sign_file(path: Path, private_key_path: Path, password: Optional[str] = None) -> Path:
data = path.read_bytes()
private_key = _load_private_key(private_key_path, password)
signature, alg = _sign_bytes(data, private_key)
payload = {
"alg": alg,
"hash": "sha256",
"hash_b64": _sha256_b64(data),
"sig_b64": base64.b64encode(signature).decode("ascii"),
}
sig_path = sig_path_for(path)
sig_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
return sig_path
def verify_file_signature(path: Path) -> None:
sig_path = sig_path_for(path)
if not sig_path.exists():
raise SignatureError(f"missing signature file: {sig_path.name}")
signature, payload = _parse_signature(sig_path)
data = path.read_bytes()
hash_b64 = payload.get("hash_b64")
if isinstance(hash_b64, str):
current = _sha256_b64(data)
if current != hash_b64:
raise SignatureError("content hash mismatch")
public_key = _load_public_key(PUBLIC_KEY_PEM)
alg_hint = payload.get("alg") if isinstance(payload, dict) else None
_verify_bytes(data, signature, public_key, alg_hint)
def verify_payload_signature(payload: bytes, sig_b64: str, alg_hint: Optional[str] = None) -> None:
_require_crypto()
if not isinstance(sig_b64, str) or not sig_b64.strip():
raise SignatureError("missing signature")
try:
signature = base64.b64decode(sig_b64.encode("ascii"))
except Exception as exc: # noqa: BLE001
raise SignatureError("invalid base64 signature") from exc
public_key = _load_public_key(PUBLIC_KEY_PEM)
_verify_bytes(payload, signature, public_key, alg_hint)