46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
"""SHA-256 allowlist for model formats that may invoke Python deserialization."""
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
MANIFEST_PATH = Path(
|
|
os.environ.get("MONITOR_TRUSTED_MODEL_MANIFEST", PROJECT_ROOT / ".trusted-models.json")
|
|
).resolve()
|
|
|
|
|
|
def sha256_file(path):
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def trusted_hashes():
|
|
hashes = {
|
|
item.strip().lower()
|
|
for item in os.environ.get("MONITOR_TRUSTED_MODEL_SHA256", "").split(",")
|
|
if item.strip()
|
|
}
|
|
if MANIFEST_PATH.is_file():
|
|
with open(MANIFEST_PATH, "r", encoding="utf-8") as handle:
|
|
data = json.load(handle)
|
|
entries = data.get("trusted_sha256", []) if isinstance(data, dict) else []
|
|
hashes.update(str(item).strip().lower() for item in entries if item)
|
|
return {item for item in hashes if len(item) == 64 and all(c in "0123456789abcdef" for c in item)}
|
|
|
|
|
|
def require_trusted_model(path):
|
|
path = Path(path).resolve()
|
|
if path.suffix.lower() != ".pt":
|
|
return ""
|
|
actual = sha256_file(path)
|
|
if actual not in trusted_hashes():
|
|
raise PermissionError(
|
|
"untrusted PyTorch model (sha256=%s); approve it with scripts/trust_model.py first" % actual
|
|
)
|
|
return actual
|