103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
"""PlanCert and PlanProof helpers.
|
|
|
|
Lightweight, dependency-free PlanCert representation and HMAC-based signing
|
|
for constrained environments. This is intentionally small: in production one
|
|
would replace the HMAC signer with a proper asymmetric signature scheme
|
|
and a compact binary encoding (CBOR) for wire efficiency.
|
|
|
|
Provides:
|
|
- PlanCert: declarative certificate describing plan invariants and provenance
|
|
- PlanProof: verifier result (accept/reject + score + optional counterexample)
|
|
- helpers: sign_plan_cert, verify_plan_cert, make_plan_proof
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hmac
|
|
import hashlib
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
def _canonical_json(obj: Any) -> str:
|
|
# Deterministic serialization used for signatures
|
|
return json.dumps(obj, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
@dataclass
|
|
class PlanCert:
|
|
plan_id: str
|
|
preconditions: Dict[str, Any]
|
|
postconditions: Dict[str, Any]
|
|
invariants: Dict[str, Any]
|
|
worst_case_energy: Optional[float] = None
|
|
wcet_ms: Optional[int] = None
|
|
# Optional maximum actuator exposure (compact risk metric)
|
|
max_actuator_exposure: Optional[float] = None
|
|
# Criticality level (e.g. low, normal, critical) used by admission control
|
|
criticality: Optional[str] = None
|
|
# Optional rollback checkpoint: small dict containing a checkpoint_hash and revert window
|
|
rollback_checkpoint: Optional[Dict[str, Any]] = None
|
|
provenance: Optional[Dict[str, Any]] = None
|
|
signature: Optional[str] = None # hex HMAC-SHA256 signature when signed
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
d = asdict(self)
|
|
# don't include signature when computing signatures
|
|
return {k: v for k, v in d.items() if k != "signature"}
|
|
|
|
def to_json(self) -> str:
|
|
return _canonical_json(self.to_dict())
|
|
|
|
|
|
@dataclass
|
|
class PlanProof:
|
|
plan_id: str
|
|
accepted: bool
|
|
score: float
|
|
counterexample: Optional[Dict[str, Any]] = None
|
|
verifier: Optional[str] = None
|
|
signature: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
# Exclude signature field when producing the canonical dict used for
|
|
# signing/verification to avoid the signature covering itself.
|
|
d = asdict(self)
|
|
return {k: v for k, v in d.items() if k != "signature"}
|
|
|
|
def to_json(self) -> str:
|
|
return _canonical_json(self.to_dict())
|
|
|
|
|
|
def sign_plan_cert(cert: PlanCert, key: bytes) -> PlanCert:
|
|
"""Sign PlanCert with HMAC-SHA256 using symmetric key.
|
|
|
|
The signature covers all fields except `signature` itself.
|
|
"""
|
|
data = cert.to_json().encode()
|
|
sig = hmac.new(key, data, hashlib.sha256).hexdigest()
|
|
cert.signature = sig
|
|
return cert
|
|
|
|
|
|
def verify_plan_cert(cert: PlanCert, key: bytes) -> bool:
|
|
if not cert.signature:
|
|
return False
|
|
expected = hmac.new(key, cert.to_json().encode(), hashlib.sha256).hexdigest()
|
|
# Use hmac.compare_digest for timing-safe comparison
|
|
return hmac.compare_digest(expected, cert.signature)
|
|
|
|
|
|
def make_plan_proof(cert: PlanCert, accepted: bool, score: float, verifier: str, key: Optional[bytes] = None, counterexample: Optional[Dict[str, Any]] = None) -> PlanProof:
|
|
proof = PlanProof(plan_id=cert.plan_id, accepted=accepted, score=score, counterexample=counterexample, verifier=verifier)
|
|
if key is not None:
|
|
proof.signature = hmac.new(key, proof.to_json().encode(), hashlib.sha256).hexdigest()
|
|
return proof
|
|
|
|
|
|
def verify_plan_proof(proof: PlanProof, key: bytes) -> bool:
|
|
if not proof.signature:
|
|
return False
|
|
expected = hmac.new(key, proof.to_json().encode(), hashlib.sha256).hexdigest()
|
|
return hmac.compare_digest(expected, proof.signature)
|