/*
 * Copyright 2025 Yubico AB
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::fmt;
use std::fmt::Display;
use strum_macros::EnumIter;
use serde::{Deserialize, Serialize};
use yubihsmrs::object::{ObjectAlgorithm, ObjectCapability, ObjectDescriptor, ObjectDomain, ObjectHandle, ObjectType};
use yubihsmrs::Session;
use crate::traits::operation_traits::YubihsmOperations;
use crate::traits::command_traits::Command;
use crate::common::error::MgmError;
use crate::common::algorithms::MgmAlgorithm;
use crate::common::validators::aes_shares_validator;
use crate::common::types::{NewObjectSpec, EXIT_LABEL};
use crate::common::util::{contains_all, get_object_descriptors};
use crate::hsm_operations::asym::AsymmetricOperations;
use crate::hsm_operations::sym::SymmetricOperations;


// 2 object ID bytes + 2 domains bytes + 8 capabilities bytes + 8 delegated capabilities = 20 bytes
static WRAP_SPLIT_PREFIX_LEN: usize = 20;

#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumIter)]
pub enum WrapCommand {
    List,
    GetKeyProperties,
    Generate,
    Import,
    Delete,
    GetPublicKey,
    ExportWrapped,
    ImportWrapped,
    GetRandom,
    Exit,
}

impl Command for WrapCommand {
    fn label(&self) -> &'static str {
        match self {
            Self::List => "List",
            Self::GetKeyProperties => "Get Object Properties",
            Self::Generate => "Generate Wrap Key",
            Self::Import => "Import Wrap Key",
            Self::Delete => "Delete Wrap Key",
            Self::GetPublicKey => "Get Public Key",
            Self::ExportWrapped => "Export objects under wrap",
            Self::ImportWrapped => "Import wrapped object",
            Self::GetRandom => "Generate pseudo random number",
            Self::Exit => EXIT_LABEL,
        }
    }

    fn description(&self) -> &'static str {
        match self {
            Self::List => "List all wrap keys stored on the YubiHSM",
            Self::GetKeyProperties => "Get properties of a wrap key stored on the YubiHSM",
            Self::Generate => "Generate a new wrap key inside the YubiHSM",
            Self::Import => "Import a wrap key into the YubiHSM",
            Self::Delete => "Delete a wrap key from the YubiHSM",
            Self::GetPublicKey => "Retrieve the public key portion of an RSA wrap key stored on the YubiHSM",
            Self::ExportWrapped => "Writes files ending with .yhw to specified directory",
            Self::ImportWrapped => "Import wrapped object from a .yhw file",
            Self::GetRandom => "Get pseudo random bytes generated by the YubiHSM",
            Self::Exit => "Exit YubiHSM Manager",
        }
    }

    fn required_capabilities(&self) -> &'static [ObjectCapability] {
        match self {
            Self::List | Self::GetKeyProperties | Self::Exit => &[],
            Self::Generate => &[ObjectCapability::GenerateWrapKey],
            Self::Import => &[ObjectCapability::PutWrapKey, ObjectCapability::PutPublicWrapKey],
            Self::Delete => &[ObjectCapability::DeleteWrapKey, ObjectCapability::DeletePublicWrapKey],
            Self::GetPublicKey => &[],
            Self::ExportWrapped => &[ObjectCapability::ExportWrapped],
            Self::ImportWrapped => &[ObjectCapability::ImportWrapped],
            Self::GetRandom => &[ObjectCapability::GetPseudoRandom],
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum WrapKeyType {
    #[default]
    Aes,
    Rsa,
    RsaPublic,
}

pub struct WrapKeyShares {
    pub shares: u8,
    pub threshold: u8,
    pub shares_data: Vec<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum WrapType {
    #[default]
    Object,
    Key,
}

impl Display for WrapType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            WrapType::Object => write!(f, "YubiHSM2 Object"),
            WrapType::Key => write!(f, "Key data only"),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct WrapOpSpec {
    pub wrapkey_id: u16,
    pub wrapkey_type: WrapKeyType,
    pub wrap_type: WrapType,
    pub include_ed_seed: bool,
    pub aes_algorithm: Option<ObjectAlgorithm>,
    pub oaep_algorithm: Option<ObjectAlgorithm>,
    pub mgf1_algorithm: Option<ObjectAlgorithm>,
}

pub struct WrappedData {
    pub object_id: u16,
    pub object_type: ObjectType,
    pub wrapkey_id: u16,
    pub wrapped_data: Vec<u8>,
    pub error: Option<MgmError>,
}

const AES_WRAP_KEY_CAPABILITIES: [ObjectCapability; 3] = [
    ObjectCapability::ExportWrapped,
    ObjectCapability::ImportWrapped,
    ObjectCapability::ExportableUnderWrap,
];

const RSA_WRAP_KEY_CAPABILITIES: [ObjectCapability; 2] = [
    ObjectCapability::ImportWrapped,
    ObjectCapability::ExportableUnderWrap,
];

const PUBLIC_WRAP_KEY_CAPABILITIES: [ObjectCapability; 2] = [
    ObjectCapability::ExportWrapped,
    ObjectCapability::ExportableUnderWrap,
];

pub struct WrapOperations;

impl YubihsmOperations for WrapOperations {

    fn context(&self) -> &'static str {
        WrapOperations::WRAP_CONTEXT
    }

    fn get_all_objects(&self, session: &Session) -> Result<Vec<ObjectDescriptor>, MgmError> {
        let mut keys = session.list_objects_with_filter(
            0, ObjectType::WrapKey, "", ObjectAlgorithm::ANY, &Vec::new())?;
        keys.extend_from_slice(&session.list_objects_with_filter(
            0, ObjectType::PublicWrapKey, "", ObjectAlgorithm::ANY,
            &Vec::new())?);
        get_object_descriptors(session, &keys)
    }

    fn get_generation_algorithms(&self) -> Vec<MgmAlgorithm> {
        MgmAlgorithm::WRAP_KEY_ALGORITHMS.to_vec()
    }

    fn get_object_capabilities(
        &self,
        object_type: Option<ObjectType>,
        object_algorithm: Option<ObjectAlgorithm>) -> Result<Vec<ObjectCapability>, MgmError> {

        if object_type.is_none() || object_algorithm.is_none() {
            return Err(MgmError::InvalidInput(
                "Missing object type and/or object algorithm".to_string(),
            ));
        }

        let key_type = Self::get_wrapkey_type(
            &object_type.unwrap(),
            &object_algorithm.unwrap(),
        )?;

        match key_type {
            WrapKeyType::Aes => Ok(AES_WRAP_KEY_CAPABILITIES.to_vec()),
            WrapKeyType::Rsa => Ok(RSA_WRAP_KEY_CAPABILITIES.to_vec()),
            WrapKeyType::RsaPublic => Ok(PUBLIC_WRAP_KEY_CAPABILITIES.to_vec()),
        }
    }

    fn generate(&self, session: &Session, spec: &NewObjectSpec) -> Result<u16, MgmError> {
        let id = session
            .generate_wrap_key(
                spec.id,
                &spec.label,
                &spec.domains,
                &spec.capabilities,
                spec.algorithm,
                &spec.delegated_capabilities)?;
        Ok(id)
    }

    fn import(&self, session: &Session, spec: &NewObjectSpec) -> Result<u16, MgmError> {
        let id = match spec.object_type {
            ObjectType::WrapKey => session.import_wrap_key(
                spec.id,
                &spec.label,
                &spec.domains,
                &spec.capabilities,
                spec.algorithm,
                &spec.delegated_capabilities,
                &spec.data[0])?,
            ObjectType::PublicWrapKey => session.import_public_wrap_key(
                spec.id,
                &spec.label,
                &spec.domains,
                &spec.capabilities,
                spec.algorithm,
                &spec.delegated_capabilities,
                &spec.data[0])?,
            _ => {
                return Err(MgmError::InvalidInput(
                    "ImportObjectSpec has invalid object type for wrap key".to_string(),
                ));
            }
        };
        Ok(id)    }
}

impl WrapOperations {

    pub const WRAP_CONTEXT: &'static str = "wrap";

    pub fn get_algorithm_from_keylen(keylen: usize) -> Result<ObjectAlgorithm, MgmError> {
        match keylen {
            16 => Ok(ObjectAlgorithm::Aes128CcmWrap),
            24 => Ok(ObjectAlgorithm::Aes192CcmWrap),
            32 => Ok(ObjectAlgorithm::Aes256CcmWrap),
            256 => Ok(ObjectAlgorithm::Rsa2048),
            384 => Ok(ObjectAlgorithm::Rsa3072),
            512 => Ok(ObjectAlgorithm::Rsa4096),
            _ => Err(MgmError::Error(format!("Unsupported key length {}", keylen))),
        }
    }

    pub fn get_wrapkey_type(object_type: &ObjectType, algorithm: &ObjectAlgorithm) -> Result<WrapKeyType, MgmError> {
        match object_type {
            ObjectType::PublicWrapKey => Ok(WrapKeyType::RsaPublic),
            ObjectType::WrapKey => match algorithm {
                ObjectAlgorithm::Aes128CcmWrap |
                ObjectAlgorithm::Aes192CcmWrap |
                ObjectAlgorithm::Aes256CcmWrap => Ok(WrapKeyType::Aes),
                ObjectAlgorithm::Rsa2048 |
                ObjectAlgorithm::Rsa3072 |
                ObjectAlgorithm::Rsa4096 => Ok(WrapKeyType::Rsa),
                _ => {
                    Err(MgmError::InvalidInput(
                        "Unsupported algorithm for wrap key".to_string(),
                    ))
                }
            },
            _ => Err(MgmError::InvalidInput(
                "Unsupported object type for wrap key".to_string(),
            )),
        }
    }

    fn get_wrapkeys_by_type(session:&Session, key_type:WrapKeyType) -> Result<Vec<ObjectDescriptor>, MgmError> {
        if key_type == WrapKeyType::RsaPublic {
            let keys = session.list_objects_with_filter(
                0, ObjectType::PublicWrapKey, "", ObjectAlgorithm::ANY,
                &Vec::new())?;
            Ok(get_object_descriptors(session, keys.as_slice())?)
        } else {
            let keys = session.list_objects_with_filter(
                0, ObjectType::WrapKey, "", ObjectAlgorithm::ANY, &Vec::new())?;
            let mut keys = get_object_descriptors(session, keys.as_slice())?;
            if key_type == WrapKeyType::Aes {
                keys.retain(|k| !AsymmetricOperations::is_rsa_key_algorithm(k.algorithm()));
            } else if key_type == WrapKeyType::Rsa {
                keys.retain(|k| AsymmetricOperations::is_rsa_key_algorithm(k.algorithm()));
            }
            Ok(keys)
        }
    }

    pub fn get_rsa_wrapkeys(session:&Session) -> Result<Vec<ObjectDescriptor>, MgmError> {
        let mut keys = Self.get_all_objects(session)?;
        keys.retain(|k| AsymmetricOperations::is_rsa_key_algorithm(k.algorithm()));
        Ok(keys)
    }

    pub fn get_wrapping_keys(session:&Session, authkey: &ObjectDescriptor) -> Result<Vec<ObjectDescriptor>, MgmError> {
        if !authkey.capabilities().contains(&ObjectCapability::ExportWrapped) {
            return Ok(Vec::new());
        }
        let mut keys = Self::get_wrapkeys_by_type(session, WrapKeyType::Aes)?;
        keys.extend_from_slice(&Self::get_wrapkeys_by_type(session, WrapKeyType::RsaPublic)?);
        keys.retain(|k| k.capabilities().contains(&ObjectCapability::ExportWrapped));
        if keys.is_empty() {
            return Err(MgmError::Error(
                format!("No wrap keys with {:?} capability were found", ObjectCapability::ExportWrapped)));
        }
        Ok(keys)
    }

    pub fn get_unwrapping_keys(session:&Session, authkey: &ObjectDescriptor) -> Result<Vec<ObjectDescriptor>, MgmError> {
        if !authkey.capabilities().contains(&ObjectCapability::ImportWrapped) {
            return Ok(Vec::new());
        }
        let mut keys = Self::get_wrapkeys_by_type(session, WrapKeyType::Aes)?;
        keys.extend_from_slice(&Self::get_wrapkeys_by_type(session, WrapKeyType::Rsa)?);
        keys.retain(|k| k.capabilities().contains(&ObjectCapability::ImportWrapped));
        if keys.is_empty() {
            return Err(MgmError::Error(
                format!("No wrap keys with {:?} capability were found", ObjectCapability::ImportWrapped)));
        }
        Ok(keys)
    }

    pub fn get_exportable_objects(session: &Session, wrap_key: &ObjectDescriptor, wrap_type: WrapType) -> Result<Vec<ObjectDescriptor>, MgmError> {
        if wrap_key.delegated_capabilities().is_none() {
            return Ok(Vec::new());
        }
        let delegated = wrap_key.delegated_capabilities().as_ref().unwrap();

        let objects = session.list_objects_with_filter(
            0,
            ObjectType::Any,
            "",
            ObjectAlgorithm::ANY,
            &[ObjectCapability::ExportableUnderWrap])?;
        let mut objects = get_object_descriptors(session, objects.as_slice())?;
        objects.retain(|obj| contains_all(delegated, obj.capabilities()));
        objects.retain(|obj| wrap_key.domains().iter().any(|d| obj.domains().contains(d)));
        if wrap_type == WrapType::Key {
            objects.retain(|obj| obj.object_type() == &ObjectType::AsymmetricKey || obj.object_type() == &ObjectType::SymmetricKey);
        }
        Ok(objects)
    }

    pub fn get_rsa_unwrapped_key_algorithms() -> Vec<MgmAlgorithm> {
        let mut algos = Vec::new();
        algos.extend_from_slice(AsymmetricOperations.get_generation_algorithms().as_slice());
        algos.extend_from_slice(SymmetricOperations.get_generation_algorithms().as_slice());
        algos
    }

    pub fn split_wrap_key(wrap_key:&NewObjectSpec, threshold:u8, shares:u8) -> Result<WrapKeyShares, MgmError> {
        let mut split_key = WrapKeyShares {
            shares,
            threshold,
            shares_data: Vec::new(),
        };

        let mut data = Vec::<u8>::new();
        data.push(((wrap_key.id >> 8) & 0xff) as u8);
        data.push((wrap_key.id & 0xff) as u8);
        data.append(&mut ObjectDomain::bytes_from_slice(wrap_key.domains.as_slice()));
        data.append(&mut ObjectCapability::bytes_from_slice(wrap_key.capabilities.as_slice()));
        data.append(&mut ObjectCapability::bytes_from_slice(wrap_key.delegated_capabilities.as_slice()));
        data.extend_from_slice(wrap_key.data[0].as_slice());

        let raw_shares = vsss_rs::Gf256::split_array(
            threshold as usize,
            shares as usize,
            &data,
            rand::rngs::OsRng,
        )?;

        split_key.shares_data = raw_shares
            .iter()
            .map(|share| {
                format!("{}-{}-{}", threshold, share[0], hex::encode(&share[1..]))
            })
            .collect();

        Ok(split_key)
    }

    pub fn get_wrapkey_from_shares(shares:Vec<String>) -> Result<NewObjectSpec, MgmError> {

        aes_shares_validator(&shares)?;

        let vsss_shares: Vec<Vec<u8>> = shares
            .iter()
            .map(|s| {
                let parts: Vec<&str> = s.trim().split('-').collect();
                let Ok(share_id) = parts[1].parse::<u8>() else {
                    return Err(MgmError::InvalidInput(format!("Unable to parse share ID: {}", s)));
                };
                let share_data = hex::decode(parts[2])?;
                let mut share = Vec::with_capacity(1 + share_data.len());
                share.push(share_id);
                share.extend_from_slice(&share_data);
                Ok(share)
            })
            .collect::<Result<Vec<_>, MgmError>>()?;

        let data = vsss_rs::Gf256::combine_array(&vsss_shares)?;
        if data.len() < WRAP_SPLIT_PREFIX_LEN {
            return Err(MgmError::InvalidInput("Reconstructed wrap key share data is too short".to_string()));
        }

        let key_len = data.len() - WRAP_SPLIT_PREFIX_LEN;

        let mut wrapkey_spec = NewObjectSpec::default();
        wrapkey_spec.object_type = ObjectType::WrapKey;
        wrapkey_spec.algorithm = WrapOperations::get_algorithm_from_keylen(key_len)?;
        wrapkey_spec.id = ((u16::from(data[0])) << 8) | u16::from(data[1]);
        wrapkey_spec.domains = ObjectDomain::from_bytes(&data[2..4])?;
        wrapkey_spec.capabilities = ObjectCapability::from_bytes(&data[4..12])?;
        wrapkey_spec.delegated_capabilities = ObjectCapability::from_bytes(&data[12..20])?;

        wrapkey_spec.data.push(data[20..].to_vec());

        Ok(wrapkey_spec)
    }

    pub fn export_wrapped(session: &Session, wrap_op_spec: &WrapOpSpec, export_objects: &Vec<ObjectHandle>) -> Result<Vec<WrappedData>, MgmError> {
        let mut wrapped = Vec::new();
        for object in export_objects {
            let mut w = WrappedData {
                object_id: object.object_id,
                object_type: object.object_type,
                wrapkey_id: wrap_op_spec.wrapkey_id,
                wrapped_data: Vec::new(),
                error: None,
            };

            let res = match wrap_op_spec.wrapkey_type {
                WrapKeyType::Aes => {
                    let format: u8 = if wrap_op_spec.include_ed_seed { 1 } else { 0 };
                    session.export_wrapped_ex(wrap_op_spec.wrapkey_id, object.object_type, object.object_id, format)
                },
                WrapKeyType::RsaPublic => {
                    let aes_algo = match wrap_op_spec.aes_algorithm {
                        Some(algo) => algo,
                        None => ObjectAlgorithm::Aes256,
                    };

                    let oaep_algo = match wrap_op_spec.oaep_algorithm {
                        Some(algo) => algo,
                        None => ObjectAlgorithm::RsaOaepSha256,
                    };

                    let oaep_label = Self::get_oaep_label(&oaep_algo)?;
                    let mgf1_algo = match wrap_op_spec.mgf1_algorithm {
                        Some(algo) => algo,
                        None => ObjectAlgorithm::Mgf1Sha256,
                    };

                    match wrap_op_spec.wrap_type {
                        WrapType::Object => session.export_rsa_wrapped_object(
                            wrap_op_spec.wrapkey_id,
                            object.object_type,
                            object.object_id,
                            aes_algo,
                            oaep_algo,
                            mgf1_algo,
                            &oaep_label,
                        ),
                        WrapType::Key => session.export_rsa_wrapped_key(
                            wrap_op_spec.wrapkey_id,
                            object.object_type,
                            object.object_id,
                            aes_algo,
                            oaep_algo,
                            mgf1_algo,
                            &oaep_label,
                        ),
                    }
                },
                WrapKeyType::Rsa => {
                    return Err(MgmError::InvalidInput("Private RSA key cannot be used for wrapping".to_string()))
                }
            };

            match res {
                Ok(bytes) => w.wrapped_data = bytes,
                Err(err) => w.error = Some(MgmError::LibYubiHsm(err))
            }
            wrapped.push(w);
        }
        Ok(wrapped)
    }

    pub fn import_wrapped(session: &Session, wrap_op_spec: &WrapOpSpec, wrapped: &[u8], new_key_spec: Option<NewObjectSpec>) -> Result<ObjectHandle, MgmError> {
        let data = wrapped.to_vec();

        let handle = match wrap_op_spec.wrapkey_type {
            WrapKeyType::Aes => session.import_wrapped(wrap_op_spec.wrapkey_id, &data)?,
            WrapKeyType::Rsa => {
                let oaep_algo = match wrap_op_spec.oaep_algorithm {
                    Some(algo) => algo,
                    None => ObjectAlgorithm::RsaOaepSha256,
                };
                let oaep_label = Self::get_oaep_label(&oaep_algo)?;
                let mgf1_algo = match wrap_op_spec.mgf1_algorithm {
                    Some(algo) => algo,
                    None => ObjectAlgorithm::Mgf1Sha256,
                };
                match wrap_op_spec.wrap_type {
                    WrapType::Object => session.import_rsa_wrapped_object(
                        wrap_op_spec.wrapkey_id,
                        oaep_algo,
                        mgf1_algo,
                        &oaep_label,
                        &data,
                    )?,
                    WrapType::Key => {
                        if new_key_spec.is_none() {
                            return Err(MgmError::InvalidInput("Object parameters for the imported wrapped key is missing".to_string()));
                        }
                        let key_spec = new_key_spec.unwrap();
                        session.import_rsa_wrapped_key(
                            wrap_op_spec.wrapkey_id,
                            key_spec.object_type,
                            key_spec.id,
                            key_spec.algorithm,
                            &key_spec.label,
                            &key_spec.domains,
                            &key_spec.capabilities,
                            oaep_algo,
                            mgf1_algo,
                            &oaep_label,
                            &data,
                        )?
                    }
                }
            }
            WrapKeyType::RsaPublic => {
                return Err(MgmError::InvalidInput("Public RSA key cannot be used for unwrapped".to_string()))
            }
        };

        Ok(handle)
    }

    fn get_oaep_label(algorithm: &ObjectAlgorithm) -> Result<Vec<u8>, MgmError> {
        let oaep_label: &[u8] = &[];
        AsymmetricOperations::get_hashed_bytes(algorithm, oaep_label)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── Helper: build a NewObjectSpec that looks like a wrap key ──
    fn make_wrapkey_spec(algo: ObjectAlgorithm, key_bytes: Vec<u8>) -> NewObjectSpec {
        NewObjectSpec {
            id: 0x0010,
            object_type: ObjectType::WrapKey,
            label: "test-wrap".to_string(),
            algorithm: algo,
            domains: vec![ObjectDomain::One],
            capabilities: vec![ObjectCapability::ExportWrapped],
            delegated_capabilities: vec![ObjectCapability::ExportableUnderWrap],
            data: vec![key_bytes],
        }
    }

    // ══════════════════════════════════════════════
    //  get_unwrapped_key_algorithms
    // ══════════════════════════════════════════════

    #[test]
    fn test_unwrapped_key_algorithms_non_empty() {
        let algos = WrapOperations::get_rsa_unwrapped_key_algorithms();
        let expected_len =
            MgmAlgorithm::RSA_KEY_ALGORITHMS.len() +
            MgmAlgorithm::EC_KEY_ALGORITHMS.len() +
            MgmAlgorithm::ED_KEY_ALGORITHMS.len() +
            MgmAlgorithm::AES_KEY_ALGORITHMS.len();
        assert_eq!(algos.len(), expected_len);
        assert!(algos.iter().all(|algo|
            MgmAlgorithm::RSA_KEY_ALGORITHMS.contains(algo) ||
            MgmAlgorithm::EC_KEY_ALGORITHMS.contains(algo) ||
            MgmAlgorithm::ED_KEY_ALGORITHMS.contains(algo) ||
            MgmAlgorithm::AES_KEY_ALGORITHMS.contains(algo)));
    }

    // ══════════════════════════════════════════════
    //  split_wrap_key + get_wrapkey_from_shares round-trip
    // ══════════════════════════════════════════════

    #[test]
    fn test_split_produces_correct_share_count() {
        let key = vec![0xABu8; 16]; // 128-bit key
        let spec = make_wrapkey_spec(ObjectAlgorithm::Aes128CcmWrap, key);
        let result = WrapOperations::split_wrap_key(&spec, 2, 3).unwrap();
        assert_eq!(result.shares, 3);
        assert_eq!(result.threshold, 2);
        assert_eq!(result.shares_data.len(), 3);
    }

    #[test]
    fn test_split_roundtrip_aes128() {
        let key = vec![0x42u8; 16];
        let spec = make_wrapkey_spec(ObjectAlgorithm::Aes128CcmWrap, key.clone());
        let split = WrapOperations::split_wrap_key(&spec, 2, 3).unwrap();

        // Reconstruct from any 2 of the 3 shares
        let subset = vec![split.shares_data[0].clone(), split.shares_data[1].clone()];
        let recovered = WrapOperations::get_wrapkey_from_shares(subset).unwrap();

        assert_eq!(recovered.id, spec.id);
        assert_eq!(recovered.object_type, spec.object_type);
        assert_eq!(recovered.label, ""); // Label is not included in the split data
        assert_eq!(recovered.algorithm, spec.algorithm);
        assert_eq!(recovered.domains, spec.domains);
        assert_eq!(recovered.capabilities, spec.capabilities);
        assert_eq!(recovered.delegated_capabilities, spec.delegated_capabilities);
        assert_eq!(recovered.data[0], key);
    }

    #[test]
    fn test_split_roundtrip_aes192() {
        let key = vec![0x47u8; 24];
        let spec = make_wrapkey_spec(ObjectAlgorithm::Aes192, key.clone());
        let split = WrapOperations::split_wrap_key(&spec, 3, 3).unwrap();
        assert_eq!(split.shares_data.len(), 3);

        let subset = vec![
            split.shares_data[0].clone(),
            split.shares_data[1].clone(),
            split.shares_data[2].clone(),
        ];
        let recovered = WrapOperations::get_wrapkey_from_shares(subset).unwrap();
        assert_eq!(recovered.data[0], key);
    }

    #[test]
    fn test_split_roundtrip_aes256() {
        let key = vec![0x99u8; 32];
        let spec = make_wrapkey_spec(ObjectAlgorithm::Aes256CcmWrap, key.clone());
        let split = WrapOperations::split_wrap_key(&spec, 3, 5).unwrap();
        assert_eq!(split.shares_data.len(), 5);

        let subset = vec![
            split.shares_data[0].clone(),
            split.shares_data[2].clone(),
            split.shares_data[4].clone(),
        ];
        let recovered = WrapOperations::get_wrapkey_from_shares(subset).unwrap();
        assert_eq!(recovered.data[0], key);
    }
}