322 lines
12 KiB
Python
322 lines
12 KiB
Python
import base64
|
|
import copy
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
TIER_ORDER = ("starter", "basic", "pro", "max")
|
|
TIER_ALIASES = {
|
|
"entry": "basic",
|
|
"entry tier": "basic",
|
|
"mid": "pro",
|
|
"mid tier": "pro",
|
|
"advanced": "max",
|
|
"advanced tier": "max",
|
|
"starter tier": "starter",
|
|
"basic tier": "basic",
|
|
"pro tier": "pro",
|
|
"max tier": "max",
|
|
}
|
|
STAGING_MAX_TIER = "basic"
|
|
SUPPORTED_SIGNATURE_ALGS = {"rsa-sha256", "sha256-rsa"}
|
|
|
|
|
|
def utcnow_iso():
|
|
return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def parse_time(value):
|
|
text = (value or "").strip()
|
|
if not text:
|
|
return None
|
|
if len(text) == 10:
|
|
text += "T00:00:00Z"
|
|
if text.endswith("Z"):
|
|
text = text[:-1] + "+00:00"
|
|
try:
|
|
return datetime.fromisoformat(text)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def normalize_tier(value, default="basic"):
|
|
raw = (value or "").strip().lower()
|
|
if raw in TIER_ORDER:
|
|
return raw
|
|
return TIER_ALIASES.get(raw, default)
|
|
|
|
|
|
def tier_rank(tier):
|
|
return TIER_ORDER.index(normalize_tier(tier))
|
|
|
|
|
|
def tier_lte(left, right):
|
|
return tier_rank(left) <= tier_rank(right)
|
|
|
|
|
|
def min_tier(*tiers):
|
|
items = [normalize_tier(t) for t in tiers if t]
|
|
if not items:
|
|
return STAGING_MAX_TIER
|
|
return sorted(items, key=tier_rank)[0]
|
|
|
|
|
|
def read_json_file(path):
|
|
try:
|
|
data = json.loads(Path(path).read_text())
|
|
return data if isinstance(data, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def write_json_file(path, payload):
|
|
p = Path(path)
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
|
p.write_text(json.dumps(payload, indent=2))
|
|
|
|
|
|
def _canonical_payload(payload):
|
|
clean = copy.deepcopy(payload or {})
|
|
for key in ("signature", "signature_status", "verification_status"):
|
|
clean.pop(key, None)
|
|
return json.dumps(clean, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
|
|
|
|
|
def verify_signature(payload, public_key_path):
|
|
if not payload:
|
|
return False, "missing"
|
|
signature = (payload.get("signature") or "").strip()
|
|
if not signature:
|
|
return False, "missing_signature"
|
|
alg = (payload.get("signature_alg") or "").strip().lower()
|
|
if alg not in SUPPORTED_SIGNATURE_ALGS:
|
|
return False, "unsupported_signature_alg"
|
|
key_path = Path(public_key_path)
|
|
if not key_path.exists():
|
|
return False, "missing_public_key"
|
|
canonical = _canonical_payload(payload).encode()
|
|
try:
|
|
sig_bytes = base64.b64decode(signature, validate=True)
|
|
except Exception:
|
|
return False, "bad_signature_encoding"
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False) as data_file, tempfile.NamedTemporaryFile(delete=False) as sig_file:
|
|
data_file.write(canonical)
|
|
sig_file.write(sig_bytes)
|
|
data_file.flush()
|
|
sig_file.flush()
|
|
data_path = data_file.name
|
|
sig_path = sig_file.name
|
|
try:
|
|
result = subprocess.run(
|
|
["openssl", "dgst", "-sha256", "-verify", str(key_path), "-signature", sig_path, data_path],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
return result.returncode == 0, "valid" if result.returncode == 0 else "invalid_signature"
|
|
finally:
|
|
for tmp in (data_path, sig_path):
|
|
try:
|
|
os.unlink(tmp)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def public_license_record(record):
|
|
public = {k: v for k, v in (record or {}).items() if k != "signature"}
|
|
if record.get("license_key"):
|
|
public["license_key_prefix"] = str(record.get("license_key"))[:12]
|
|
public["license_key_present"] = True
|
|
else:
|
|
public["license_key_present"] = False
|
|
return public
|
|
|
|
|
|
def evaluate_license(payload, public_key_path, machine_id=None, now=None):
|
|
now_dt = parse_time(now) if isinstance(now, str) else now or datetime.now(timezone.utc)
|
|
base = {
|
|
"schema": "cezen.license.v2",
|
|
"status": "missing",
|
|
"valid": False,
|
|
"allowed_tier": STAGING_MAX_TIER,
|
|
"feature_overrides": {},
|
|
"install_type": "field-staging",
|
|
"needs_activation": True,
|
|
"license_record": {},
|
|
"notes": [],
|
|
}
|
|
if not payload:
|
|
return base
|
|
|
|
record = payload if isinstance(payload, dict) else {}
|
|
base["license_record"] = public_license_record(record)
|
|
base["install_type"] = (record.get("install_type") or "licensed").strip() or "licensed"
|
|
base["allowed_tier"] = normalize_tier(record.get("allowed_tier") or record.get("tier"), STAGING_MAX_TIER)
|
|
overrides = record.get("feature_overrides") or {}
|
|
base["feature_overrides"] = overrides if isinstance(overrides, dict) else {}
|
|
|
|
ok, sig_status = verify_signature(record, public_key_path)
|
|
if not ok:
|
|
base["status"] = "invalid_signature" if sig_status != "missing_signature" else "missing"
|
|
base["notes"].append(sig_status)
|
|
return base
|
|
|
|
issued_at = parse_time(record.get("issued_at"))
|
|
if issued_at and issued_at > now_dt:
|
|
base["status"] = "not_yet_valid"
|
|
base["notes"].append("issued_at_in_future")
|
|
return base
|
|
|
|
expires_at = parse_time(record.get("expires_at"))
|
|
if expires_at and expires_at < now_dt:
|
|
base["status"] = "expired"
|
|
base["notes"].append("license_expired")
|
|
return base
|
|
|
|
machine_binding = record.get("machine_binding") or {}
|
|
if machine_id and isinstance(machine_binding, dict):
|
|
bound_id = (machine_binding.get("machine_id") or "").strip()
|
|
if bound_id and bound_id != machine_id:
|
|
base["status"] = "machine_mismatch"
|
|
base["notes"].append("machine_binding_mismatch")
|
|
return base
|
|
|
|
base["status"] = "valid"
|
|
base["valid"] = True
|
|
base["needs_activation"] = False
|
|
return base
|
|
|
|
|
|
def evaluate_override(payload, public_key_path, now=None):
|
|
if not payload:
|
|
return {"status": "missing", "valid": False, "allow_hardware_mismatch": False, "max_override_tier": None}
|
|
record = payload if isinstance(payload, dict) else {}
|
|
ok, sig_status = verify_signature(record, public_key_path)
|
|
if not ok:
|
|
return {
|
|
"status": "invalid_signature" if sig_status != "missing_signature" else "missing",
|
|
"valid": False,
|
|
"allow_hardware_mismatch": False,
|
|
"max_override_tier": None,
|
|
}
|
|
expires_at = parse_time(record.get("expires_at"))
|
|
now_dt = parse_time(now) if isinstance(now, str) else now or datetime.now(timezone.utc)
|
|
if expires_at and expires_at < now_dt:
|
|
return {"status": "expired", "valid": False, "allow_hardware_mismatch": False, "max_override_tier": None}
|
|
return {
|
|
"status": "valid",
|
|
"valid": True,
|
|
"allow_hardware_mismatch": bool(record.get("allow_hardware_mismatch")),
|
|
"max_override_tier": normalize_tier(record.get("max_override_tier"), "max"),
|
|
"override_record": public_license_record(record),
|
|
}
|
|
|
|
|
|
def build_tier_options(license_eval, hardware_tier, override_eval=None):
|
|
override_eval = override_eval or {}
|
|
licensed_cap = normalize_tier(license_eval.get("allowed_tier"), STAGING_MAX_TIER)
|
|
hardware_cap = normalize_tier(hardware_tier, STAGING_MAX_TIER)
|
|
override_cap = normalize_tier(override_eval.get("max_override_tier"), "max") if override_eval.get("allow_hardware_mismatch") else None
|
|
options = []
|
|
for tier in TIER_ORDER:
|
|
state = "enabled"
|
|
reason = ""
|
|
if not tier_lte(tier, licensed_cap):
|
|
state = "disabled_by_license"
|
|
reason = f"Not included in current license ({licensed_cap})."
|
|
elif not tier_lte(tier, hardware_cap):
|
|
if override_cap and tier_lte(tier, override_cap):
|
|
state = "override_required"
|
|
reason = f"Requires Cezen override; hardware recommendation is {hardware_cap}."
|
|
else:
|
|
state = "disabled_by_hardware"
|
|
reason = f"Requires larger hardware; recommendation is {hardware_cap}."
|
|
options.append({"tier": tier, "state": state, "reason": reason, "selectable": state in {"enabled", "override_required"}})
|
|
return options
|
|
|
|
|
|
def provisioned_components_from_record(record):
|
|
record = record or {}
|
|
components = record.get("components") or {}
|
|
if isinstance(components, dict):
|
|
return components
|
|
skip_roles = set(record.get("skip_roles") or [])
|
|
return {
|
|
"ollama": "ollama" not in skip_roles,
|
|
"jupyterlab": "jupyterlab" not in skip_roles,
|
|
"chromadb": "chromadb" not in skip_roles,
|
|
"vllm": "vllm" not in skip_roles,
|
|
"mlflow": "mlflow" not in skip_roles,
|
|
"minio": "minio" not in skip_roles,
|
|
"monitoring": "monitoring" not in skip_roles,
|
|
"k3s": "k3s" not in skip_roles,
|
|
}
|
|
|
|
|
|
def resolve_effective_features(base_features, feature_overrides, provisioned_components, hardware_features):
|
|
features = copy.deepcopy(base_features or {})
|
|
for key, value in (feature_overrides or {}).items():
|
|
features[key] = value
|
|
|
|
provisioned_components = provisioned_components or {}
|
|
hardware_features = hardware_features or {}
|
|
|
|
if features.get("rag") and not provisioned_components.get("chromadb", True):
|
|
features["rag"] = False
|
|
if features.get("gpu_inference") and not (
|
|
hardware_features.get("ollama_gpu") or hardware_features.get("vllm")
|
|
):
|
|
features["gpu_inference"] = False
|
|
if features.get("fine_tuning") and not (
|
|
provisioned_components.get("jupyterlab", False) and hardware_features.get("fine_tuning_qlora")
|
|
):
|
|
features["fine_tuning"] = False
|
|
if features.get("deepspeed") and not hardware_features.get("distributed_training"):
|
|
features["deepspeed"] = False
|
|
return features
|
|
|
|
|
|
def collect_entitlement(tier_matrix, license_record, install_record, feasibility, public_key_path, machine_id=None):
|
|
license_eval = evaluate_license(license_record, public_key_path, machine_id=machine_id)
|
|
override_eval = evaluate_override((install_record or {}).get("override_record") or read_json_file((install_record or {}).get("override_path", "")), public_key_path)
|
|
licensed_tier = normalize_tier(license_eval.get("allowed_tier"), STAGING_MAX_TIER)
|
|
provisioned_tier = normalize_tier((install_record or {}).get("provisioned_tier") or (install_record or {}).get("selected_tier"), STAGING_MAX_TIER)
|
|
hardware_tier = normalize_tier(((feasibility or {}).get("recommendation") or {}).get("recommended_tier"), STAGING_MAX_TIER)
|
|
components = provisioned_components_from_record(install_record)
|
|
tier_defaults = copy.deepcopy((tier_matrix or {}).get(licensed_tier, {}))
|
|
effective_features = resolve_effective_features(
|
|
tier_defaults.get("features", {}),
|
|
license_eval.get("feature_overrides", {}),
|
|
components,
|
|
(feasibility or {}).get("features", {}),
|
|
)
|
|
|
|
states = []
|
|
if tier_rank(licensed_tier) > tier_rank(provisioned_tier):
|
|
states.append("licensed_not_provisioned")
|
|
if tier_rank(provisioned_tier) > tier_rank(licensed_tier):
|
|
states.append("provisioned_above_license")
|
|
if tier_rank(provisioned_tier) > tier_rank(hardware_tier) and not override_eval.get("allow_hardware_mismatch"):
|
|
states.append("hardware_below_provisioned")
|
|
if license_eval.get("status") != "valid":
|
|
states.append(license_eval.get("status"))
|
|
|
|
return {
|
|
"license_status": license_eval.get("status"),
|
|
"license_valid": license_eval.get("valid"),
|
|
"license_record": public_license_record(license_record),
|
|
"licensed_tier": licensed_tier,
|
|
"provisioned_tier": provisioned_tier,
|
|
"hardware_recommended_tier": hardware_tier,
|
|
"effective_features": effective_features,
|
|
"feature_overrides": license_eval.get("feature_overrides", {}),
|
|
"provisioned_components": components,
|
|
"compliance_state": states or ["licensed_and_provisioned"],
|
|
"override": override_eval,
|
|
"install_type": license_eval.get("install_type") or (install_record or {}).get("install_type") or "field-staging",
|
|
}
|