aegisnet-hsm 0.1.0

use anyhow::{Context, Result};
use cryptoki::context::{CInitializeArgs, Pkcs11};
use cryptoki::object::Attribute;
use cryptoki::session::{Session, UserType};
use cryptoki::types::AuthPin;
use serde::{Deserialize, Serialize};

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct XmssSignature {
    pub signature: Vec<u8>,
    pub state_counter: u64,
}

pub struct HsmAtomicSigner {
    _pkcs11: Pkcs11,
    session: Session,
    key_handle: cryptoki::object::ObjectHandle,
    state_counter: u64,
}

impl HsmAtomicSigner {
    pub fn new(pkcs11_lib_path: &str, pin: &str, key_label: &str) -> Result<Self> {
        let pkcs11 = Pkcs11::new(pkcs11_lib_path).context("Failed to load PKCS#11 library")?;

        pkcs11
            .initialize(CInitializeArgs::OsThreads)
            .context("Failed to initialize PKCS#11")?;

        let slots = pkcs11
            .get_slots_with_token()
            .context("No slots with token found")?;

        let slot = *slots.first().context("No token slot found")?;

        let session = pkcs11
            .open_rw_session(slot)
            .context("Failed to open RW session")?;

        session
            .login(UserType::User, Some(&AuthPin::new(pin.to_string())))
            .context("Failed to login to HSM")?;

        let key_template = vec![Attribute::Label(key_label.as_bytes().to_vec())];
        let key_objects = session
            .find_objects(&key_template)
            .context("Failed to search for key")?;

        let key_handle = key_objects
            .into_iter()
            .next()
            .context("Signing key not found")?;

        Ok(Self {
            _pkcs11: pkcs11,
            session,
            key_handle,
            state_counter: 0,
        })
    }

    pub fn atomic_sign(&mut self, message: &[u8]) -> Result<XmssSignature> {
        let mechanism = cryptoki::mechanism::Mechanism::Sha256RsaPkcs;
        let signature = self
            .session
            .sign(&mechanism, self.key_handle, message)
            .context("HSM signing failed")?;

        self.state_counter += 1;

        Ok(XmssSignature {
            signature,
            state_counter: self.state_counter,
        })
    }
}