use std::borrow::Cow;

use anyhow::Context;

use sequoia_openpgp as openpgp;
use openpgp::crypto::Decryptor as _;
use openpgp::crypto::ecdh::aes_key_unwrap;
use openpgp::crypto::mem::Protected;
use openpgp::crypto::mpi;
use openpgp::crypto::{Curve, SessionKey, SymmetricAlgorithm};
use openpgp::crypto;
use openpgp::packet::key::PublicParts;
use openpgp::packet::key::UnspecifiedRole;
use openpgp::packet::key;
use openpgp::packet;
use openpgp::types::HashAlgorithm;
use openpgp::types::PublicKeyAlgorithm;

use cryptoki::mechanism::{Mechanism, MechanismType};
use cryptoki::mechanism::elliptic_curve::EcKdf;
use cryptoki::mechanism::elliptic_curve::Ecdh1DeriveParams;
use cryptoki::mechanism::hkdf::{HkdfParams, HkdfSalt};
use cryptoki::object::Attribute;
use cryptoki::object::AttributeType;
use cryptoki::object::KeyType;
use cryptoki::object::ObjectClass;
use cryptoki::session::UserType;
use cryptoki::types::AuthPin;
use cryptoki::types::Ulong;

use crate::Result;
use crate::Prompt;
use crate::key_info::KeyInfo;
use crate::objects::Key;

const TRACE: bool = false;

/// A decryptor for a PKCS#11 key.
///
/// This implements [`sequoia_openpgp::crypto::Decryptor`] and hence can
/// be used wherever a Sequoia method needs to decrypt data.
///
///   [`sequoia_openpgp::crypto::Decryptor`]: https://docs.rs/sequoia-openpgp/latest/sequoia_openpgp/crypto/trait.Decryptor.html
pub struct Decryptor<'key, 'prompt> {
    key: Cow<'key, Key>,
    openpgp_key: packet::Key<PublicParts, UnspecifiedRole>,
    user_pin: Option<AuthPin>,
    prompt: Box<dyn Prompt + 'prompt>,
}

impl<'key, 'prompt> Decryptor<'key, 'prompt> {
    /// Creates a `Decryptor`.
    ///
    /// This is a bit ugly, because we can't stash a PKCS#11
    /// `Session`: a struct that implements `Decryptor` has to
    /// implement `Sync`, but a `Session` is not `Sync`.
    ///
    /// If the user PIN is known, pass it through to avoid prompting
    /// the user.  (If `user_pin` is wrong, we don't prompt the user,
    /// but immediately return an error.)  Otherwise, the user is
    /// prompted.
    pub fn new<P>(prompt: P, key: &'key Key, user_pin: Option<AuthPin>)
        -> Result<Self>
    where
        P: Prompt + 'prompt,
    {
        // Before we prompt for the PIN, make sure we can convert the
        // public key to an OpenPGP key.
        let openpgp_key = key.to_openpgp_key()
            .context("Converting public key to OpenPGP format")?
            .role_into_unspecified();

        Ok(Self {
            key: Cow::Borrowed(key),
            openpgp_key,
            user_pin,
            prompt: Box::new(prompt),
        })
    }

    /// Creates a `Decryptor`.
    ///
    /// Like [`Decryptor::new`], but takes ownership of `key`.
    pub fn new_owned<P>(prompt: P, key: Key, user_pin: Option<AuthPin>)
        -> Result<Self>
    where
        P: Prompt + 'prompt,
    {
        // Before we prompt for the PIN, make sure we can convert the
        // public key to an OpenPGP key.
        let openpgp_key = key.to_openpgp_key()
            .context("Converting public key to OpenPGP format")?
            .role_into_unspecified();

        Ok(Self {
            key: Cow::Owned(key),
            openpgp_key,
            user_pin,
            prompt: Box::new(prompt),
        })
    }
}

impl crypto::Decryptor for Decryptor<'_, '_> {
    fn public(&self) -> &packet::Key<PublicParts, UnspecifiedRole> {
        &self.openpgp_key
    }

    fn decrypt(&mut self, ciphertext: &mpi::Ciphertext,
               plaintext_len: Option<usize>)
        -> Result<crypto::SessionKey>
    {
        // We wrap the real implementation so that we can more easily
        // trace any errors.
        tracer!(TRACE, "Decryptor::decrypt");
        self.decrypt_internal(ciphertext, plaintext_len)
            .map_err(|err| {
                t!("Decryption error: {}", err);
                err
            })
    }
}

impl Decryptor<'_, '_> {
    fn decrypt_internal(&mut self,
                        ciphertext: &mpi::Ciphertext,
                        plaintext_len: Option<usize>)
        -> Result<crypto::SessionKey>
    {
        tracer!(TRACE, "Decryptor::decrypt_internal");
        t!("ciphertext: {:?}, plaintext_len: {:?}",
           ciphertext, plaintext_len);

        let (session, _object) = self.key.open_ro_session()?;

        // Login as the user.
        if let Some(user_pin) = self.user_pin.as_ref() {
            // PIN was supplied.
            session.login(UserType::User, Some(&user_pin))
                .context("Logging in to token")?;
        } else {
            // Prompt the user for the PIN.
            loop {
                let user_pin = self.prompt.prompt_to_unlock("the user PIN")?;
                let user_pin = String::from_utf8(user_pin.map(|p| p.to_vec()))
                    .context("Parsing user PIN")?;
                let user_pin = AuthPin::new(user_pin.into());

                match session.login(UserType::User, Some(&user_pin)) {
                    Ok(_) => {
                        self.user_pin = Some(user_pin);
                        break;
                    }
                    Err(err) => {
                        eprintln!("Error logging in: {}", err);
                    }
                }
            }
        }

        // The default public session only allows enumerating public
        // key objects.  We now look up the private key object using
        // the user session.  By convention, the public and private
        // keys have the same ID (CKA_ID).
        let key_info = KeyInfo::new(&session, &self.key)?;
        let composite_info = if let Some(composite)
            = self.key.composite_key().as_ref()
        {
            Some(KeyInfo::new(&session, composite)?)
        } else {
            None
        };

        let decrypt_mlkem = |pk_algo,
                             key_info: KeyInfo,
                             composite_info: Option<KeyInfo>,
                             ecdh_public: &[u8],
                             ecdh_part: &[u8],
                             mlkem_part: &[u8],
                             esk: &[u8]|
        {
            // Decrypt the ML-KEM key share.

            let (mlkem_info, ecdh_info) = key_info.check_composite(
                &self.openpgp_key, composite_info,
                KeyType::ML_KEM, KeyType::EC_MONTGOMERY)?;

            let attributes = &[
                Attribute::Class(ObjectClass::SECRET_KEY),
                Attribute::KeyType(KeyType::GENERIC_SECRET),
                // Don't persist the session key on the token.
                Attribute::Token(false),
                // Mark the session key bundle as extractable.
                Attribute::Extractable(true),
                Attribute::Sensitive(false),
            ];

            let mlkem_share_handle = session.decapsulate_key(
                &Mechanism::MlKem, mlkem_info.private_key,
                attributes,
                &mlkem_part[..])
                .context("Decrypting ML-KEM part")?;

            // Extract the share from the token.
            let mut mlkem_share = None;
            for attribute in session.get_attributes(
                mlkem_share_handle,
                &[ AttributeType::Value ])?
            {
                if let Attribute::Value(v) = attribute {
                    mlkem_share = Some(v);
                } else {
                    return Err(anyhow::anyhow!(
                        "Token returned unrequested attribute: {:?}",
                        attribute));
                }
            }

            let Some(mlkem_share) = mlkem_share else {
                return Err(anyhow::anyhow!(
                    "Failed to extract ML-KEM share from token"));
            };

            // Decrypt the ECDH key share.

            let kdf = EcKdf::null();
            let ecdh1_derive_params = Ecdh1DeriveParams::new(
                kdf, &ecdh_part[..]);

            let attributes = &[
                Attribute::Class(ObjectClass::SECRET_KEY),
                Attribute::KeyType(KeyType::GENERIC_SECRET),
                // Don't persist the KEK on the token.
                Attribute::Token(false),
                // Mark the session key bundle as extractable.
                Attribute::Extractable(true),
                Attribute::Sensitive(false),
            ];

            t!("Deriving ECDH share from {} bytes", ecdh_part.len());
            let ecdh_share_handle = session.derive_key(
                &Mechanism::Ecdh1Derive(ecdh1_derive_params),
                ecdh_info.private_key,
                attributes)?;
            t!("Derived ECDH share on token");

            // Extract the share from the token.
            let mut ecdh_share = None;
            for attribute in session.get_attributes(
                ecdh_share_handle,
                &[ AttributeType::Value ])?
            {
                if let Attribute::Value(v) = attribute {
                    ecdh_share = Some(v);
                } else {
                    return Err(anyhow::anyhow!(
                        "Token returned unrequested attribute: {:?}",
                        attribute));
                }
            }

            let Some(ecdh_share) = ecdh_share else {
                return Err(anyhow::anyhow!(
                    "Failed to extract ECDH share from token"));
            };
            t!("Extracted ECDH share from token, got {} bytes",
               ecdh_share.len());

            let kek = multi_key_combine(
                &mlkem_share,
                &ecdh_share,
                &ecdh_part[..],
                &ecdh_public[..],
                pk_algo)
                .context("Combining ML-KEM and ECDH key parts")?;

            Ok(aes_key_unwrap(SymmetricAlgorithm::AES256,
                              kek.as_protected(),
                              esk)
               .context("Unwrapping KEK")?.into())
        };

        match (key_info.key_type, self.openpgp_key.mpis(), ciphertext) {
            (KeyType::EC_MONTGOMERY,
             mpi::PublicKey::MLKEM768_X25519 {
                 ecdh: ecdh_public,
                 ..
             },
             mpi::Ciphertext::MLKEM768_X25519 {
                 ecdh: ecdh_part,
                 mlkem: mlkem_part,
                 esk,
             })
                | (KeyType::ML_KEM,
                   mpi::PublicKey::MLKEM768_X25519 {
                       ecdh: ecdh_public,
                       ..
                   },
                   mpi::Ciphertext::MLKEM768_X25519 {
                       ecdh: ecdh_part,
                       mlkem: mlkem_part,
                       esk,
                   }) =>
            {
                decrypt_mlkem(PublicKeyAlgorithm::MLKEM768_X25519,
                              key_info, composite_info,
                              &ecdh_public[..], &ecdh_part[..],
                              &mlkem_part[..], &esk[..])
            }

            (KeyType::EC_MONTGOMERY,
             mpi::PublicKey::MLKEM1024_X448 {
                 ecdh: ecdh_public,
                 ..
             },
             mpi::Ciphertext::MLKEM1024_X448 {
                 ecdh: ecdh_part,
                 mlkem: mlkem_part,
                 esk,
             })
                | (KeyType::ML_KEM,
                   mpi::PublicKey::MLKEM1024_X448 {
                       ecdh: ecdh_public,
                       ..
                   },
                   mpi::Ciphertext::MLKEM1024_X448 {
                       ecdh: ecdh_part,
                       mlkem: mlkem_part,
                       esk,
                   }) =>
            {
                decrypt_mlkem(PublicKeyAlgorithm::MLKEM1024_X448,
                              key_info, composite_info,
                              &ecdh_public[..], &ecdh_part[..],
                              &mlkem_part[..], &esk[..])
            }

            (KeyType::EC_MONTGOMERY,
             mpi::PublicKey::X25519 { u },
             mpi::Ciphertext::X25519 { e, key }) =>
            {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info,
                    KeyType::EC_MONTGOMERY)?;

                // Derive the shared key from our private key on the
                // token and the ephemeral key.

                // For the key derivation, we use PKCS#11's
                // CKM_ECDH1_DERIVE mechanism.  Per PKCS#11, v3.2,
                // draft 13, Section 6.3.17 Elliptic Curve
                // Diffie-Hellman key derivation:
                //
                //   The Elliptic Curve Diffie-Hellman (ECDH) key
                //   derivation mechanism, denoted CKM_ECDH1_DERIVE,
                //   is a mechanism for key derivation based on the
                //   Diffie-Hellman version of the Elliptic Curve key
                //   agreement scheme, as defined in [ANSI X9.63] for
                //   short Weierstrass EC keys and [RFC 7748] for
                //   Montgomery keys, where each party contributes one
                //   key pair all using the same EC domain parameters.

                let ecdh1_derive_params = Ecdh1DeriveParams::new(
                    EcKdf::null(), &e[..]);

                let attributes = &[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    // Don't persist the shared key on the token.
                    Attribute::Token(false),
                    // Make the shared key as extractable.
                    Attribute::Extractable(true),
                    Attribute::Sensitive(false),
                ];

                let shared_key_handle = session.derive_key(
                    &Mechanism::Ecdh1Derive(ecdh1_derive_params),
                    key_info.private_key,
                    attributes)?;
                t!("Derived shared key on token");

                // Extract the shared key from the token.
                let mut shared_key = None;
                for attribute in session.get_attributes(
                    shared_key_handle,
                    &[ AttributeType::Value ])?
                {
                    if let Attribute::Value(v) = attribute {
                        shared_key = Some(v);
                    } else {
                        return Err(anyhow::anyhow!(
                            "Token returned unrequested attribute: {:?}",
                            attribute));
                    }
                }

                let shared_key = if let Some(shared_key) = shared_key {
                    shared_key
                } else {
                    return Err(anyhow::anyhow!(
                        "Failed to export shared key form token"));
                };

                t!("Extracted shared key: {} bytes", shared_key.len());

                // Compute the wrap key.
                let wrap_algo = SymmetricAlgorithm::AES128;
                let mut ikm = vec![0; 32 + 32 + 32];

                // Yes clippy, this operation will always return
                // zero.  This is the intended outcome.  Chill.
                #[allow(clippy::erasing_op)]
                ikm[0 * 32..1 * 32].copy_from_slice(&e[..]);
                ikm[1 * 32..2 * 32].copy_from_slice(&u[..]);
                ikm[2 * 32..3 * 32].copy_from_slice(&shared_key[..]);

                // XXX: We should chain the operations so that
                // intermediate values don't leave the token.  But, it
                // seems that cryptoki doesn't expose
                // CKM_CONCATENATE_BASE_AND_KEY and
                // CKM_CONCATENATE_BASE_AND_DATA.  So, in the mean
                // time, we do the concatenation here.

                // Load the IKM.
                let ikm = session.create_object(&[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    Attribute::Token(false),
                    Attribute::Derive(true),
                    Attribute::Value(ikm.into()),
                ])?;

                // Prepare HKDF/SHA256.
                let hkdf_params = HkdfParams::new(MechanismType::SHA256_HMAC,
                                                  Some(HkdfSalt::Null),
                                                  Some(b"OpenPGP X25519"));

                let attributes = &[
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    // Don't persist the shared key on the token.
                    Attribute::Token(false),
                    // Make the shared key as extractable.
                    Attribute::Extractable(true),
                    Attribute::Sensitive(false),
                    Attribute::ValueLen(wrap_algo.key_size()?.try_into()?),
                ];

                let kek_handle = session.derive_key(
                    &Mechanism::HkdfDerive(hkdf_params),
                    ikm,
                    attributes)?;
                session.destroy_object(ikm)?;
                t!("Executed HKDF on the token");

                // XXX: We shouldn't extract, just chain aes unwrap.

                // Extract the OKM from the token.
                let mut kek = None;
                for attribute in session.get_attributes(
                    kek_handle,
                    &[ AttributeType::Value ])?
                {
                    if let Attribute::Value(v) = attribute {
                        kek = Some(v);
                    } else {
                        return Err(anyhow::anyhow!(
                            "Token returned unrequested attribute: {:?}",
                            attribute));
                    }
                }

                let kek = kek.map(SessionKey::from)
                    .ok_or(anyhow::anyhow!(
                        "Failed to export shared key form token"))?;

                // Decrypt the session key.
                Ok(aes_key_unwrap(wrap_algo, kek.as_protected(),
                                  key)?.into())
            },

            (KeyType::EC_MONTGOMERY,
             mpi::PublicKey::X448 { u },
             mpi::Ciphertext::X448 { e, key }) =>
            {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info,
                    KeyType::EC_MONTGOMERY)?;

                // Derive the shared key from our private key on the
                // token and the ephemeral key.

                // For the key derivation, we use PKCS#11's
                // CKM_ECDH1_DERIVE mechanism.  Per PKCS#11, v3.2,
                // draft 13, Section 6.3.17 Elliptic Curve
                // Diffie-Hellman key derivation:
                //
                //   The Elliptic Curve Diffie-Hellman (ECDH) key
                //   derivation mechanism, denoted CKM_ECDH1_DERIVE,
                //   is a mechanism for key derivation based on the
                //   Diffie-Hellman version of the Elliptic Curve key
                //   agreement scheme, as defined in [ANSI X9.63] for
                //   short Weierstrass EC keys and [RFC 7748] for
                //   Montgomery keys, where each party contributes one
                //   key pair all using the same EC domain parameters.

                let ecdh1_derive_params = Ecdh1DeriveParams::new(
                    EcKdf::null(), &e[..]);

                let attributes = &[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    // Don't persist the shared key on the token.
                    Attribute::Token(false),
                    // Make the shared key as extractable.
                    Attribute::Extractable(true),
                    Attribute::Sensitive(false),
                ];

                let shared_key_handle = session.derive_key(
                    &Mechanism::Ecdh1Derive(ecdh1_derive_params),
                    key_info.private_key,
                    attributes)?;
                t!("Derived shared key on token");

                // Extract the shared key from the token.
                let mut shared_key = None;
                for attribute in session.get_attributes(
                    shared_key_handle,
                    &[ AttributeType::Value ])?
                {
                    if let Attribute::Value(v) = attribute {
                        shared_key = Some(v);
                    } else {
                        return Err(anyhow::anyhow!(
                            "Token returned unrequested attribute: {:?}",
                            attribute));
                    }
                }

                let shared_key = if let Some(shared_key) = shared_key {
                    shared_key
                } else {
                    return Err(anyhow::anyhow!(
                        "Failed to export shared key form token"));
                };

                t!("Extracted shared key: {} bytes", shared_key.len());

                // Compute the wrap key.
                let wrap_algo = SymmetricAlgorithm::AES256;
                let mut ikm = vec![0; 56 + 56 + 56];

                // Yes clippy, this operation will always return
                // zero.  This is the intended outcome.  Chill.
                #[allow(clippy::erasing_op)]
                ikm[0 * 56..1 * 56].copy_from_slice(&e[..]);
                ikm[1 * 56..2 * 56].copy_from_slice(&u[..]);
                ikm[2 * 56..3 * 56].copy_from_slice(&shared_key[..]);

                // XXX: We should chain the operations so that
                // intermediate values don't leave the token.  But, it
                // seems that cryptoki doesn't expose
                // CKM_CONCATENATE_BASE_AND_KEY and
                // CKM_CONCATENATE_BASE_AND_DATA.  So, in the mean
                // time, we do the concatenation here.

                // Load the IKM.
                let ikm = session.create_object(&[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    Attribute::Token(false),
                    Attribute::Derive(true),
                    Attribute::Value(ikm.into()),
                ])?;

                // Prepare HKDF/SHA512.
                let hkdf_params = HkdfParams::new(MechanismType::SHA512_HMAC,
                                                  Some(HkdfSalt::Null),
                                                  Some(b"OpenPGP X448"));

                let attributes = &[
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    // Don't persist the shared key on the token.
                    Attribute::Token(false),
                    // Make the shared key as extractable.
                    Attribute::Extractable(true),
                    Attribute::Sensitive(false),
                    Attribute::ValueLen(wrap_algo.key_size()?.try_into()?),
                ];

                let kek_handle = session.derive_key(
                    &Mechanism::HkdfDerive(hkdf_params),
                    ikm,
                    attributes)?;
                session.destroy_object(ikm)?;
                t!("Executed HKDF on the token");

                // XXX: We shouldn't extract, just chain aes unwrap.

                // Extract the OKM from the token.
                let mut kek = None;
                for attribute in session.get_attributes(
                    kek_handle,
                    &[ AttributeType::Value ])?
                {
                    if let Attribute::Value(v) = attribute {
                        kek = Some(v);
                    } else {
                        return Err(anyhow::anyhow!(
                            "Token returned unrequested attribute: {:?}",
                            attribute));
                    }
                }

                let kek = kek.map(SessionKey::from)
                    .ok_or(anyhow::anyhow!(
                        "Failed to export shared key form token"))?;

                // Decrypt the session key.
                Ok(aes_key_unwrap(wrap_algo, kek.as_protected(),
                                  key)?.into())
            },

            (KeyType::EC_MONTGOMERY,
             mpi::PublicKey::ECDH { curve, hash, sym, .. },
             mpi::Ciphertext::ECDH { e, key, .. }) =>
            {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info,
                    KeyType::EC_MONTGOMERY)?;

                // 1.) Get the shared secret using the private key on
                // the token, 2.) derive the KEK from the shared
                // secret, 3.) use the KEK to decrypt the session key
                // bundle (session key plus metadata), then 4.)
                // extract the session key bundle.

                // The decryption algorithm is defined in RFC 9580,
                // Section 11.5. ECDH Algorithm:
                //
                //   This section describes the One-Pass
                //   Diffie-Hellman method, which is a combination of
                //   the ECC Diffie-Hellman method that establishes a
                //   shared secret and the key derivation method that
                //   processes the shared secret into a derived key.
                //
                //   https://www.rfc-editor.org/rfc/rfc9580.html#name-ecdh-algorithm

                // For 1) getting the shared secret, and 2) deriving
                // the KEK, we use PKCS#11's CKM_ECDH1_DERIVE
                // mechanism.  Per PKCS#11, v3.2, draft 13, Section
                // 6.3.17 Elliptic Curve Diffie-Hellman key
                // derivation:
                //
                //   The Elliptic Curve Diffie-Hellman (ECDH) key
                //   derivation mechanism, denoted CKM_ECDH1_DERIVE,
                //   is a mechanism for key derivation based on the
                //   Diffie-Hellman version of the Elliptic Curve key
                //   agreement scheme, as defined in [ANSI X9.63] for
                //   short Weierstrass EC keys and [RFC 7748] for
                //   Montgomery keys, where each party contributes one
                //   key pair all using the same EC domain parameters.

                // Configure the KDF.  According to
                //
                //   11.4. Key Derivation Function
                //
                //   A key derivation function (KDF) is necessary to
                //   implement EC encryption. The Concatenation Key
                //   Derivation Function (Approved Alternative 1)
                //   [SP800-56A] with the KDF hash function that is
                //   SHA2-256 [FIPS180] or stronger is REQUIRED.
                //
                //   https://www.rfc-editor.org/rfc/rfc9580.html#section-11.4
                let param = make_ecdh_derive_param(
                    self.public(), curve, hash, sym);
                let kdf = match hash {
                    HashAlgorithm::SHA256 => EcKdf::sha256_sp800(&param),
                    HashAlgorithm::SHA384 => EcKdf::sha384_sp800(&param),
                    HashAlgorithm::SHA512 => EcKdf::sha512_sp800(&param),
                    HashAlgorithm::SHA1 // Too weak.
                        | HashAlgorithm::SHA224 // Too weak.
                        | _ =>
                    {
                        return Err(anyhow::anyhow!(
                            "Can't decrypt message: using {} with the \
                             ECDH KDF is not supported",
                            hash));
                    }
                };

                let ephemeral_key = e.decode_point(&Curve::Cv25519)?.0;

                let ecdh1_derive_params = Ecdh1DeriveParams::new(
                    kdf, ephemeral_key);

                // The derived key's symmetric algorithm is encoded in
                // the public key.
                let sym_key_size = sym.key_size()
                    .with_context(|| {
                        format!("Symmetric algorithm {:?} not support",
                                sym)
                    })? as u64;

                let attributes = &[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::AES),
                    Attribute::ValueLen(Ulong::new(sym_key_size.try_into()?)),
                    // Don't persist the KEK on the token.
                    Attribute::Token(false),
                    // Mark the KEK as unwrap capable.
                    Attribute::Unwrap(true),
                ];

                t!("Deriving {:?} KEK using {:?} as the hash function",
                   sym, hash);
                let kek_handle = session.derive_key(
                    &Mechanism::Ecdh1Derive(ecdh1_derive_params),
                    key_info.private_key,
                    attributes)?;
                t!("Derived KEK on token");

                // We now 3.) use the KEK to decrypt the session key
                // bundle.
                //
                // From PKCS#11, v3.2, draft 13, Section
                // 6.16.3.17 AES Key Wrap:
                //
                //   The mechanisms support only single-part
                //   operations, i.e. single part wrapping and
                //   unwrapping, and single-part encryption and
                //   decryption.
                //
                //   • CKM_AES_KEY_WRAP
                //
                //   For unwrapping, the mechanism decrypts the
                //   wrapped key. In case of a secret key, it
                //   truncates the result according to the
                //   CKA_KEY_TYPE attribute of the template and, if it
                //   has one and the key type supports it, the
                //   CKA_VALUE_LEN attribute of the template. The
                //   length specified in the template must not be less
                //   than n-7 bytes, where n is the length of the
                //   wrapped key. In case of a private key, the
                //   mechanism parses the encoding as defined in
                //   section 6.7 and ignores trailing zero bytes.

                let attributes = &[
                    Attribute::Class(ObjectClass::SECRET_KEY),
                    Attribute::KeyType(KeyType::GENERIC_SECRET),
                    // Don't persist the session key on the token.
                    Attribute::Token(false),
                    // Mark the session key bundle as extractable.
                    Attribute::Extractable(true),
                    Attribute::Sensitive(false),
                ];

                t!("Unwrapping the session key bundle on the token");
                let session_key_handle = session.unwrap_key(
                    &Mechanism::AesKeyWrap,
                    kek_handle,
                    key,
                    attributes)
                    .context("Unwrapping session key")?;
                t!("Unwrapped the session key on the token");

                // Finally, we 4.) extract the session key bundle.

                let mut session_key_bundle = None;
                for attribute in session.get_attributes(
                    session_key_handle,
                    &[ AttributeType::Value ])?
                {
                    if let Attribute::Value(v) = attribute {
                        session_key_bundle = Some(v);
                    } else {
                        return Err(anyhow::anyhow!(
                            "Token returned unrequested attribute: {:?}",
                            attribute));
                    }
                }

                // The "session key" isn't just a session key.  For v3
                // PKESKs, we get:
                //
                //   symm_alg_ID || session key || checksum || pkcs5_padding
                //
                // (The algorithm ID is one byte and the checksum is
                // two bytes.)
                //
                // For v6 PKESKs, the symmetric algo is not included.
                //
                // For details, see
                // https://www.rfc-editor.org/rfc/rfc9580.html#section-11.5-20.4.1
                //
                // As hinted at above, this interface operates below
                // the OpenPGP level (we don't know if this is a v3
                // PKESK, a v6 PKESK, or something else entirely).  As
                // such, we only remove the padding, and treat the
                // rest as an opaque blob.  The caller does the rest.
                let session_key_bundle = if let Some(session_key_bundle)
                    = session_key_bundle
                {
                    session_key_bundle
                } else {
                    return Err(anyhow::anyhow!(
                        "Failed to export session key form token"));
                };

                let session_key_bundle = pkcs5_unpad(session_key_bundle)?;

                t!("Unwrapped {} bytes of session key", session_key_bundle.len());

                // The following may be useful for debugging, but
                // leaks the session key.
                //
                //t!("Returning {} bytes: {}",
                //   session_key_bundle.len(),
                //   session_key_bundle.iter()
                //       .map(|k| format!("{:02X}", k))
                //       .collect::<Vec<_>>()
                //       .join(""));

                Ok(session_key_bundle.into())
            },

            (KeyType::RSA,
             mpi::PublicKey::RSA { n, .. },
             mpi::Ciphertext::RSA { c }) =>
            {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info,
                    KeyType::RSA)?;

                let value = c.value_padded((n.bits() + 7) / 8)
                    // value_padded returns an error if the target
                    // size is smaller than the input.  This is quite
                    // unusual, but let's not fail here and instead
                    // just pass the ciphertext through as is and
                    // propagate the backend's error.
                    .unwrap_or_else(|_| c.value().into());

                let shared_key = session.decrypt(
                    &Mechanism::RsaPkcs, key_info.private_key, &value)?;

                Ok(shared_key.into())
            },

            (_, _, _) => {
                Err(anyhow::anyhow!("{} keys are not supported",
                                    self.openpgp_key.pk_algo()))
            },
        }
    }
}

// This is based on pkcs5_unpad from sequoia-openpgp:
//
// https://gitlab.com/sequoia-pgp/sequoia/-/blob/6773ab9a/openpgp/src/crypto/ecdh.rs#L201
/// Removes PKCS5 padding from a session key.
///
/// See [Section 8 of RFC 6637].
///
///   [Section 8 of RFC 6637]: https://tools.ietf.org/html/rfc6637#section-8
fn pkcs5_unpad(mut buf: Vec<u8>) -> Result<Protected> {
    if buf.len() > 0xff {
        return Err(anyhow::anyhow!("Can't unpad: message too large").into());
    }
    let mut good = true;
    let mut target_len = 0;

    let padding = buf[buf.len() - 1];
    if padding == 0 || padding as usize > buf.len() {
        // The padding can't be "0" or longer than the buffer.
        good = false
    } else {
        target_len = buf.len() - padding as usize;
        for &b in &buf[target_len..] {
            good = b == padding && good;
        }
    }
    if good {
        buf.truncate(target_len);
        Ok(buf.into())
    } else {
        let sk: Protected = buf.into();
        drop(sk);
        Err(anyhow::anyhow!("Can't unpad: bad padding").into())
    }
}

// Copied from sequoia-openpgp:
// https://gitlab.com/sequoia-pgp/sequoia/-/blob/6773ab9a/openpgp/src/crypto/ecdh.rs#L383
fn make_ecdh_derive_param<P, R>(recipient: &key::Key<P, R>,
                                curve: &Curve, hash: &HashAlgorithm,
                                sym: &SymmetricAlgorithm)
    -> Vec<u8>
    where P: key::KeyParts,
          R: key::KeyRole
{
    // Param = curve_OID_len || curve_OID ||
    // public_key_alg_ID || 03 || 01 || KDF_hash_ID ||
    // KEK_alg_ID for AESKeyWrap || "Anonymous Sender    " ||
    // recipient_fingerprint;
    let fp = recipient.fingerprint();
    let mut param = Vec::with_capacity(
        1 + curve.oid().len()        // Length and Curve OID,
            + 1                      // Public key algorithm ID,
            + 4                      // KDF parameters,
            + 20                     // "Anonymous Sender    ",
            + fp.as_bytes().len());  // Recipients key fingerprint.
    param.push(curve.oid().len() as u8);
    param.extend_from_slice(curve.oid());
    param.push(PublicKeyAlgorithm::ECDH.into());
    // KDF parameters.
    param.push(3); // Octet count of the following parameters.
    param.push(1); // 1-octet value 0x01, reserved for future extensions.
    param.push((*hash).into());
    param.push((*sym).into());
    param.extend_from_slice(b"Anonymous Sender    ");
    param.extend_from_slice(fp.as_bytes());
    assert_eq!(param.len(),
               1 + curve.oid().len()    // Length and Curve OID,
               + 1                      // Public key algorithm ID,
               + 4                      // KDF parameters,
               + 20                     // "Anonymous Sender    ",
               + fp.as_bytes().len());  // Recipients key fingerprint.

    param
}

/// Combines PQC and classical algorithms.
///
/// See [Section 4.2.1 of draft-ietf-openpgp-pqc-08].
///
/// [Section 4.2.1 of draft-ietf-openpgp-pqc-08]: https://www.ietf.org/archive/id/draft-ietf-openpgp-pqc-08.html#kem-key-combiner
// Copied from sequoia-openpgp:
// https://gitlab.com/sequoia-pgp/sequoia/-/blob/03595533/openpgp/src/crypto/asymmetric.rs#L480
pub(crate) fn multi_key_combine(mlkem_key: &[u8],
                                ecdh_key: &[u8],
                                ecdh_ciphertext: &[u8],
                                ecdh_public: &[u8],
                                pk_algo: PublicKeyAlgorithm)
                                -> Result<SessionKey>
{
    //   multiKeyCombine(
    //       mlkemKeyShare, ecdhKeyShare,
    //       ecdhCipherText, ecdhPublicKey,
    //       algId
    //   )
    //
    //   Input:
    //   mlkemKeyShare   - the ML-KEM key share encoded as an octet string
    //   ecdhKeyShare    - the ECDH key share encoded as an octet string
    //   ecdhCipherText  - the ECDH ciphertext encoded as an octet string
    //   ecdhPublicKey   - the ECDH public key of the recipient as an octet string
    //   algId           - the OpenPGP algorithm ID of the public-key encryption algorithm
    //
    // KEK = SHA3-256(
    //           mlkemKeyShare || ecdhKeyShare ||
    //           ecdhCipherText || ecdhPublicKey ||
    //           algId || domSep || len(domSep)
    //       )

    let mut hash = HashAlgorithm::SHA3_256.context()?.for_digest();
    hash.update(mlkem_key);
    hash.update(ecdh_key);
    hash.update(ecdh_ciphertext);
    hash.update(ecdh_public);
    hash.update(&[pk_algo.into()]);
    // Domain separation and length octet.
    hash.update(b"OpenPGPCompositeKDFv1\x15");

    let mut kek = SessionKey::from(vec![0; 32]);
    hash.digest(&mut kek)?;
    Ok(kek)
}
