#![deny(missing_docs)]
//! Declarative TPM.
//!
//! This crate allows operating on TPM objects using declarative descriptions of keys
//! and the associated artifacts.
//!
//! See binaries in `bin` for examples on how are these functions defined.

use serde::{Deserialize, Serialize};
use std::convert::{TryFrom, TryInto};
use std::str::FromStr;
use tss_esapi::attributes::object::ObjectAttributesBuilder;
use tss_esapi::attributes::session::SessionAttributesBuilder;
use tss_esapi::constants::session_type::SessionType;
use tss_esapi::constants::tss::*;
use tss_esapi::constants::CommandCode;
use tss_esapi::handles::{KeyHandle, PersistentTpmHandle, TpmHandle};
use tss_esapi::interface_types::algorithm::{
    EccSchemeAlgorithm, PublicAlgorithm, RsaSchemeAlgorithm,
};
use tss_esapi::interface_types::ecc::EccCurve;
use tss_esapi::interface_types::key_bits::RsaKeyBits;
use tss_esapi::interface_types::resource_handles::Hierarchy;
use tss_esapi::interface_types::session_handles::PolicySession;
use tss_esapi::structures::SymmetricDefinitionObject;
use tss_esapi::structures::{
    Auth, Digest, EccParameter, EccScheme, KeyDerivationFunctionScheme, Private, PublicBuilder,
    PublicEccParametersBuilder, PublicRsaParametersBuilder, RsaExponent, RsaScheme, Signature,
    SignatureScheme,
};
use tss_esapi::Result;

use tss_esapi::constants::tss::TPM2_ST_HASHCHECK;
use tss_esapi::structures::{
    Data, EccPoint, Public, PublicKeyRsa, RsaDecryptionScheme, SymmetricDefinition,
};
use tss_esapi::tss2_esys::TPMT_TK_HASHCHECK;
use tss_esapi::{Context, Tcti};

// Re-export as this is part of our public API.
pub use tss_esapi::interface_types::algorithm::HashingAlgorithm;

/// Description of a key.
///
/// The details of the key are specified by the `spec` field.
#[derive(Debug, Serialize, Deserialize)]
pub struct Description {
    /// Specification of the described key.
    pub spec: Specification,
}

/// Key details specification.
///
/// The specification of the contents of the key as well as by which provider it is managed.
/// While this crate defines only the TPM provider other crates could define other provider
/// types.
#[derive(Debug, Serialize, Deserialize)]
pub struct Specification {
    /// Provider to use for the key.
    pub provider: Provider,
    /// Algorithm used by the key.
    pub algo: AlgorithmSpec,
    /// Optional: key private bits.
    pub private: Option<PrivateKeyMaterial>,
    /// List of key capabilities.
    pub capabilities: Vec<Capability>,
    /// Authentication value to use the key.
    pub auth: String,
}

/// Key private material.
///
/// Defines data layout and contents of the private key definitions for both RSA and EC keys.
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PrivateKeyMaterial {
    /// RSA key material.
    Rsa(PrivateRsaKeyMaterial),
    /// EC key material.
    Ec(EcParameter),
}

/// RSA private key material.
///
/// Defines the contents of RSA private key: `prime` and `modulus` fields.
#[derive(Debug, Serialize, Deserialize)]
pub struct PrivateRsaKeyMaterial {
    /// RSA prime factor.
    pub prime: String,
    /// RSA modulus factor.
    pub modulus: RsaPublic,
}

impl From<&PrivateRsaKeyMaterial> for tss_esapi_sys::TPM2B_PRIVATE_KEY_RSA {
    fn from(private_rsa: &PrivateRsaKeyMaterial) -> Self {
        let key_prime = hex::decode(&private_rsa.prime).unwrap();
        let mut result = Self::default();
        result.buffer[..key_prime.len()].clone_from_slice(&key_prime[..key_prime.len()]);
        result.size = key_prime.len().try_into().unwrap();
        result
    }
}

/// EC private key material.
///
/// Defines private key material for EC: `parameter` and both `points`.
#[derive(Debug, Serialize, Deserialize)]
pub struct EcParameter {
    /// EC parameter.
    pub parameter: String,
    /// EC point.
    pub points: EcPublic,
}

impl From<&EcParameter> for tss_esapi_sys::TPM2B_ECC_PARAMETER {
    fn from(param: &EcParameter) -> Self {
        let parameter = hex::decode(&param.parameter).unwrap();
        let mut parameter_buffer = [0u8; 128];
        parameter_buffer[..parameter.len()].clone_from_slice(&parameter);
        Self {
            size: parameter.len() as u16,
            buffer: parameter_buffer,
        }
    }
}

/// Key capability.
///
/// Defines for which usages is this key intended. All keys need to be either `Decrypt` or `Sign`
/// keys. Some keys can also be `Restrict`ed keys.
///
/// Keys that are both `Decrypt` and `Sign` are called storage keys and are used to load other
/// keys in the TPM. Restricted decryption keys will not reveal the plaintext to non-TPM code.
///
/// Restricted signing keys will sign only digests produced in the TPM. They are used mostly
/// to sign data produced by the TPM, e.g. the attestation data.
#[derive(Debug, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Capability {
    /// Allows decryption using this key.
    Decrypt,
    /// Allows signing using this key.
    Sign,
    /// This key is restricted.
    Restrict,
    /// This is an unknown, unsupported capability.
    #[serde(other)]
    Unknown,
}

/// Algorithm specification.
///
/// If a new key is about to be created the specifics of the algorithm need to be defined.
/// This object specifies details for each supported asymmetric algorithm.
#[derive(Debug, Serialize, Deserialize)]
pub enum AlgorithmSpec {
    /// RSA bits and optional exponent.
    #[serde(rename = "RSA")]
    Rsa {
        /// Number of bits for this key.
        bits: u16,
        /// The exponent if it's different from 0x010001.
        exponent: Option<u32>,
    },
    /// EC curve.
    #[serde(rename = "EC")]
    Ec {
        /// Elliptic Curve to use.
        curve: EcCurve,
    },
}

/// EC curve.
///
/// EC curve that is supported by this library.
#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub enum EcCurve {
    /// NIST P-256 curve.
    #[serde(rename = "NIST-P256")]
    NistP256,
    /// NIST P-384 curve.
    #[serde(rename = "NIST-P384")]
    NistP384,
}

impl From<&EcCurve> for EccCurve {
    fn from(curve: &EcCurve) -> Self {
        match curve {
            EcCurve::NistP256 => EccCurve::NistP256,
            EcCurve::NistP384 => EccCurve::NistP384,
        }
    }
}

/// Cryptographic provider.
///
/// Specifies cryptographic provider that should be used for all cryptographic operations.
#[derive(Debug, Serialize, Deserialize)]
pub struct Provider {
    /// Defines details for the TPM provider.
    pub tpm: TpmProvider,
}

/// TPM cryptographic provider.
///
/// Specifies details of the target TPM as well as several TPM-specific key properties.
#[derive(Debug, Serialize, Deserialize)]
pub struct TpmProvider {
    /// Connection string for the provider. E.g. `mssim:` for the simulator or
    /// `device:/dev/tpmrm0` for a hardware device.
    pub tcti: String,
    /// Numeric handle of the key if it is persisted.
    pub handle: Option<u32>,
    /// Numeric handle of the parent of this key if the parent is persisted.
    pub parent: Option<u32>,
    /// Wrapped private bytes of the key.
    pub private: Option<String>,
    /// Wrapped unique/public bytes of the key.
    pub unique: Option<PublicKeyBytes>,
    /// Key has been wrapped by [`wrap`].
    pub wrapped: Option<WrappedKey>,
    /// Policy of key usage.
    pub policy: Option<String>,
}

/// Key that has been wrapped.
///
/// Some cryptographic providers can export a wrapped (encrypted) key. This structure
/// defines details of these keys.
#[derive(Debug, Serialize, Deserialize)]
pub struct WrappedKey {
    /// Secret value that has been wrapped.
    pub secret: String,
    /// Private value that has been wrapped.
    pub private: String,
    /// Encrypted data for wrapping.
    pub data: String,
}

/// Public key bits.
///
/// For export of public keys this structure is used with one variant per one supported
/// asymmetric algorithm.
#[derive(Debug, Serialize, Deserialize)]
pub enum PublicKeyBytes {
    /// RSA public key details.
    RSA(RsaPublic),
    /// EC public details.
    EC(EcPublic),
}

/// RSA public key.
///
/// Defines raw `bytes` of the public key of the RSA key.
#[derive(Debug, Serialize, Deserialize)]
pub struct RsaPublic {
    /// RSA raw bytes.
    pub bytes: String,
}

impl From<&RsaPublic> for PublicKeyRsa {
    fn from(public: &RsaPublic) -> Self {
        let public_modulus = hex::decode(&public.bytes).unwrap();
        PublicKeyRsa::try_from(public_modulus).unwrap()
    }
}

/// EC public point.
///
/// Defines the coordinates of EC's public point.
#[derive(Debug, Serialize, Deserialize)]
pub struct EcPublic {
    /// Point's `x` coordinate.
    pub x: String,
    /// Point's `y` coordinate.
    pub y: String,
}

impl From<&EcPublic> for EccPoint {
    fn from(public: &EcPublic) -> Self {
        let x = hex::decode(&public.x).unwrap();
        let y = hex::decode(&public.y).unwrap();

        EccPoint::new(
            EccParameter::try_from(x).unwrap(),
            EccParameter::try_from(y).unwrap(),
        )
    }
}

/// Create key given by the `spec`.
///
/// Creates given key and returns the `PublicBuilder` for that key. If the key was a private key import
/// this function additionally returns a Sensitive value.
pub fn create(
    spec: &Specification,
) -> Result<(PublicBuilder, Option<tss_esapi_sys::TPM2B_SENSITIVE>)> {
    let is_fixed = spec.private.is_none()
        && spec.provider.tpm.wrapped.is_none()
        && spec.provider.tpm.policy.is_none();
    let attributes = ObjectAttributesBuilder::new()
        .with_fixed_tpm(is_fixed)
        .with_fixed_parent(is_fixed)
        .with_sensitive_data_origin(true)
        .with_user_with_auth(true)
        .with_decrypt(spec.capabilities.contains(&Capability::Decrypt))
        .with_sign_encrypt(spec.capabilities.contains(&Capability::Sign))
        .with_restricted(spec.capabilities.contains(&Capability::Restrict))
        .build()?;

    let mut builder = PublicBuilder::new()
        .with_name_hashing_algorithm(HashingAlgorithm::Sha256)
        .with_object_attributes(attributes);

    if let Some(policy) = &spec.provider.tpm.policy {
        builder = builder.with_auth_policy(Digest::try_from(hex::decode(policy).unwrap())?);
    }

    if let Some(unique) = &spec.provider.tpm.unique {
        builder = match unique {
            PublicKeyBytes::RSA(ref bytes) => builder.with_rsa_unique_identifier(bytes.into()),
            PublicKeyBytes::EC(ref bytes) => builder.with_ecc_unique_identifier(bytes.into()),
        }
    } else {
        builder = builder
            .with_rsa_unique_identifier(Default::default())
            .with_ecc_unique_identifier(Default::default());
    }

    builder = match &spec.algo {
        AlgorithmSpec::Rsa { bits, exponent } => {
            let mut rsa_params_builder = PublicRsaParametersBuilder::new();
            if spec.capabilities.contains(&Capability::Restrict) {
                rsa_params_builder =
                    rsa_params_builder.with_symmetric(SymmetricDefinitionObject::AES_256_CFB);
            }
            rsa_params_builder = rsa_params_builder
                .with_scheme(if spec.capabilities.contains(&Capability::Decrypt) {
                    RsaScheme::Null
                } else if spec.capabilities.contains(&Capability::Sign) {
                    RsaScheme::create(RsaSchemeAlgorithm::RsaSsa, Some(HashingAlgorithm::Sha256))
                        .unwrap()
                } else {
                    panic!("Key needs to be for decryption or for signing")
                })
                .with_key_bits(RsaKeyBits::try_from(*bits).unwrap())
                .with_exponent(RsaExponent::try_from(exponent.unwrap_or(0)).unwrap())
                .with_is_signing_key(spec.capabilities.contains(&Capability::Sign))
                .with_is_decryption_key(spec.capabilities.contains(&Capability::Decrypt))
                .with_restricted(spec.capabilities.contains(&Capability::Restrict));

            let rsa_params = rsa_params_builder.build()?;

            builder
                .with_public_algorithm(PublicAlgorithm::Rsa)
                .with_rsa_parameters(rsa_params)
        }
        AlgorithmSpec::Ec { ref curve } => {
            let mut ecc_builder = PublicEccParametersBuilder::new()
                .with_ecc_scheme(if spec.capabilities.contains(&Capability::Decrypt) {
                    EccScheme::Null
                } else if spec.capabilities.contains(&Capability::Sign) {
                    EccScheme::create(
                        EccSchemeAlgorithm::EcDsa,
                        Some(HashingAlgorithm::Sha256),
                        None,
                    )
                    .unwrap()
                } else {
                    panic!("Key needs to be for decryption or for signing")
                })
                .with_curve(curve.into())
                .with_is_signing_key(spec.capabilities.contains(&Capability::Sign))
                .with_is_decryption_key(spec.capabilities.contains(&Capability::Decrypt))
                .with_restricted(spec.capabilities.contains(&Capability::Restrict))
                .with_key_derivation_function_scheme(KeyDerivationFunctionScheme::Null);
            if spec.capabilities.contains(&Capability::Restrict) {
                ecc_builder = ecc_builder.with_symmetric(SymmetricDefinitionObject::AES_256_CFB);
            }

            builder = builder
                .with_public_algorithm(PublicAlgorithm::Ecc)
                .with_ecc_parameters(ecc_builder.build()?);

            builder
        }
    };

    let private = match spec.private {
        Some(PrivateKeyMaterial::Rsa(ref private_rsa)) => {
            let rsa: tss_esapi_sys::TPM2B_PRIVATE_KEY_RSA = private_rsa.into();

            builder = builder.with_rsa_unique_identifier((&private_rsa.modulus).into());
            Some(tss_esapi_sys::TPM2B_SENSITIVE {
                size: rsa.size,
                sensitiveArea: tss_esapi_sys::TPMT_SENSITIVE {
                    sensitiveType: TPM2_ALG_RSA,
                    authValue: Default::default(),
                    seedValue: Default::default(),
                    sensitive: tss_esapi_sys::TPMU_SENSITIVE_COMPOSITE { rsa },
                },
            })
        }
        Some(PrivateKeyMaterial::Ec(ref param)) => {
            let ecc: tss_esapi_sys::TPM2B_ECC_PARAMETER = param.into();
            builder = builder.with_ecc_unique_identifier((&(param.points)).into());
            Some(tss_esapi_sys::TPM2B_SENSITIVE {
                size: ecc.size,
                sensitiveArea: tss_esapi_sys::TPMT_SENSITIVE {
                    sensitiveType: TPM2_ALG_ECC,
                    authValue: Default::default(),
                    seedValue: Default::default(),
                    sensitive: tss_esapi_sys::TPMU_SENSITIVE_COMPOSITE { ecc },
                },
            })
        }
        _ => None,
    };

    Ok((builder, private))
}

/// Converts `specification` to raw key handle.
///
/// TPM functions require a key handle as a key identifier. This function converts the key
/// specification to the internal key handle object.
pub fn convert_to_key_handle(
    context: &mut Context,
    specification: &Specification,
) -> Result<KeyHandle> {
    let key_handle = if let (public, Some(private)) = create(specification)? {
        context.load_external(private.try_into()?, public.build()?, Hierarchy::Null)?
    } else if let Some(handle) = specification.provider.tpm.handle {
        let persistent_tpm_handle = PersistentTpmHandle::new(handle)?;

        let handle = context.execute_without_session(|ctx| {
            ctx.tr_from_tpm_public(TpmHandle::Persistent(persistent_tpm_handle))
                .expect("Need handle")
        });

        handle.into()
    } else if let (Some(parent), Some(private)) = (
        specification.provider.tpm.parent,
        &specification.provider.tpm.private,
    ) {
        let persistent_tpm_handle = PersistentTpmHandle::new(parent)?;

        let handle = context.execute_without_session(|ctx| {
            ctx.tr_from_tpm_public(TpmHandle::Persistent(persistent_tpm_handle))
                .expect("Need handle")
        });

        let key_handle: KeyHandle = handle.into();

        context.tr_set_auth(
            key_handle.into(),
            Auth::try_from(specification.auth.as_bytes())?,
        )?;

        context.load(
            key_handle,
            Private::try_from(hex::decode(private).unwrap())?,
            create(specification)?.0.build()?,
        )?
    } else {
        panic!("Cannot load key");
    };

    context.tr_set_auth(
        key_handle.into(),
        Auth::try_from(specification.auth.as_bytes())?,
    )?;

    Ok(key_handle)
}

/// Read public key.
///
/// Refreshes the key specification with data available in the cryptographic provider
/// most notably the public key.
pub fn read_key(spec: &mut Specification) -> Result<()> {
    let tcti = Tcti::from_str(&spec.provider.tpm.tcti)?;

    let mut context = Context::new(tcti)?;

    let session = context.start_auth_session(
        None,
        None,
        None,
        SessionType::Hmac,
        SymmetricDefinition::AES_256_CFB,
        HashingAlgorithm::Sha256,
    )?;
    let (session_attr, session_mask) = SessionAttributesBuilder::new()
        .with_decrypt(true)
        .with_encrypt(true)
        .build();
    context
        .tr_sess_set_attributes(session.unwrap(), session_attr, session_mask)
        .unwrap();
    context.set_sessions((session, None, None));

    let key_handle = convert_to_key_handle(&mut context, spec)?;

    let (public, _, _) = context.read_public(key_handle)?;

    let public_key = match &public {
        Public::Rsa { unique, .. } => PublicKeyBytes::RSA(RsaPublic {
            bytes: hex::encode(unique.value()),
        }),
        Public::Ecc { unique, .. } => PublicKeyBytes::EC(EcPublic {
            x: hex::encode(unique.x().value()),
            y: hex::encode(unique.y().value()),
        }),
        _ => panic!("Unsupported key type."),
    };

    spec.provider.tpm.unique = Some(public_key);
    spec.provider.tpm.policy = hex::encode(
        match &public {
            tss_esapi::structures::Public::Rsa { auth_policy, .. } => auth_policy,
            tss_esapi::structures::Public::Ecc { auth_policy, .. } => auth_policy,
            _ => panic!("Unsupported key type."),
        }
        .value(),
    )
    .into();

    Ok(())
}

/// Decrypt a buffer of data.
///
/// Use RSA decrypt function to decrypt a piece of `ciphertext` into a vector
/// of bytes.
///
/// *Note*: for EC keys use [`derive`].
pub fn decrypt(spec: &Specification, ciphertext: &[u8]) -> Result<Vec<u8>> {
    let tcti = Tcti::from_str(&spec.provider.tpm.tcti)?;

    let mut context = Context::new(tcti)?;

    let session = context.start_auth_session(
        None,
        None,
        None,
        SessionType::Hmac,
        SymmetricDefinition::AES_256_CFB,
        HashingAlgorithm::Sha256,
    )?;
    let (session_attr, session_mask) = SessionAttributesBuilder::new()
        .with_decrypt(true)
        .with_encrypt(true)
        .build();
    context
        .tr_sess_set_attributes(session.unwrap(), session_attr, session_mask)
        .unwrap();
    context.set_sessions((session, None, None));

    let key_handle = convert_to_key_handle(&mut context, spec)?;

    let cipher_text = PublicKeyRsa::try_from(ciphertext)?;

    let plain_text = context.rsa_decrypt(
        key_handle,
        cipher_text,
        RsaDecryptionScheme::RsaEs,
        Data::default(),
    )?;

    Ok(plain_text.to_vec())
}

/// Derives a common point.
///
/// Use EC algorithm to derive a common point on a curve for both the key
/// defined by `spec` and the public Z point defined in `data`.
pub fn derive(spec: &Specification, data: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
    let tcti = Tcti::from_str(&spec.provider.tpm.tcti)?;

    let mut context = Context::new(tcti)?;

    let session = context.start_auth_session(
        None,
        None,
        None,
        SessionType::Hmac,
        SymmetricDefinition::AES_256_CFB,
        HashingAlgorithm::Sha256,
    )?;
    let (session_attr, session_mask) = SessionAttributesBuilder::new()
        .with_decrypt(true)
        .with_encrypt(true)
        .build();
    context
        .tr_sess_set_attributes(session.unwrap(), session_attr, session_mask)
        .unwrap();
    context.set_sessions((session, None, None));

    let key_handle = convert_to_key_handle(&mut context, spec)?;

    let size = data.len() / 2;

    let z_point = context.ecdh_z_gen(
        key_handle,
        EccPoint::new(
            EccParameter::try_from(&data[..size])?,
            EccParameter::try_from(&data[size..])?,
        ),
    )?;

    Ok((z_point.x().value().to_vec(), z_point.y().value().to_vec()))
}

/// Sign a hash of data.
///
/// Uses key defined by `spec` to sign a digest specified by `hash` and returns raw
/// signatures bytes as a vector. For RSA the vector specifies raw signature values.
/// For EC the vector is a concatenation of signature's R and S values.
///
/// Note: most TPMs only support SHA256.
pub fn sign(spec: &Specification, hash_algo: HashingAlgorithm, hash: &[u8])
    -> Result<Vec<u8>>
{
    let tcti = Tcti::from_str(&spec.provider.tpm.tcti)?;

    let mut context = Context::new(tcti)?;

    let session = context.start_auth_session(
        None,
        None,
        None,
        SessionType::Hmac,
        SymmetricDefinition::AES_256_CFB,
        hash_algo,
    )?;
    let (session_attr, session_mask) = SessionAttributesBuilder::new()
        .with_decrypt(true)
        .with_encrypt(true)
        .build();
    context
        .tr_sess_set_attributes(session.unwrap(), session_attr, session_mask)
        .unwrap();
    context.set_sessions((session, None, None));

    let key_handle = convert_to_key_handle(&mut context, spec)?;

    let scheme = SignatureScheme::Null;
    let validation = TPMT_TK_HASHCHECK {
        tag: TPM2_ST_HASHCHECK,
        hierarchy: TPM2_RH_NULL,
        digest: Default::default(),
    }
    .try_into()?;

    let digest = Digest::try_from(hash)?;

    let signature = context.sign(key_handle, digest, scheme, validation)?;

    Ok(match signature {
        Signature::RsaSsa(ref signature) => Vec::from(signature.signature().value()),
        Signature::EcDsa(signature) => {
            let mut sig = vec![];
            sig.extend(signature.signature_r().value());
            sig.extend(signature.signature_s().value());
            sig
        }
        _ => panic!("Unsupported signature scheme."),
    })
}

/// Wrap (encrypt) a key.
///
/// Wraps the key specified by `spec` with the key specified by `parent`. Effectively this
/// encrypts the private portion of the `spec` key with public key of the `parent` key.
/// The `spec` key needs to allow wrapping (e.g. using PolicyDuplicationSelect) and the `parent`
/// key needs to support import of wrapped keys (e.g. using Decrypt+Restricted capabilities).
pub fn wrap(spec: &mut Specification, parent: &Specification) -> Result<()> {
    let tcti = Tcti::from_str(&spec.provider.tpm.tcti)?;

    let mut context = Context::new(tcti)?;

    // create a policy digest that allows key duplication
    let trial_session = context
        .start_auth_session(
            None,
            None,
            None,
            SessionType::Trial,
            SymmetricDefinition::AES_256_CFB,
            HashingAlgorithm::Sha256,
        )?
        .expect("Start auth session returned a NONE handle");

    let (policy_auth_session_attributes, policy_auth_session_attributes_mask) =
        SessionAttributesBuilder::new()
            .with_decrypt(true)
            .with_encrypt(true)
            .build();

    context.tr_sess_set_attributes(
        trial_session,
        policy_auth_session_attributes,
        policy_auth_session_attributes_mask,
    )?;

    let policy_session = PolicySession::try_from(trial_session)?;

    context.policy_auth_value(policy_session)?;

    context.policy_command_code(policy_session, CommandCode::Duplicate)?;
    let digest = context.policy_get_digest(policy_session)?;
    // end of: create policy digest

    let session = context.start_auth_session(
        None,
        None,
        None,
        SessionType::Hmac,
        SymmetricDefinition::AES_256_CFB,
        HashingAlgorithm::Sha256,
    )?;
    let (session_attr, session_mask) = SessionAttributesBuilder::new()
        .with_decrypt(true)
        .with_encrypt(true)
        .build();
    context.tr_sess_set_attributes(session.unwrap(), session_attr, session_mask)?;
    context.set_sessions((session, None, None));

    let key_handle = if let (public, Some(private)) = create(spec)? {
        context.load_external(
            private.try_into()?,
            public.with_auth_policy(digest).build()?,
            Hierarchy::Null,
        )?
    } else {
        panic!("Can import only private keys");
    };

    let parent_handle =
        context.load_external_public(create(parent)?.0.build()?, Hierarchy::Null)?;

    let (public, _, _) = context.read_public(key_handle)?;

    let public_key = match &public {
        Public::Rsa { unique, .. } => PublicKeyBytes::RSA(RsaPublic {
            bytes: hex::encode(unique.value()),
        }),
        Public::Ecc { unique, .. } => PublicKeyBytes::EC(EcPublic {
            x: hex::encode(unique.x().value()),
            y: hex::encode(unique.y().value()),
        }),
        _ => panic!("Unsupported key type."),
    };

    let auth_policy = match &public {
        tss_esapi::structures::Public::Rsa { auth_policy, .. } => auth_policy,
        tss_esapi::structures::Public::Ecc { auth_policy, .. } => auth_policy,
        _ => panic!("Unsupported key type."),
    }
    .value();

    spec.provider.tpm.unique = Some(public_key);
    spec.private = None;

    context.set_sessions((None, None, None));

    let policy_auth_session = context
        .start_auth_session(
            None,
            None,
            None,
            SessionType::Policy,
            SymmetricDefinition::AES_256_CFB,
            HashingAlgorithm::Sha256,
        )?
        .expect("Start auth session returned a NONE handle");
    let (policy_auth_session_attributes, policy_auth_session_attributes_mask) =
        SessionAttributesBuilder::new()
            .with_decrypt(true)
            .with_encrypt(true)
            .build();
    context.tr_sess_set_attributes(
        policy_auth_session,
        policy_auth_session_attributes,
        policy_auth_session_attributes_mask,
    )?;

    let policy_session = PolicySession::try_from(policy_auth_session)?;

    context.policy_auth_value(policy_session)?;

    context.policy_command_code(policy_session, CommandCode::Duplicate)?;
    context.set_sessions((Some(policy_auth_session), None, None));

    let (data, private, secret) = context.duplicate(
        key_handle.into(),
        parent_handle.into(),
        None,
        SymmetricDefinitionObject::Null,
    )?;

    spec.provider.tpm.wrapped = Some(WrappedKey {
        private: hex::encode(private.value()),
        secret: hex::encode(secret.value()),
        data: hex::encode(data.value()),
    });
    spec.provider.tpm.policy = Some(hex::encode(auth_policy));
    spec.provider.tpm.parent = parent.provider.tpm.handle;

    Ok(())
}
