synta-python 0.2.4

//! Python bindings for PKCS#11 token management.
//!
//! Exposes the `synta.pkcs11` submodule with:
//! - `SlotInfo` — information about a PKCS#11 slot / token
//! - `KeyInfo`  — information about a private key stored on a token
//! - `Pkcs11Token` — session handle for a single named token
//! - `list_slots(module=None)` — enumerate all token slots

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

use synta_certificate::crypto::{KeySpec, TokenManager};
use synta_certificate::pkcs11_mgmt::Pkcs11Manager;
use synta_certificate::{merge_object_label, Pkcs11Uri, Pkcs11UriAttributes};

use crate::crypto_keys::PyPrivateKey;

// ── SlotInfo ──────────────────────────────────────────────────────────────────

/// Information about a PKCS#11 token slot.
///
/// Returned by :func:`list_slots`.
#[pyclass(frozen, name = "SlotInfo")]
pub struct PySlotInfo {
    pub(crate) slot_id: u64,
    pub(crate) token_label: String,
    pub(crate) manufacturer_id: String,
    pub(crate) model: String,
    pub(crate) serial_number: String,
    pub(crate) flags: u64,
}

#[pymethods]
impl PySlotInfo {
    /// Numeric PKCS#11 slot ID (``CK_SLOT_ID``).
    #[getter]
    fn slot_id(&self) -> u64 {
        self.slot_id
    }

    /// Token label (`CKA_LABEL` on the token object), space-padding removed.
    #[getter]
    fn token_label(&self) -> &str {
        &self.token_label
    }

    /// Manufacturer ID field from ``CK_TOKEN_INFO``.
    #[getter]
    fn manufacturer_id(&self) -> &str {
        &self.manufacturer_id
    }

    /// Model field from ``CK_TOKEN_INFO``.
    #[getter]
    fn model(&self) -> &str {
        &self.model
    }

    /// Serial number field from ``CK_TOKEN_INFO``.
    #[getter]
    fn serial_number(&self) -> &str {
        &self.serial_number
    }

    /// Subset of ``CKF_*`` flag bits from ``CK_TOKEN_INFO`` (PKCS#11 v3.0 §4.9).
    ///
    /// Four bits are captured:
    ///
    /// - ``0x0002`` (``CKF_WRITE_PROTECTED``) — token is read-only
    /// - ``0x0004`` (``CKF_LOGIN_REQUIRED``) — a PIN must be presented before accessing private objects
    /// - ``0x0100`` (``CKF_PROTECTED_AUTHENTICATION_PATH``) — PIN entry is via an on-device keypad
    /// - ``0x0400`` (``CKF_TOKEN_INITIALIZED``) — token has been initialised
    ///
    /// All other bits are zero.
    #[getter]
    fn flags(&self) -> u64 {
        self.flags
    }

    fn __repr__(&self) -> String {
        format!(
            "SlotInfo(slot_id={}, token_label={:?}, manufacturer_id={:?})",
            self.slot_id, self.token_label, self.manufacturer_id
        )
    }
}

// ── KeyInfo ───────────────────────────────────────────────────────────────────

/// Information about a private key object stored on a PKCS#11 token.
///
/// Returned by :meth:`Pkcs11Token.list_keys`.
#[pyclass(frozen, name = "KeyInfo")]
pub struct PyKeyInfo {
    pub(crate) label: String,
    pub(crate) id: Vec<u8>,
    pub(crate) key_type: String,
    pub(crate) key_bits: u32,
}

#[pymethods]
impl PyKeyInfo {
    /// Key label (``CKA_LABEL``).
    #[getter]
    fn label(&self) -> &str {
        &self.label
    }

    /// Raw ``CKA_ID`` bytes; may be empty.
    #[getter]
    fn id<'py>(&self, py: Python<'py>) -> Bound<'py, pyo3::types::PyBytes> {
        pyo3::types::PyBytes::new(py, &self.id)
    }

    /// Human-readable key type: ``"RSA"``, ``"EC"``, ``"Ed"``, ``"ML-DSA"``,
    /// ``"ML-KEM"``, or ``"Unknown"``.
    #[getter]
    fn key_type(&self) -> &str {
        &self.key_type
    }

    /// RSA modulus length in bits, derived from ``CKA_MODULUS``.
    ///
    /// ``0`` for all non-RSA key types (EC, Ed25519, Ed448, ML-DSA, ML-KEM) because
    /// those algorithms do not expose ``CKA_MODULUS``.
    #[getter]
    fn key_bits(&self) -> u32 {
        self.key_bits
    }

    fn __repr__(&self) -> String {
        format!(
            "KeyInfo(label={:?}, key_type={:?}, key_bits={})",
            self.label, self.key_type, self.key_bits
        )
    }
}

// ── Pkcs11Token ───────────────────────────────────────────────────────────────

/// A handle for PKCS#11 management operations on a single named token.
///
/// Example:
///     >>> import synta.pkcs11
///     >>> token = synta.pkcs11.Pkcs11Token(
///     ...     "pkcs11:token=MySoftHSM2Token0?pin-value=1234"
///     ... )
///     >>> token.list_keys()
///     [KeyInfo(label='caKey', key_type='RSA', key_bits=2048)]
#[pyclass(frozen, name = "Pkcs11Token")]
pub struct PyPkcs11Token {
    mgr: Pkcs11Manager,
    /// Verbatim URI string (used for repr and for building per-operation URIs).
    uri: String,
    /// Decoded attributes from `uri`, cached so methods do not need to re-parse.
    attrs: Pkcs11UriAttributes,
}

fn make_mgr(uri: &Pkcs11Uri, module: Option<&str>) -> PyResult<Pkcs11Manager> {
    match module {
        Some(path) => Pkcs11Manager::new(path).map_err(|e| PyValueError::new_err(format!("{e}"))),
        None => Pkcs11Manager::from_uri(uri).map_err(|e| PyValueError::new_err(format!("{e}"))),
    }
}

/// Redact the `pin-value=…` component from a PKCS#11 URI for safe display.
fn redact_pin(uri: &str) -> String {
    let body = uri.strip_prefix("pkcs11:").unwrap_or(uri);
    let (path, query) = body.split_once('?').unwrap_or((body, ""));
    if query.is_empty() {
        return uri.to_owned();
    }
    let redacted: Vec<&str> = query
        .split('&')
        .map(|seg| {
            if seg.starts_with("pin-value=") {
                "pin-value=***"
            } else {
                seg
            }
        })
        .collect();
    format!("pkcs11:{path}?{}", redacted.join("&"))
}

#[pymethods]
impl PyPkcs11Token {
    /// Create a PKCS#11 token handle.
    ///
    /// ``uri`` must be a ``pkcs11:`` URI identifying the token; include a
    /// ``?pin-value=<PIN>`` query component to authenticate automatically.
    ///
    /// ``module`` optionally specifies the path to the PKCS#11 shared library.
    /// When omitted, the library is resolved from the ``PKCS11_MODULE_PATH``
    /// environment variable or the system default (``p11-kit-proxy.so``).
    ///
    /// Example:
    ///     >>> token = synta.pkcs11.Pkcs11Token(
    ///     ...     "pkcs11:token=TestToken?pin-value=1234",
    ///     ...     module="/usr/lib64/pkcs11/libkryoptic_pkcs11.so",
    ///     ... )
    #[new]
    #[pyo3(signature = (uri, module = None))]
    fn new(uri: &str, module: Option<&str>) -> PyResult<Self> {
        let parsed = Pkcs11Uri::parse(uri).ok_or_else(|| {
            PyValueError::new_err(format!("invalid PKCS#11 URI: {:?}", redact_pin(uri)))
        })?;
        let mgr = make_mgr(&parsed, module)?;
        let attrs = parsed.attrs;
        Ok(Self {
            mgr,
            uri: uri.to_owned(),
            attrs,
        })
    }

    /// Return ``True`` if a private key matching ``object_uri``'s label exists
    /// on this token.
    ///
    /// ``object_uri`` must be a ``pkcs11:`` URI with an ``object=<label>``
    /// component.  The token name and PIN are inherited from the URI used to
    /// construct this :class:`Pkcs11Token` when not present in ``object_uri``.
    ///
    /// Example:
    ///     >>> token.find_key("pkcs11:object=caKey")
    ///     True
    fn find_key(&self, object_uri: &str) -> PyResult<bool> {
        let mut uri = Pkcs11Uri::parse(object_uri).ok_or_else(|| {
            PyValueError::new_err(format!("invalid PKCS#11 URI: {:?}", redact_pin(object_uri)))
        })?;
        // Inherit token and PIN from the construction URI when absent in object_uri.
        if uri.attrs.token.is_none() {
            uri.attrs.token = self.attrs.token.clone();
        }
        if !uri.attrs.has_pin() {
            uri.attrs
                .set_pin_value(self.attrs.pin_value().map(|s| s.to_owned()));
        }
        self.mgr
            .find_key(&uri)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// List all private keys on this token.
    ///
    /// Returns a list of :class:`KeyInfo` objects.
    ///
    /// Example:
    ///     >>> keys = token.list_keys()
    ///     >>> for k in keys:
    ///     ...     print(k.label, k.key_type, k.key_bits)
    fn list_keys(&self) -> PyResult<Vec<PyKeyInfo>> {
        let token_name = self
            .attrs
            .token
            .as_deref()
            .ok_or_else(|| PyValueError::new_err("Pkcs11Token URI must contain 'token='"))?;
        let pin = self.attrs.pin_value();
        let keys = self
            .mgr
            .list_keys(token_name, pin)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        Ok(keys
            .into_iter()
            .map(|k| PyKeyInfo {
                label: k.label,
                id: k.id,
                key_type: k.key_type,
                key_bits: k.key_bits,
            })
            .collect())
    }

    /// Destroy the private key (and matching public key) identified by
    /// ``object_uri`` from this token.
    ///
    /// The token name and PIN are inherited from the URI used to construct this
    /// :class:`Pkcs11Token` when not present in ``object_uri``.
    ///
    /// Example:
    ///     >>> token.delete_key("pkcs11:object=caKey")
    fn delete_key(&self, object_uri: &str) -> PyResult<()> {
        let mut uri = Pkcs11Uri::parse(object_uri).ok_or_else(|| {
            PyValueError::new_err(format!("invalid PKCS#11 URI: {:?}", redact_pin(object_uri)))
        })?;
        // Inherit token and PIN from the construction URI when absent in object_uri.
        if uri.attrs.token.is_none() {
            uri.attrs.token = self.attrs.token.clone();
        }
        if !uri.attrs.has_pin() {
            uri.attrs
                .set_pin_value(self.attrs.pin_value().map(|s| s.to_owned()));
        }
        self.mgr
            .delete_key(&uri)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// Generate a key pair on this token and return the private key.
    ///
    /// ``key_type`` is ``"rsa"``, ``"ec"``, ``"ed25519"``, ``"ed448"``,
    /// ``"ml-dsa"``, or ``"ml-kem"``.
    /// ``param`` is the RSA bit-length (``int``) for RSA keys, the curve name
    /// (``str``) for EC keys (``"P-256"``, ``"P-384"``, ``"P-521"``), or the
    /// parameter-set name for PQC keys (``"ML-DSA-44"``, ``"ML-DSA-65"``,
    /// ``"ML-DSA-87"``; ``"ML-KEM-512"``, ``"ML-KEM-768"``, ``"ML-KEM-1024"``).
    /// ``label`` is the ``CKA_LABEL`` to assign to the generated key objects.
    /// ``extractable`` controls ``CKA_EXTRACTABLE`` on the private key; default
    /// is ``False`` (hardware-resident key that cannot be exported).  Set to
    /// ``True`` only when key export is explicitly required.
    ///
    /// The token and PIN are taken from the URI used to construct this
    /// :class:`Pkcs11Token`.
    ///
    /// Example:
    ///     >>> key = token.generate_key_pair("rsa", 4096, "caKey")
    ///     >>> key = token.generate_key_pair("ec", "P-256", "ecKey")
    ///     >>> key = token.generate_key_pair("ml-dsa", "ML-DSA-65", "dsaKey")
    ///     >>> key = token.generate_key_pair("ml-kem", "ML-KEM-768", "kemKey")
    ///     >>> key = token.generate_key_pair("rsa", 2048, "backupKey", extractable=True)
    #[pyo3(signature = (key_type, param, label, extractable = false))]
    fn generate_key_pair(
        &self,
        key_type: &str,
        param: Bound<'_, PyAny>,
        label: &str,
        extractable: bool,
    ) -> PyResult<PyPrivateKey> {
        let spec = build_key_spec(key_type, &param)?;

        // Merge the label into the base URI (preserving token= and pin-value=).
        let full_uri_str = merge_object_label(&self.uri, label);
        let full_uri = Pkcs11Uri::parse(&full_uri_str).ok_or_else(|| {
            PyValueError::new_err(format!("failed to build URI with label '{label}'"))
        })?;

        let inner = self
            .mgr
            .generate_key_pair_in_token(&spec, &full_uri, extractable)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        Ok(PyPrivateKey { inner })
    }

    fn __repr__(&self) -> String {
        format!("Pkcs11Token(uri={:?})", redact_pin(&self.uri))
    }
}

// ── Module-level list_slots ───────────────────────────────────────────────────

/// List all PKCS#11 token slots using the system module.
///
/// ``module`` optionally specifies the PKCS#11 shared library path.  When
/// omitted, the library is resolved from ``PKCS11_MODULE_PATH`` or the system
/// default (``p11-kit-proxy.so``).
///
/// Example:
///     >>> import synta.pkcs11
///     >>> slots = synta.pkcs11.list_slots()
///     >>> for s in slots:
///     ...     print(s.slot_id, s.token_label, s.manufacturer_id)
#[pyfunction]
#[pyo3(signature = (module = None))]
pub fn list_slots(module: Option<&str>) -> PyResult<Vec<PySlotInfo>> {
    let mgr = if let Some(path) = module {
        Pkcs11Manager::new(path).map_err(|e| PyValueError::new_err(format!("{e}")))?
    } else {
        Pkcs11Manager::from_env().map_err(|e| PyValueError::new_err(format!("{e}")))?
    };
    let slots = mgr
        .list_slots()
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    Ok(slots
        .into_iter()
        .map(|s| PySlotInfo {
            slot_id: s.slot_id,
            token_label: s.token_label,
            manufacturer_id: s.manufacturer_id,
            model: s.model,
            serial_number: s.serial_number,
            flags: s.flags,
        })
        .collect())
}

// ── Helpers ───────────────────────────────────────────────────────────────────

fn build_key_spec(key_type: &str, param: &Bound<'_, PyAny>) -> PyResult<KeySpec> {
    match key_type.to_ascii_lowercase().as_str() {
        "rsa" => {
            let bits: u32 = param.extract().map_err(|e| {
                PyValueError::new_err(format!("RSA key requires an integer bit-length as param: {e}"))
            })?;
            Ok(KeySpec::Rsa(bits))
        }
        "ec" => {
            let curve: String = param.extract().map_err(|e| {
                PyValueError::new_err(format!(
                    "EC key requires a curve name string as param (e.g. 'P-256'): {e}"
                ))
            })?;
            Ok(KeySpec::Ec(curve))
        }
        "ed25519" => Ok(KeySpec::Ed25519),
        "ed448" => Ok(KeySpec::Ed448),
        "ml-dsa" => {
            let ps: String = param.extract().map_err(|e| {
                PyValueError::new_err(format!(
                    "ML-DSA key requires a parameter-set string (e.g. 'ML-DSA-65'): {e}"
                ))
            })?;
            Ok(KeySpec::MlDsa(ps))
        }
        "ml-kem" => {
            let ps: String = param.extract().map_err(|e| {
                PyValueError::new_err(format!(
                    "ML-KEM key requires a parameter-set string (e.g. 'ML-KEM-768'): {e}"
                ))
            })?;
            Ok(KeySpec::MlKem(ps))
        }
        other => Err(PyValueError::new_err(format!(
            "unsupported key type {other:?}; expected 'rsa', 'ec', 'ed25519', 'ed448', 'ml-dsa', or 'ml-kem'"
        ))),
    }
}

// ── Module registration ───────────────────────────────────────────────────────

pub fn register_pkcs11_module(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    let py = parent.py();
    let m = PyModule::new(py, "pkcs11")?;

    m.add_class::<PySlotInfo>()?;
    m.add_class::<PyKeyInfo>()?;
    m.add_class::<PyPkcs11Token>()?;
    m.add_function(wrap_pyfunction!(list_slots, &m)?)?;

    crate::install_submodule(
        parent,
        &m,
        "synta.pkcs11",
        Some("PKCS#11 token management: slot listing, key find/list/delete/generate."),
    )
}