pax_coder_gate_pq.py

#!/usr/bin/env python3
"""

PAX-Coder Post-Quantum Protected Execution Gate



ML-DSA-44 (CRYSTALS-Dilithium, NIST FIPS 204) upgrade of pax_coder_gate.py.



Drop-in replacement: same 5-step gate, same exit codes, same environment

variables — but signature verification uses ML-DSA-44 instead of Ed25519.



Ed25519 is broken by Shor's algorithm. Every capability token signed with

Ed25519 is a liability against a quantum adversary in harvest-now-decrypt-later

mode. ML-DSA-44 provides 128-bit post-quantum security (NIST level 2).



Key size changes:

  Ed25519 public key:  32 bytes  (64 hex chars in PEM)

  ML-DSA-44 public key: 1312 bytes

  Ed25519 signature:   64 bytes  (128 hex chars)

  ML-DSA-44 signature: 2420 bytes (4840 hex chars)



Token format:

  PAX_CAPABILITY_TOKEN = <json_payload>|<mldsa44_signature_hex>



Authority public key:

  PAX_MLDSA_PUBLIC_KEY_HEX  env var  OR  sovereign/mldsa_authority.pub (hex file)



Exit codes (unchanged from Ed25519 gate):

  0 = AUTHORIZED

  1 = INTEGRITY_FAILED

  2 = AUTHORIZATION_DENIED

  3 = SCRIPT_ERROR



Authors: Ahmad Ali Parr, Jessica L. Williams (SNAPKITTYWEST)

Part of: worm-engines LOCKER / PAX-Coder sovereign stack

"""

import hashlib
import json
import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from pydantic import BaseModel, Field, ValidationError

# ── ML-DSA-44 constants (NIST FIPS 204) ──────────────────────────────────────

MLDSA_PK_BYTES  = 1312
MLDSA_SIG_BYTES = 2420
MLDSA_PK_HEX    = MLDSA_PK_BYTES  * 2   # 2624 hex chars
MLDSA_SIG_HEX   = MLDSA_SIG_BYTES * 2   # 4840 hex chars

# ── Exit codes ────────────────────────────────────────────────────────────────

EXIT_AUTHORIZED        = 0
EXIT_INTEGRITY_FAILED  = 1
EXIT_DENIED            = 2
EXIT_ERROR             = 3

# ── ML-DSA-44 Python binding ──────────────────────────────────────────────────
# We use the `dilithium-py` package (pip install dilithium-py) which implements
# ML-DSA-44 in pure Python matching NIST FIPS 204.
# Fallback: subprocess to the sovereign-trinity-kernel Dex .so via cffi.

def _import_mldsa():
    """Try to import dilithium-py. Returns verify function or None."""
    try:
        from dilithium_py.dilithium import Dilithium2  # ML-DSA-44 = Dilithium2
        return Dilithium2
    except ImportError:
        return None

MLDSA = _import_mldsa()

def mldsa_verify(public_key_bytes: bytes, message: bytes, signature: bytes) -> bool:
    """

    Verify an ML-DSA-44 signature.



    Args:

        public_key_bytes: 1312-byte ML-DSA-44 public key

        message:          message that was signed

        signature:        2420-byte ML-DSA-44 signature



    Returns:

        True if valid, False if invalid

    """
    if len(public_key_bytes) != MLDSA_PK_BYTES:
        return False
    if len(signature) != MLDSA_SIG_BYTES:
        return False

    if MLDSA is not None:
        try:
            return MLDSA.verify(public_key_bytes, message, signature)
        except Exception:
            return False

    # Fallback: call sovereign-trinity-kernel Rust verifier via subprocess
    # (requires pax_verify_mldsa binary built from worm-engines)
    binary = os.environ.get("PAX_MLDSA_VERIFY_BIN", "pax_verify_mldsa")
    try:
        import tempfile
        with tempfile.TemporaryDirectory() as tmpdir:
            pk_path  = Path(tmpdir) / "pk.bin"
            msg_path = Path(tmpdir) / "msg.bin"
            sig_path = Path(tmpdir) / "sig.bin"
            pk_path.write_bytes(public_key_bytes)
            msg_path.write_bytes(message)
            sig_path.write_bytes(signature)
            result = subprocess.run(
                [binary, str(pk_path), str(msg_path), str(sig_path)],
                capture_output=True, timeout=10
            )
            return result.returncode == 0
    except Exception:
        return False

# ── Pydantic models ───────────────────────────────────────────────────────────

class CapabilityPayload(BaseModel):
    node_id:    str = Field(min_length=1)
    release_id: str = Field(min_length=1)
    commit:     str = Field(min_length=1)
    nonce:      str = Field(min_length=1)
    expires_at: str = Field(min_length=1)

class ReleaseMetadata(BaseModel):
    project:              str = ""
    repository:           str = ""
    release_version:      str = ""
    git_commit:           str = ""
    node_id:              str = ""
    manifest_sha256:      str = ""
    release_timestamp_utc: str = ""

class AuthorizationRecord(BaseModel):
    authorization_id:     str = ""
    node_id:              str = ""
    authorization_status: str = ""

# ── Gate ──────────────────────────────────────────────────────────────────────

class PaxCoderGatePQ:
    """

    PAX-Coder post-quantum execution gate (ML-DSA-44).



    Identical flow to PaxCoderGate but Step 5 uses ML-DSA-44.

    """

    def __init__(self, repo_root: Optional[Path] = None, quiet: bool = False):
        if repo_root is None:
            repo_root = Path(__file__).resolve().parent
        self.repo_root    = Path(repo_root)
        self.sovereign_dir = self.repo_root / "sovereign"
        self.quiet        = quiet
        self._messages: list[str] = []

    def log(self, msg: str) -> None:
        self._messages.append(msg)
        if not self.quiet:
            print(msg)

    def run(self) -> int:
        self.log("==========================================")
        self.log("PAX-CODER PQ GATE (ML-DSA-44 / FIPS 204)")
        self.log("==========================================")
        self.log("")

        r = self._verify_release_integrity()
        if r != EXIT_AUTHORIZED: return r

        r = self._verify_node_authorization()
        if r != EXIT_AUTHORIZED: return r

        raw = self._get_capability_token()
        if raw is None: return EXIT_DENIED

        payload, sig_hex = self._parse_capability(raw)
        if payload is None: return EXIT_DENIED

        r = self._validate_capability(payload)
        if r != EXIT_AUTHORIZED: return r

        r = self._verify_mldsa_signature(payload, sig_hex)
        if r != EXIT_AUTHORIZED: return r

        self.log("")
        self.log("==========================================")
        self.log("STATUS: AUTHORIZATION_GRANTED (ML-DSA-44)")
        self.log("==========================================")
        self.log(f"Capability valid until: {payload.expires_at}")
        return EXIT_AUTHORIZED

    # ── Step 1: Release integrity ─────────────────────────────────────────────

    def _verify_release_integrity(self) -> int:
        self.log("[1/5] Verifying release integrity...")
        release_file = self.sovereign_dir / "release.json"
        if not release_file.exists():
            self.log("FAILED: sovereign/release.json not found")
            return EXIT_INTEGRITY_FAILED
        try:
            release = ReleaseMetadata(**json.loads(release_file.read_text()))
        except (json.JSONDecodeError, ValidationError) as e:
            self.log(f"FAILED: Cannot parse release.json: {e}")
            return EXIT_INTEGRITY_FAILED

        current = self._get_git_commit()
        if current is None:
            self.log("FAILED: Cannot determine git commit")
            return EXIT_ERROR
        if current != release.git_commit:
            self.log("FAILED: Release integrity check failed")
            return EXIT_INTEGRITY_FAILED

        manifest = self.sovereign_dir / "manifest.json"
        if manifest.exists() and release.manifest_sha256:
            if self._sha256_file(manifest) != release.manifest_sha256:
                self.log("FAILED: Manifest hash mismatch")
                return EXIT_INTEGRITY_FAILED

        self.log("  OK: Release integrity verified")
        return EXIT_AUTHORIZED

    # ── Step 2: Node authorization ────────────────────────────────────────────

    def _verify_node_authorization(self) -> int:
        self.log("[2/5] Verifying node authorization...")
        auth_file = self.sovereign_dir / "authorization.json"
        if not auth_file.exists():
            self.log("DENIED: sovereign/authorization.json not found")
            return EXIT_DENIED
        try:
            auth = AuthorizationRecord(**json.loads(auth_file.read_text()))
        except (json.JSONDecodeError, ValidationError) as e:
            self.log(f"DENIED: Cannot parse authorization.json: {e}")
            return EXIT_DENIED
        if auth.authorization_status.upper() != "ACTIVE":
            self.log(f"DENIED: Node status is {auth.authorization_status}")
            return EXIT_DENIED
        self.log("  OK: Node is ACTIVE")
        return EXIT_AUTHORIZED

    # ── Step 3: Capability possession ────────────────────────────────────────

    def _get_capability_token(self) -> Optional[str]:
        self.log("[3/5] Locating capability token...")
        token = os.environ.get("PAX_CAPABILITY_TOKEN")
        if token:
            self.log("  OK: Token found in PAX_CAPABILITY_TOKEN")
            return token
        cap_file = self.sovereign_dir / "capability.token"
        if cap_file.exists():
            self.log("  OK: Token found in sovereign/capability.token")
            return cap_file.read_text().strip()
        self.log("DENIED: No capability token found")
        return None

    # ── Step 4: Validate capability ───────────────────────────────────────────

    def _parse_capability(self, raw: str):
        self.log("[4/5] Parsing capability token...")
        parts = raw.strip().split("|")
        if len(parts) != 2:
            self.log("DENIED: Token format invalid (expected JSON|sig_hex)")
            return None, None
        try:
            data    = json.loads(parts[0])
            payload = CapabilityPayload(**data)
        except (json.JSONDecodeError, ValidationError) as e:
            self.log(f"DENIED: Token payload invalid: {e}")
            return None, None
        sig_hex = parts[1].strip()
        if len(sig_hex) != MLDSA_SIG_HEX:
            self.log(f"DENIED: Signature must be {MLDSA_SIG_HEX} hex chars (ML-DSA-44)")
            self.log(f"        Got {len(sig_hex)} chars")
            return None, None
        try:
            bytes.fromhex(sig_hex)
        except ValueError:
            self.log("DENIED: Signature contains non-hex characters")
            return None, None
        self.log("  OK: Token parsed")
        return payload, sig_hex

    def _validate_capability(self, payload: CapabilityPayload) -> int:
        try:
            exp = datetime.fromisoformat(payload.expires_at.replace("Z", "+00:00"))
            if exp < datetime.now(timezone.utc):
                self.log(f"DENIED: Capability expired at {payload.expires_at}")
                return EXIT_DENIED
        except ValueError:
            self.log("DENIED: Invalid expires_at format")
            return EXIT_DENIED
        self.log("  OK: Capability not expired")
        return EXIT_AUTHORIZED

    # ── Step 5: ML-DSA-44 signature verification ──────────────────────────────

    def _verify_mldsa_signature(self, payload: CapabilityPayload, sig_hex: str) -> int:
        self.log("[5/5] Verifying ML-DSA-44 signature (NIST FIPS 204)...")

        # Load authority public key
        pk_hex = os.environ.get("PAX_MLDSA_PUBLIC_KEY_HEX")
        if not pk_hex:
            pk_file = self.sovereign_dir / "mldsa_authority.pub"
            if pk_file.exists():
                pk_hex = pk_file.read_text().strip()
        if not pk_hex:
            self.log("DENIED: No ML-DSA-44 authority public key found")
            self.log("        Set PAX_MLDSA_PUBLIC_KEY_HEX or create sovereign/mldsa_authority.pub")
            return EXIT_DENIED

        if len(pk_hex) != MLDSA_PK_HEX:
            self.log(f"DENIED: Authority public key must be {MLDSA_PK_HEX} hex chars (ML-DSA-44)")
            return EXIT_DENIED

        try:
            pk_bytes  = bytes.fromhex(pk_hex)
            sig_bytes = bytes.fromhex(sig_hex)
        except ValueError:
            self.log("DENIED: Key or signature contains non-hex characters")
            return EXIT_DENIED

        # Message: canonical JSON of payload (sorted keys, no whitespace)
        message = json.dumps(
            {k: getattr(payload, k) for k in sorted(payload.model_fields)},
            separators=(",", ":"), sort_keys=True
        ).encode("utf-8")

        if mldsa_verify(pk_bytes, message, sig_bytes):
            self.log("  OK: ML-DSA-44 signature VALID")
            return EXIT_AUTHORIZED
        else:
            self.log("DENIED: ML-DSA-44 signature INVALID")
            return EXIT_DENIED

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _get_git_commit(self) -> Optional[str]:
        try:
            r = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                capture_output=True, text=True,
                cwd=self.repo_root, timeout=10
            )
            return r.stdout.strip() if r.returncode == 0 else None
        except Exception:
            return None

    @staticmethod
    def _sha256_file(path: Path) -> str:
        h = hashlib.sha256()
        h.update(path.read_bytes())
        return h.hexdigest()


# ── CLI ───────────────────────────────────────────────────────────────────────

def main() -> int:
    import argparse
    parser = argparse.ArgumentParser(description="PAX-Coder PQ Gate (ML-DSA-44)")
    parser.add_argument("--quiet", action="store_true")
    parser.add_argument("--repo-root", type=Path, default=None)
    args = parser.parse_args()

    gate = PaxCoderGatePQ(repo_root=args.repo_root, quiet=args.quiet)
    return gate.run()


if __name__ == "__main__":
    sys.exit(main())