libsoliton/soliton_py/src/identity.rs

//! Identity key management: keygen, sign, verify, fingerprint.

use crate::errors::to_py_err;
use pyo3::prelude::*;
use pyo3::types::PyBytes;

/// An identity keypair (X-Wing + Ed25519 + ML-DSA-65).
///
/// Secret key material is zeroized when this object is garbage collected.
/// For deterministic cleanup, call `close()` explicitly or use as a context
/// manager: `with Identity.generate() as id: ...`
#[pyclass]
pub struct Identity {
    pk: soliton::identity::IdentityPublicKey,
    sk: Option<soliton::identity::IdentitySecretKey>,
}

#[pymethods]
impl Identity {
    /// Generate a fresh identity keypair.
    #[staticmethod]
    fn generate() -> PyResult<Self> {
        let id = soliton::identity::generate_identity().map_err(to_py_err)?;
        Ok(Self {
            pk: id.public_key,
            sk: Some(id.secret_key),
        })
    }

    /// Reconstruct from serialized public + secret key bytes.
    #[staticmethod]
    fn from_bytes(pk_bytes: &[u8], sk_bytes: &[u8]) -> PyResult<Self> {
        let pk = soliton::identity::IdentityPublicKey::from_bytes(pk_bytes.to_vec())
            .map_err(to_py_err)?;
        let sk = soliton::identity::IdentitySecretKey::from_bytes(sk_bytes.to_vec())
            .map_err(to_py_err)?;
        Ok(Self { pk, sk: Some(sk) })
    }

    /// Reconstruct a public-key-only identity (cannot sign).
    #[staticmethod]
    fn from_public_bytes(pk_bytes: &[u8]) -> PyResult<Self> {
        let pk = soliton::identity::IdentityPublicKey::from_bytes(pk_bytes.to_vec())
            .map_err(to_py_err)?;
        Ok(Self { pk, sk: None })
    }

    /// Public key bytes.
    fn public_key<'py>(&self, py: Python<'py>) -> Py<PyBytes> {
        PyBytes::new(py, self.pk.as_bytes()).into()
    }

    /// Secret key bytes. Raises if this is a public-key-only identity.
    fn secret_key<'py>(&self, py: Python<'py>) -> PyResult<Py<PyBytes>> {
        let sk = self
            .sk
            .as_ref()
            .ok_or_else(|| crate::errors::InvalidDataError::new_err("no secret key"))?;
        Ok(PyBytes::new(py, sk.as_bytes()).into())
    }

    /// SHA3-256 fingerprint of the public key (32 bytes).
    fn fingerprint<'py>(&self, py: Python<'py>) -> Py<PyBytes> {
        let fp = soliton::primitives::sha3_256::hash(self.pk.as_bytes());
        PyBytes::new(py, &fp).into()
    }

    /// Hex-encoded fingerprint.
    fn fingerprint_hex(&self) -> String {
        soliton::primitives::sha3_256::fingerprint_hex(self.pk.as_bytes())
    }

    /// Hybrid sign (Ed25519 + ML-DSA-65).
    fn sign<'py>(&self, py: Python<'py>, message: &[u8]) -> PyResult<Py<PyBytes>> {
        let sk = self
            .sk
            .as_ref()
            .ok_or_else(|| crate::errors::InvalidDataError::new_err("no secret key"))?;
        let sig = soliton::identity::hybrid_sign(sk, message).map_err(to_py_err)?;
        Ok(PyBytes::new(py, sig.as_bytes()).into())
    }

    /// Verify a hybrid signature against this identity's public key.
    fn verify(&self, message: &[u8], signature: &[u8]) -> PyResult<()> {
        let sig = soliton::identity::HybridSignature::from_bytes(signature.to_vec())
            .map_err(to_py_err)?;
        soliton::identity::hybrid_verify(&self.pk, message, &sig).map_err(to_py_err)
    }

    /// Zeroize the secret key immediately.
    fn close(&mut self) {
        self.sk = None;
    }

    fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    fn __exit__(
        &mut self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) {
        self.close();
    }
}

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Identity>()?;
    Ok(())
}