ZKAC/src/python.rs

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use rand::rngs::OsRng;

use crate::credential::{self, bbs, SIGNATURE_LEN};
use crate::credential::registry as reg;
use crate::issuance;
use crate::transport::handshake::HANDSHAKE_MSG_LEN;

fn to_py_err(e: crate::Error) -> PyErr {
    PyValueError::new_err(e.to_string())
}

fn to_32(bytes: &[u8], name: &str) -> PyResult<[u8; 32]> {
    bytes.try_into().map_err(|_| PyValueError::new_err(format!("{name} must be 32 bytes")))
}

fn bytes_to_hex(b: &[u8]) -> String {
    b.iter().map(|byte| format!("{byte:02x}")).collect()
}

// ── Ristretto Keypair (transport identity) ───────────────────────────

#[pyclass(name = "Keypair")]
pub struct PyKeypair {
    inner: Option<credential::Keypair>,
}

#[pymethods]
impl PyKeypair {
    #[new]
    fn new() -> Self {
        PyKeypair {
            inner: Some(credential::Keypair::generate(&mut OsRng)),
        }
    }

    fn public_key(&self) -> PyResult<PyPublicKey> {
        let kp = self.inner.as_ref().ok_or_else(|| {
            PyValueError::new_err("keypair was consumed by Node")
        })?;
        Ok(PyPublicKey {
            inner: *kp.public(),
        })
    }

    fn sign<'py>(&self, py: Python<'py>, msg: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let kp = self.inner.as_ref().ok_or_else(|| {
            PyValueError::new_err("keypair was consumed by Node")
        })?;
        let sig = kp.sign(msg);
        Ok(PyBytes::new(py, &sig.to_bytes()))
    }

    fn secret_key_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let kp = self.inner.as_ref().ok_or_else(|| {
            PyValueError::new_err("keypair was consumed by Node")
        })?;
        Ok(PyBytes::new(py, &kp.secret_key_bytes()))
    }

    #[staticmethod]
    fn from_secret_key(bytes: &[u8]) -> PyResult<Self> {
        if bytes.len() != 32 {
            return Err(PyValueError::new_err("secret key must be 32 bytes"));
        }
        let arr: [u8; 32] = bytes.try_into().unwrap();
        let inner = credential::Keypair::from_secret_key_bytes(&arr).map_err(to_py_err)?;
        Ok(PyKeypair {
            inner: Some(inner),
        })
    }
}

// ── Ristretto PublicKey ──────────────────────────────────────────────

#[pyclass(name = "PublicKey")]
#[derive(Clone)]
pub struct PyPublicKey {
    inner: credential::PublicKey,
}

#[pymethods]
impl PyPublicKey {
    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner.to_bytes())
    }

    #[staticmethod]
    fn from_bytes(bytes: &[u8]) -> PyResult<Self> {
        if bytes.len() != 32 {
            return Err(PyValueError::new_err("public key must be 32 bytes"));
        }
        let arr: [u8; 32] = bytes.try_into().unwrap();
        Ok(PyPublicKey {
            inner: credential::PublicKey::from_bytes(arr).map_err(to_py_err)?,
        })
    }

    fn __eq__(&self, other: &PyPublicKey) -> bool {
        self.inner == other.inner
    }

    fn __hash__(&self) -> u64 {
        let b = self.inner.to_bytes();
        u64::from_le_bytes(b[..8].try_into().unwrap())
    }

    fn verify(&self, msg: &[u8], signature: &[u8]) -> PyResult<bool> {
        if signature.len() != SIGNATURE_LEN {
            return Err(PyValueError::new_err("signature must be 64 bytes"));
        }
        let sig_arr: [u8; SIGNATURE_LEN] = signature.try_into().unwrap();
        let sig = credential::Signature::from_bytes(&sig_arr).map_err(to_py_err)?;
        match self.inner.verify(msg, &sig) {
            Ok(()) => Ok(true),
            Err(_) => Ok(false),
        }
    }

    fn __repr__(&self) -> String {
        let hex = bytes_to_hex(&self.inner.to_bytes());
        format!("PublicKey({hex})")
    }
}

// ── BBS+ Issuer ──────────────────────────────────────────────────────

#[pyclass(name = "BbsIssuer")]
pub struct PyBbsIssuer {
    inner: bbs::IssuerKeyPair,
}

#[pymethods]
impl PyBbsIssuer {
    #[new]
    fn new() -> PyResult<Self> {
        let inner = bbs::IssuerKeyPair::generate().map_err(to_py_err)?;
        Ok(Self { inner })
    }

    #[staticmethod]
    fn from_secret_key(bytes: &[u8]) -> PyResult<Self> {
        let inner = bbs::IssuerKeyPair::from_secret_key_bytes(bytes).map_err(to_py_err)?;
        Ok(Self { inner })
    }

    fn public_key(&self) -> PyBbsPublicKey {
        PyBbsPublicKey { inner: self.inner.public_key() }
    }

    fn secret_key_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner.secret_key_bytes())
    }

    fn issue_blind<'py>(
        &self,
        py: Python<'py>,
        commitment_with_proof: &[u8],
        role_id: &[u8],
        epoch: u64,
    ) -> PyResult<Bound<'py, PyBytes>> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);
        let sig = self.inner.issue_blind(commitment_with_proof, &rid, epoch).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &sig))
    }
}

// ── BBS+ Public Key ──────────────────────────────────────────────────

#[pyclass(name = "BbsPublicKey")]
#[derive(Clone)]
pub struct PyBbsPublicKey {
    inner: bbs::IssuerPublicKey,
}

#[pymethods]
impl PyBbsPublicKey {
    fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner.to_bytes())
    }

    #[staticmethod]
    fn from_bytes(bytes: &[u8]) -> PyResult<Self> {
        let inner = bbs::IssuerPublicKey::from_bytes(bytes).map_err(to_py_err)?;
        Ok(Self { inner })
    }
}

// ── BBS+ Blind Request ──────────────────────────────────────────────

#[pyclass(name = "BlindRequest")]
pub struct PyBlindRequest {
    commitment_with_proof: Vec<u8>,
    prover_blind_bytes: Vec<u8>,
    member_secret: Vec<u8>,
}

#[pymethods]
impl PyBlindRequest {
    fn commitment_with_proof<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.commitment_with_proof)
    }

    fn prover_blind<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.prover_blind_bytes)
    }

    fn member_secret<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.member_secret)
    }
}

#[pyfunction]
fn prepare_blind_request() -> PyResult<PyBlindRequest> {
    let req = bbs::prepare_blind_request().map_err(to_py_err)?;
    Ok(PyBlindRequest {
        commitment_with_proof: req.commitment_with_proof,
        prover_blind_bytes: req.prover_blind.to_bytes().to_vec(),
        member_secret: req.member_secret,
    })
}

// ── BBS+ Credential ─────────────────────────────────────────────────

#[pyclass(name = "Credential")]
pub struct PyCredential {
    inner: bbs::Credential,
}

#[pymethods]
impl PyCredential {
    #[staticmethod]
    fn finalize(
        blind_sig_bytes: &[u8],
        member_secret: &[u8],
        prover_blind: &[u8],
        role_id: &[u8],
        epoch: u64,
        pk: &PyBbsPublicKey,
    ) -> PyResult<Self> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        if prover_blind.len() != 32 {
            return Err(PyValueError::new_err("prover_blind must be 32 bytes"));
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);

        let pb_arr: [u8; 32] = prover_blind.try_into().unwrap();
        let blind_factor = zkryptium::bbsplus::commitment::BlindFactor::from_bytes(&pb_arr)
            .map_err(|e| PyValueError::new_err(e.to_string()))?;

        let inner = bbs::Credential::finalize(
            blind_sig_bytes,
            member_secret.to_vec(),
            blind_factor,
            rid,
            epoch,
            &pk.inner,
        ).map_err(to_py_err)?;

        Ok(Self { inner })
    }

    fn present<'py>(&self, py: Python<'py>, nonce: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let pres = self.inner.present(nonce).map_err(to_py_err)?;
        Ok(PyBytes::new(py, pres.to_bytes()))
    }

    fn role_id<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, self.inner.role_id())
    }

    fn epoch(&self) -> u64 {
        self.inner.epoch()
    }
}

// ── BBS+ Role ID helper ─────────────────────────────────────────────

#[pyfunction]
fn role_id<'py>(py: Python<'py>, name: &str) -> Bound<'py, PyBytes> {
    let rid = bbs::role_id(name);
    PyBytes::new(py, &rid)
}

// ── RoleRegistry (BBS+ server-side) ──────────────────────────────────

#[pyclass(name = "RoleRegistry")]
pub struct PyRoleRegistry {
    inner: credential::RoleRegistry,
}

#[pymethods]
impl PyRoleRegistry {
    #[new]
    fn new() -> Self {
        PyRoleRegistry {
            inner: credential::RoleRegistry::new(),
        }
    }

    fn register_role(&mut self, role_id: &[u8], pk: &PyBbsPublicKey, epoch: u64) -> PyResult<()> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);
        self.inner.register_role(rid, pk.inner.clone(), epoch);
        Ok(())
    }

    fn set_epoch(&mut self, role_id: &[u8], epoch: u64) -> PyResult<()> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);
        self.inner.set_epoch(&rid, epoch).map_err(to_py_err)
    }

    fn verify_presentation(&self, role_id: &[u8], proof_bytes: &[u8], nonce: &[u8]) -> PyResult<bool> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        if proof_bytes.len() > crate::node::MAX_BBS_AUTH_PROOF_BYTES {
            return Ok(false);
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);
        let pres = bbs::Presentation::from_bytes(proof_bytes.to_vec());
        match self.inner.verify_presentation(&rid, &pres, nonce) {
            Ok(()) => Ok(true),
            Err(crate::Error::InvalidPresentation | crate::Error::RoleNotRegistered) => Ok(false),
            Err(e) => Err(to_py_err(e)),
        }
    }

    fn has_role(&self, role_id: &[u8]) -> PyResult<bool> {
        if role_id.len() != 32 {
            return Err(PyValueError::new_err("role_id must be 32 bytes"));
        }
        let mut rid = [0u8; 32];
        rid.copy_from_slice(role_id);
        Ok(self.inner.has_role(&rid))
    }
}

// ── Registry State (client-managed) ──────────────────────────────────

#[pyclass(name = "RegistryState")]
pub struct PyRegistryState {
    inner_bytes: Vec<u8>,
}

#[pymethods]
impl PyRegistryState {
    #[staticmethod]
    fn build(
        admin_issuer_pk: &PyBbsPublicKey,
        issuance_pk: &[u8],
        version: u64,
        prev_state_hash: &[u8],
        roles: Vec<(Vec<u8>, PyBbsPublicKey, u64)>,
    ) -> PyResult<Self> {
        if issuance_pk.len() != 32 {
            return Err(PyValueError::new_err("issuance_pk must be 32 bytes"));
        }
        if prev_state_hash.len() != 32 {
            return Err(PyValueError::new_err("prev_state_hash must be 32 bytes"));
        }
        let mut iss_pk = [0u8; 32];
        iss_pk.copy_from_slice(issuance_pk);
        let mut prev = [0u8; 32];
        prev.copy_from_slice(prev_state_hash);

        let entries: PyResult<Vec<reg::RoleEntry>> = roles.into_iter().map(|(rid, pk, epoch)| {
            if rid.len() != 32 {
                return Err(PyValueError::new_err("role_id must be 32 bytes"));
            }
            let mut role_id = [0u8; 32];
            role_id.copy_from_slice(&rid);
            Ok(reg::RoleEntry { role_id, issuer_pk: pk.inner.clone(), epoch })
        }).collect();

        let state = reg::RegistryState::new(
            admin_issuer_pk.inner.clone(), iss_pk, version, prev, entries?,
        );
        Ok(PyRegistryState { inner_bytes: state.serialize() })
    }

    fn serialize<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner_bytes)
    }

    #[staticmethod]
    fn deserialize(data: &[u8]) -> PyResult<Self> {
        let _ = reg::RegistryState::deserialize(data).map_err(to_py_err)?;
        Ok(PyRegistryState { inner_bytes: data.to_vec() })
    }

    fn registry_id<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let state = reg::RegistryState::deserialize(&self.inner_bytes).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &state.registry_id))
    }

    fn version(&self) -> PyResult<u64> {
        let state = reg::RegistryState::deserialize(&self.inner_bytes).map_err(to_py_err)?;
        Ok(state.version)
    }

    fn state_hash<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        let h = reg::RegistryState::state_hash(&self.inner_bytes);
        PyBytes::new(py, &h)
    }

    fn certify<'py>(&self, py: Python<'py>, admin_credential: &PyCredential) -> PyResult<Bound<'py, PyBytes>> {
        let state = reg::RegistryState::deserialize(&self.inner_bytes).map_err(to_py_err)?;
        let cert = state.certify(&admin_credential.inner, &self.inner_bytes).map_err(to_py_err)?;
        Ok(PyBytes::new(py, cert.to_bytes()))
    }

    #[staticmethod]
    fn verify_cert(admin_issuer_pk: &PyBbsPublicKey, state_cert: &[u8], state_bytes: &[u8]) -> PyResult<bool> {
        let cert = bbs::Presentation::from_bytes(state_cert.to_vec());
        match reg::RegistryState::verify_cert(&admin_issuer_pk.inner, &cert, state_bytes) {
            Ok(()) => Ok(true),
            Err(_) => Ok(false),
        }
    }
}

#[pyfunction]
fn registry_id<'py>(py: Python<'py>, admin_issuer_pk: &PyBbsPublicKey) -> Bound<'py, PyBytes> {
    let rid = reg::registry_id(&admin_issuer_pk.inner);
    PyBytes::new(py, &rid)
}

#[pyfunction]
fn admin_role_id<'py>(py: Python<'py>) -> Bound<'py, PyBytes> {
    let rid = reg::admin_role_id();
    PyBytes::new(py, &rid)
}

// ── Registry Manager (server-side) ──────────────────────────────────

#[pyclass(name = "RegistryManager")]
pub struct PyRegistryManager {
    inner: crate::registry_manager::RegistryManager,
}

#[pymethods]
impl PyRegistryManager {
    #[new]
    fn new() -> Self {
        PyRegistryManager {
            inner: crate::registry_manager::RegistryManager::new(),
        }
    }

    fn create<'py>(&mut self, py: Python<'py>, state_bytes: &[u8], state_cert: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let rid = self.inner.create(state_bytes, state_cert).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &rid))
    }

    /// Load certified registry state from disk (any version); see [`RegistryManager::restore`].
    fn restore<'py>(&mut self, py: Python<'py>, state_bytes: &[u8], state_cert: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let rid = self.inner.restore(state_bytes, state_cert).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &rid))
    }

    fn update(&mut self, registry_id: &[u8], state_bytes: &[u8], state_cert: &[u8]) -> PyResult<()> {
        let rid = to_32(registry_id, "registry_id")?;
        self.inner.update(&rid, state_bytes, state_cert).map_err(to_py_err)
    }

    fn get<'py>(&self, py: Python<'py>, registry_id: &[u8]) -> PyResult<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> {
        let rid = to_32(registry_id, "registry_id")?;
        let (state_bytes, cert_bytes) = self.inner.get(&rid).map_err(to_py_err)?;
        Ok((PyBytes::new(py, state_bytes), PyBytes::new(py, cert_bytes)))
    }

    fn has_registry(&self, registry_id: &[u8]) -> PyResult<bool> {
        let rid = to_32(registry_id, "registry_id")?;
        Ok(self.inner.has_registry(&rid))
    }

    fn verify_admin(&self, registry_id: &[u8], proof_bytes: &[u8], nonce: &[u8]) -> PyResult<bool> {
        if proof_bytes.len() > crate::node::MAX_BBS_AUTH_PROOF_BYTES {
            return Ok(false);
        }
        let rid = to_32(registry_id, "registry_id")?;
        match self.inner.verify_admin(&rid, proof_bytes, nonce) {
            Ok(()) => Ok(true),
            Err(crate::Error::InvalidPresentation | crate::Error::RegistryNotFound) => Ok(false),
            Err(e) => Err(to_py_err(e)),
        }
    }

    fn verify_presentation(&self, registry_id: &[u8], role_id: &[u8], proof_bytes: &[u8], nonce: &[u8]) -> PyResult<bool> {
        if proof_bytes.len() > crate::node::MAX_BBS_AUTH_PROOF_BYTES {
            return Ok(false);
        }
        let rid = to_32(registry_id, "registry_id")?;
        let roid = to_32(role_id, "role_id")?;
        let pres = bbs::Presentation::from_bytes(proof_bytes.to_vec());
        match self.inner.verify_presentation(&rid, &roid, &pres, nonce) {
            Ok(()) => Ok(true),
            Err(crate::Error::InvalidPresentation | crate::Error::RoleNotRegistered | crate::Error::RegistryNotFound) => Ok(false),
            Err(e) => Err(to_py_err(e)),
        }
    }

    fn queue_issuance_request(
        &mut self,
        registry_id: &[u8],
        request_id: &[u8],
        role_id: &[u8],
        eph_pk: &[u8],
        encrypted_commitment: &[u8],
    ) -> PyResult<()> {
        let rid = to_32(registry_id, "registry_id")?;
        let req_id = to_32(request_id, "request_id")?;
        let roid = to_32(role_id, "role_id")?;
        let epk = to_32(eph_pk, "eph_pk")?;
        self.inner.queue_issuance_request(&rid, req_id, roid, epk, encrypted_commitment.to_vec())
            .map_err(to_py_err)
    }

    fn take_pending_requests(&mut self, py: Python<'_>, registry_id: &[u8]) -> PyResult<Vec<PyObject>> {
        let rid = to_32(registry_id, "registry_id")?;
        let items = self.inner.take_pending_requests(&rid).map_err(to_py_err)?;
        let result: Vec<PyObject> = items.into_iter().map(|(req_id, req)| {
            let tuple = (
                PyBytes::new(py, &req_id).into_any(),
                PyBytes::new(py, &req.role_id).into_any(),
                PyBytes::new(py, &req.eph_pk).into_any(),
                PyBytes::new(py, &req.encrypted_commitment).into_any(),
            );
            tuple.into_pyobject(py).unwrap().into_any().unbind()
        }).collect();
        Ok(result)
    }

    fn grant_credential(&mut self, registry_id: &[u8], request_id: &[u8], encrypted_blind_sig: &[u8]) -> PyResult<()> {
        let rid = to_32(registry_id, "registry_id")?;
        let req_id = to_32(request_id, "request_id")?;
        self.inner.grant_credential(&rid, req_id, encrypted_blind_sig.to_vec())
            .map_err(to_py_err)
    }

    fn take_granted_credential<'py>(&mut self, py: Python<'py>, registry_id: &[u8], request_id: &[u8]) -> PyResult<Option<Bound<'py, PyBytes>>> {
        let rid = to_32(registry_id, "registry_id")?;
        let req_id = to_32(request_id, "request_id")?;
        let result = self.inner.take_granted_credential(&rid, &req_id).map_err(to_py_err)?;
        Ok(result.map(|data| PyBytes::new(py, &data)))
    }
}

// ── Issuance E2E Encryption ──────────────────────────────────────────

#[pyclass(name = "IssuanceKeypair")]
pub struct PyIssuanceKeypair {
    inner: issuance::IssuanceKeypair,
}

#[pymethods]
impl PyIssuanceKeypair {
    #[new]
    fn new() -> Self {
        PyIssuanceKeypair {
            inner: issuance::IssuanceKeypair::generate(&mut OsRng),
        }
    }

    #[staticmethod]
    fn from_secret(bytes: &[u8]) -> PyResult<Self> {
        if bytes.len() != 32 {
            return Err(PyValueError::new_err("issuance secret key must be 32 bytes"));
        }
        let arr: [u8; 32] = bytes.try_into().unwrap();
        Ok(PyIssuanceKeypair {
            inner: issuance::IssuanceKeypair::from_secret_bytes(&arr),
        })
    }

    fn public_key_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner.public_key_bytes())
    }

    fn secret_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inner.secret_bytes())
    }

    fn decrypt<'py>(&self, py: Python<'py>, eph_pk: &[u8], ciphertext: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let epk = to_32(eph_pk, "eph_pk")?;
        let plaintext = self.inner.decrypt(&epk, ciphertext).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &plaintext))
    }

    fn encrypt<'py>(&self, py: Python<'py>, eph_pk: &[u8], plaintext: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let epk = to_32(eph_pk, "eph_pk")?;
        let ciphertext = self.inner.encrypt(&epk, plaintext).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &ciphertext))
    }
}

#[pyfunction]
fn encrypt_for_admin<'py>(py: Python<'py>, admin_issuance_pk: &[u8], plaintext: &[u8]) -> PyResult<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> {
    let pk = to_32(admin_issuance_pk, "admin_issuance_pk")?;
    let (eph_pk, ciphertext) = issuance::encrypt_for_admin(OsRng, &pk, plaintext).map_err(to_py_err)?;
    Ok((PyBytes::new(py, &eph_pk), PyBytes::new(py, &ciphertext)))
}

#[pyfunction]
fn decrypt_from_admin<'py>(py: Python<'py>, eph_secret: &[u8], admin_issuance_pk: &[u8], ciphertext: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
    let sec = to_32(eph_secret, "eph_secret")?;
    let pk = to_32(admin_issuance_pk, "admin_issuance_pk")?;
    let plaintext = issuance::decrypt_from_admin(&sec, &pk, ciphertext).map_err(to_py_err)?;
    Ok(PyBytes::new(py, &plaintext))
}

// ── Session ──────────────────────────────────────────────────────────

#[pyclass(name = "Session")]
pub struct PySession {
    inner: crate::transport::Session,
}

#[pymethods]
impl PySession {
    fn transcript_hash<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, self.inner.transcript_hash())
    }

    fn encrypt<'py>(&mut self, py: Python<'py>, plaintext: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let packet = self.inner.encrypt(plaintext).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &packet))
    }

    fn decrypt<'py>(&mut self, py: Python<'py>, packet: &[u8]) -> PyResult<Bound<'py, PyBytes>> {
        let plaintext = self.inner.decrypt(packet).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &plaintext))
    }
}

// ── Node (high-level API) ────────────────────────────────────────────

#[pyclass(name = "Node")]
pub struct PyNode {
    inner: crate::node::Node,
}

#[pyclass(name = "PendingConnect")]
pub struct PyPendingConnect {
    inner: Option<crate::node::PendingConnect>,
}

#[pymethods]
impl PyNode {
    #[new]
    fn new(keypair: &mut PyKeypair) -> PyResult<Self> {
        let kp = keypair.inner.take().ok_or_else(|| {
            PyValueError::new_err("keypair already consumed")
        })?;
        Ok(PyNode {
            inner: crate::node::Node::new(kp),
        })
    }

    fn public_key(&self) -> PyPublicKey {
        PyPublicKey {
            inner: *self.inner.public_key(),
        }
    }

    fn connect(&self) -> (PyPendingConnect, Vec<u8>) {
        let (pending, msg) = self.inner.connect(OsRng);
        (
            PyPendingConnect {
                inner: Some(pending),
            },
            msg.to_vec(),
        )
    }

    /// Complete handshake: verify server identity, then produce BBS+ auth.
    fn complete_connect(
        &self,
        pending: &mut PyPendingConnect,
        response_msg: &[u8],
        identity_proof: &[u8],
        expected_server_pk: &PyPublicKey,
        credential: &PyCredential,
    ) -> PyResult<(PySession, Vec<u8>)> {
        let p = pending
            .inner
            .take()
            .ok_or_else(|| PyValueError::new_err("PendingConnect already consumed"))?;

        if response_msg.len() != HANDSHAKE_MSG_LEN {
            return Err(PyValueError::new_err("response_msg must be 32 bytes"));
        }
        let msg: [u8; HANDSHAKE_MSG_LEN] = response_msg.try_into().unwrap();

        let (session, auth_packet) = self
            .inner
            .complete_connect(
                p,
                &msg,
                identity_proof,
                &expected_server_pk.inner,
                &credential.inner,
            )
            .map_err(to_py_err)?;

        Ok((PySession { inner: session }, auth_packet))
    }

    /// Produce encrypted identity proof (server signs transcript with long-term key).
    fn prove_identity<'py>(
        &self,
        py: Python<'py>,
        session: &mut PySession,
    ) -> PyResult<Bound<'py, PyBytes>> {
        let proof = self.inner.prove_identity(&mut session.inner).map_err(to_py_err)?;
        Ok(PyBytes::new(py, &proof))
    }

    fn accept(&self, init_msg: &[u8]) -> PyResult<(PySession, Vec<u8>)> {
        if init_msg.len() != HANDSHAKE_MSG_LEN {
            return Err(PyValueError::new_err("init_msg must be 32 bytes"));
        }
        let msg: [u8; HANDSHAKE_MSG_LEN] = init_msg.try_into().unwrap();
        let (session, response) = self.inner.accept(OsRng, &msg).map_err(to_py_err)?;
        Ok((PySession { inner: session }, response.to_vec()))
    }

    /// Verify encrypted BBS+ auth packet. Returns the role_id (32 bytes) on success.
    fn verify_auth<'py>(
        &self,
        py: Python<'py>,
        session: &mut PySession,
        encrypted_auth: &[u8],
        registry: &PyRoleRegistry,
    ) -> PyResult<Bound<'py, PyBytes>> {
        let rid = self
            .inner
            .verify_auth(&mut session.inner, encrypted_auth, &registry.inner)
            .map_err(to_py_err)?;
        Ok(PyBytes::new(py, &rid))
    }

    /// Complete handshake verifying server identity only (no BBS+ auth).
    fn complete_connect_anon(
        &self,
        pending: &mut PyPendingConnect,
        response_msg: &[u8],
        identity_proof: &[u8],
        expected_server_pk: &PyPublicKey,
    ) -> PyResult<PySession> {
        let p = pending
            .inner
            .take()
            .ok_or_else(|| PyValueError::new_err("PendingConnect already consumed"))?;

        if response_msg.len() != HANDSHAKE_MSG_LEN {
            return Err(PyValueError::new_err("response_msg must be 32 bytes"));
        }
        let msg: [u8; HANDSHAKE_MSG_LEN] = response_msg.try_into().unwrap();

        let session = self
            .inner
            .complete_connect_anon(
                p,
                &msg,
                identity_proof,
                &expected_server_pk.inner,
            )
            .map_err(to_py_err)?;

        Ok(PySession { inner: session })
    }

    /// Complete handshake for a client-managed registry. Includes
    /// registry_id in the auth packet.
    fn complete_connect_managed(
        &self,
        pending: &mut PyPendingConnect,
        response_msg: &[u8],
        identity_proof: &[u8],
        expected_server_pk: &PyPublicKey,
        credential: &PyCredential,
        registry_id: &[u8],
    ) -> PyResult<(PySession, Vec<u8>)> {
        let p = pending
            .inner
            .take()
            .ok_or_else(|| PyValueError::new_err("PendingConnect already consumed"))?;

        if response_msg.len() != HANDSHAKE_MSG_LEN {
            return Err(PyValueError::new_err("response_msg must be 32 bytes"));
        }
        let msg: [u8; HANDSHAKE_MSG_LEN] = response_msg.try_into().unwrap();
        let rid = to_32(registry_id, "registry_id")?;

        let (session, auth_packet) = self
            .inner
            .complete_connect_managed(
                p,
                &msg,
                identity_proof,
                &expected_server_pk.inner,
                &credential.inner,
                &rid,
            )
            .map_err(to_py_err)?;

        Ok((PySession { inner: session }, auth_packet))
    }

    /// Verify managed-registry auth packet. Returns (registry_id, role_id).
    fn verify_auth_managed<'py>(
        &self,
        py: Python<'py>,
        session: &mut PySession,
        encrypted_auth: &[u8],
        manager: &PyRegistryManager,
    ) -> PyResult<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> {
        let (reg_id, role_id) = self
            .inner
            .verify_auth_managed(&mut session.inner, encrypted_auth, &manager.inner)
            .map_err(to_py_err)?;
        Ok((PyBytes::new(py, &reg_id), PyBytes::new(py, &role_id)))
    }
}

// ── Module ───────────────────────────────────────────────────────────

#[pymodule]
fn _zkac(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("MAX_BBS_AUTH_PROOF_BYTES", crate::node::MAX_BBS_AUTH_PROOF_BYTES)?;
    // Transport identity (ristretto255)
    m.add_class::<PyKeypair>()?;
    m.add_class::<PyPublicKey>()?;
    // BBS+ anonymous credentials
    m.add_class::<PyBbsIssuer>()?;
    m.add_class::<PyBbsPublicKey>()?;
    m.add_class::<PyBlindRequest>()?;
    m.add_class::<PyCredential>()?;
    m.add_function(wrap_pyfunction!(prepare_blind_request, m)?)?;
    m.add_function(wrap_pyfunction!(role_id, m)?)?;
    // Server registry (static, server-configured)
    m.add_class::<PyRoleRegistry>()?;
    // Client-managed registries
    m.add_class::<PyRegistryState>()?;
    m.add_class::<PyRegistryManager>()?;
    m.add_function(wrap_pyfunction!(registry_id, m)?)?;
    m.add_function(wrap_pyfunction!(admin_role_id, m)?)?;
    // E2E-encrypted issuance
    m.add_class::<PyIssuanceKeypair>()?;
    m.add_function(wrap_pyfunction!(encrypt_for_admin, m)?)?;
    m.add_function(wrap_pyfunction!(decrypt_from_admin, m)?)?;
    // Transport
    m.add_class::<PySession>()?;
    m.add_class::<PyNode>()?;
    m.add_class::<PyPendingConnect>()?;
    Ok(())
}