"""
This module provides PKCS#11 integration for pyHanko, by providing a wrapper
for `python-pkcs11 <https://github.com/danni/python-pkcs11>`_ that can be
seamlessly plugged into a :class:`~.signers.PdfSigner`.
"""

import asyncio
import binascii
import logging
import warnings
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import pkcs11
from asn1crypto import algos, core, x509
from asn1crypto.algos import RSASSAPSSParams
from cryptography.hazmat.primitives import hashes
from pyhanko.config.pkcs11 import (
    PKCS11PinEntryMode,
    PKCS11SignatureConfig,
    PKCS11SigningPinEntryMode,
    TokenCriteria,
)
from pyhanko.pdf_utils.misc import coalesce
from pyhanko.sign.general import SigningError
from pyhanko.sign.signers import Signer
from pyhanko_certvalidator.registry import CertificateStore
from pyhanko_certvalidator.util import get_pyca_cryptography_hash

try:
    from pkcs11 import (
        MGF,
        PROTECTED_AUTH,
        Attribute,
        Mechanism,
        ObjectClass,
        PKCS11Error,
        Session,
        SignMixin,
    )
    from pkcs11 import lib as p11_lib
    from pkcs11 import types as p11_types
except ImportError as e:  # pragma: nocover
    raise ImportError(
        "pyhanko.sign.pkcs11 requires pyHanko to be installed with "
        "the [pkcs11] option. You can install missing "
        "dependencies by running \"pip install 'pyHanko[pkcs11]'\".",
        e,
    )

__all__ = [
    'PKCS11Signer',
    'PKCS11SigningContext',
    'find_token',
    'open_pkcs11_session',
    'select_pkcs11_signing_params',
]

logger = logging.getLogger(__name__)


def criteria_mismatches(
    criteria: TokenCriteria | None, token: p11_types.Token
) -> list[tuple[str, str]]:
    if criteria is None:
        return []

    err_items = []

    if criteria.label is not None and token.label != criteria.label:
        err_items.append(('label', criteria.label))
    if criteria.serial is not None and token.serial != criteria.serial:
        err_items.append(('serial', criteria.serial.hex()))
    return err_items


def criteria_satisfied_by(
    criteria: TokenCriteria | None, token: p11_types.Token
) -> bool:
    return not criteria_mismatches(criteria, token)


def find_token(
    slots: list[p11_types.Slot],
    slot_no: int | None = None,
    token_criteria: TokenCriteria | None = None,
) -> p11_types.Token | None:
    """
    Internal helper method to find a token.

    :param slots:
        The list of slots.
    :param slot_no:
        Slot number to use. If not specified, the first slot containing a token
        satisfying the criteria will be used
    :param token_criteria:
        Criteria the token must satisfy.
    :return:
        A PKCS#11 token object, or ``None`` if none was found.
    """

    if token_criteria is None and slot_no is None:
        if len(slots) == 1:
            return slots[0].get_token()
        else:
            raise PKCS11Error(
                "Module has more than 1 slot; slot index or token criteria "
                "must be provided"
            )

    token = None
    if slot_no is None:
        for slot in slots:
            try:
                cur_token = slot.get_token()
                if criteria_satisfied_by(token_criteria, cur_token):
                    token = cur_token
                    break
            except PKCS11Error:
                continue
    else:
        if slot_no >= len(slots):
            raise PKCS11Error(
                f"Slot index {slot_no} too large; there are only {len(slots)}"
            )
        cur_token = slots[slot_no].get_token()
        errors = criteria_mismatches(token_criteria, cur_token)
        if errors:
            err_str = ", ".join(
                f"{field} is not {val!r}" for field, val in errors
            )
            raise PKCS11Error(
                f"Token in slot {slot_no} does not satisfy criteria; {err_str}."
            )
        else:
            token = cur_token
    return token


@dataclass(frozen=True)
class PKCS11SignatureOperationSpec:
    """
    Internal helper class to describe how to invoke a signature operation on
    a key in a PKCS #11 token.
    """

    sign_kwargs: dict[str, Any]
    """
    Keyword arguments to the ``sign`` function on the key handle.
    """

    pre_sign_transform: Callable[[bytes], bytes] | None
    """
    An optional transformation to apply to the data prior to signing.
    """

    post_sign_transform: Callable[[bytes], bytes] | None
    """
    An optional transformation to apply to the data after signing.
    """


RSA_MECH_MAP = {
    'sha1': Mechanism.SHA1_RSA_PKCS,
    'sha224': Mechanism.SHA224_RSA_PKCS,
    'sha256': Mechanism.SHA256_RSA_PKCS,
    'sha384': Mechanism.SHA384_RSA_PKCS,
    'sha512': Mechanism.SHA512_RSA_PKCS,
}

RSASSA_PSS_MECH_MAP = {
    'sha1': Mechanism.SHA1_RSA_PKCS_PSS,
    'sha224': Mechanism.SHA224_RSA_PKCS_PSS,
    'sha256': Mechanism.SHA256_RSA_PKCS_PSS,
    'sha384': Mechanism.SHA384_RSA_PKCS_PSS,
    'sha512': Mechanism.SHA512_RSA_PKCS_PSS,
}

MGF_MECH_MAP = {
    'sha1': MGF.SHA1,
    'sha224': MGF.SHA224,
    'sha256': MGF.SHA256,
    'sha384': MGF.SHA384,
    'sha512': MGF.SHA512,
}

ECDSA_MECH_MAP = {
    'sha1': Mechanism.ECDSA_SHA1,
    'sha224': Mechanism.ECDSA_SHA224,
    'sha256': Mechanism.ECDSA_SHA256,
    'sha384': Mechanism.ECDSA_SHA384,
    'sha512': Mechanism.ECDSA_SHA512,
}

DSA_MECH_MAP = {
    'sha1': Mechanism.DSA_SHA1,
    'sha224': Mechanism.DSA_SHA224,
    'sha256': Mechanism.DSA_SHA256,
    # These can't be used in CMS IIRC (since the key sizes required
    # to meaningfully use them are ridiculous),
    # but they're in the PKCS#11 spec, so let's add them for
    # completeness
    'sha384': Mechanism.DSA_SHA384,
    'sha512': Mechanism.DSA_SHA512,
}

DIGEST_MECH_MAP = {
    'sha1': Mechanism.SHA_1,
    'sha224': Mechanism.SHA224,
    'sha256': Mechanism.SHA256,
    'sha384': Mechanism.SHA384,
    'sha512': Mechanism.SHA512,
}


def select_pkcs11_signing_params(
    signature_mechanism: algos.SignedDigestAlgorithm,
    digest_algorithm: str,
    use_raw_mechanism: bool,
    sign_kwargs: dict[str, Any],
) -> PKCS11SignatureOperationSpec:
    """
    Internal helper function to set up a PKCS #11 signing operation.

    :param signature_mechanism:
        The signature mechanism to use (as an ASN.1 value)
    :param digest_algorithm:
        The digest algorithm to use
    :param use_raw_mechanism:
        Whether to attempt to use the raw mechanism on pre-hashed data.
    :param sign_kwargs:
        Override kwargs manually.
    :return:
    """
    from pkcs11.util.dsa import encode_dsa_signature
    from pkcs11.util.ec import encode_ecdsa_signature

    pre_sign_transform = None
    post_sign_transform = None
    kwargs: dict[str, Any] = {}

    try:
        signature_algo = signature_mechanism.signature_algo
    except ValueError:
        signature_algo = signature_mechanism['algorithm'].native

    if signature_algo == 'rsassa_pkcs1v15':
        if use_raw_mechanism:
            kwargs['mechanism'] = Mechanism.RSA_PKCS
            pre_sign_transform = _hash_fully(
                digest_algorithm, wrap_digest_info=True
            )
        else:
            kwargs['mechanism'] = RSA_MECH_MAP[digest_algorithm]
    elif signature_algo == 'dsa':
        if use_raw_mechanism:
            kwargs['mechanism'] = Mechanism.DSA
            pre_sign_transform = _hash_fully(
                digest_algorithm, wrap_digest_info=False
            )
        else:
            kwargs['mechanism'] = DSA_MECH_MAP[digest_algorithm]
        post_sign_transform = encode_dsa_signature
    elif signature_algo == 'ecdsa':
        if use_raw_mechanism:
            kwargs['mechanism'] = Mechanism.ECDSA
            pre_sign_transform = _hash_fully(
                digest_algorithm, wrap_digest_info=False
            )
        else:
            # TODO test these (unsupported in SoftHSMv2 right now)
            kwargs['mechanism'] = ECDSA_MECH_MAP[digest_algorithm]
        post_sign_transform = encode_ecdsa_signature
    elif signature_algo == 'rsassa_pss':
        if use_raw_mechanism:
            raise NotImplementedError("RSASSA-PSS not available in raw mode")
        params: RSASSAPSSParams = signature_mechanism['parameters']
        assert digest_algorithm == params['hash_algorithm']['algorithm'].native

        # unpack PSS parameters into PKCS#11 language
        kwargs['mechanism'] = RSASSA_PSS_MECH_MAP[digest_algorithm]

        pss_digest_param = DIGEST_MECH_MAP[digest_algorithm]

        mgf_val = params['mask_gen_algorithm']['parameters']['algorithm'].native
        pss_mgf_param = MGF_MECH_MAP[mgf_val]
        pss_salt_len = params['salt_length'].native

        kwargs['mechanism_param'] = (
            pss_digest_param,
            pss_mgf_param,
            pss_salt_len,
        )
    elif signature_algo == 'ed25519':
        if use_raw_mechanism:
            # Note: Ed25519-ph isn't the same thing.
            raise NotImplementedError("Ed25519 not available in raw mode")
        kwargs['mechanism'] = Mechanism.EDDSA
    elif signature_algo == 'ed448':
        if use_raw_mechanism:
            # Note: Ed448-ph isn't the same thing.
            raise NotImplementedError("Ed448 not available in raw mode")
        kwargs['mechanism'] = Mechanism.EDDSA
        kwargs['mechanism_param'] = (False, None)
    elif signature_algo in ('mldsa44', 'mldsa65', 'mldsa87'):
        if use_raw_mechanism:
            raise NotImplementedError("ML-DSA not available in raw mode")
        # the key object's parameter set determines which ML_DSA variant we use
        kwargs['mechanism'] = Mechanism.ML_DSA
    elif 'mechanism' not in sign_kwargs:
        raise NotImplementedError(
            f"Signature algorithm '{signature_algo}' is not supported."
        )
    kwargs.update(**sign_kwargs)
    return PKCS11SignatureOperationSpec(
        sign_kwargs=kwargs,
        pre_sign_transform=pre_sign_transform,
        post_sign_transform=post_sign_transform,
    )


def open_pkcs11_session(
    lib_location: str,
    slot_no: int | None = None,
    token_label: str | None = None,
    token_criteria: TokenCriteria | None = None,
    user_pin: str | object | None = None,
) -> Session:
    """
    Open a PKCS#11 session

    :param lib_location:
        Path to the PKCS#11 module.
    :param slot_no:
        Slot number to use. If not specified, the first slot containing a token
        labelled ``token_label`` will be used.
    :param token_label:
        .. deprecated:: 0.14.0
            Use ``token_criteria`` instead.

        Label of the token to use. If ``None``, there is no constraint.
    :param token_criteria:
        Criteria that the token should match.
    :param user_pin:
        User PIN to use, or :attr:`.PROTECTED_AUTH`. If ``None``, authentication
        is skipped.

        .. note::
            Some PKCS#11 implementations do not require PIN when the token
            is opened, but will prompt for it out-of-band when signing.
            Whether :attr:`.PROTECTED_AUTH` or ``None`` is used in this case
            depends on the implementation.
    :return:
        An open PKCS#11 session object.
    """
    lib = p11_lib(lib_location)

    if token_criteria is None and token_label is not None:
        warnings.warn(
            "'token_label' is deprecated, use 'token_criteria' instead",
            DeprecationWarning,
        )
        token_criteria = TokenCriteria(label=token_label)

    slots = lib.get_slots()
    token = find_token(slots, slot_no=slot_no, token_criteria=token_criteria)
    if token is None:
        raise PKCS11Error(
            f'No token matching criteria {token_criteria!r} found'
            if token_criteria is not None
            else 'No token found'
        )

    kwargs = {}
    if user_pin is not None:
        kwargs['user_pin'] = user_pin

    return token.open(**kwargs)  # type: ignore


def _format_pull_err_msg(
    object_kind: str,
    no_results: bool,
    label: str | None = None,
    id_value: bytes | None = None,
):
    info_strs = []
    if label is not None:
        info_strs.append(f"label '{label}'")
    if id_value is not None:
        info_strs.append(f"ID '{binascii.hexlify(id_value).decode('ascii')}'")
    qualifier = f" with {', '.join(info_strs)}" if info_strs else ""
    if no_results:
        err = f"Could not find {object_kind}{qualifier}."
    else:
        err = f"Found more than one {object_kind}{qualifier}."
    return err


def _hash_fully(digest_algorithm: str, *, wrap_digest_info: bool):
    md_spec = get_pyca_cryptography_hash(digest_algorithm)

    def _h(data: bytes) -> bytes:
        h = hashes.Hash(md_spec)
        h.update(data)
        digest = h.finalize()
        if wrap_digest_info:
            return algos.DigestInfo(
                {
                    'digest_algorithm': {
                        'algorithm': digest_algorithm.lower(),
                        'parameters': core.Null(),
                    },
                    'digest': digest,
                }
            ).dump()
        else:
            return digest

    return _h


# TODO: perhaps attempt automatic key discovery if the labels aren't provided?


class PKCS11Signer(Signer):
    """
    Signer implementation for PKCS11 devices.

    :param pkcs11_session:
        The PKCS11 session object to use.
    :param cert_label:
        The label of the certificate that will be used for signing, to
        be pulled from the PKCS#11 token.
    :param cert_id:
        ID of the certificate object that will be used for signing, to
        be pulled from the PKCS#11 token.
    :param signing_cert:
        The signer's certificate. If the signer's certificate is provided via
        this parameter, the ``cert_label`` and ``cert_id`` parameters will not
        be used to retrieve the signer's certificate.
    :param ca_chain:
        Set of other relevant certificates
        (as :class:`.asn1crypto.x509.Certificate` objects).
    :param key_label:
        The label of the key that will be used for signing.
        Defaults to the value of ``cert_label`` if left unspecified and
        ``key_id`` is also unspecified.

        .. note::
            At least one of ``key_id``, ``key_label`` and ``cert_label`` must
            be supplied.
    :param key_id:
        ID of the private key object (optional).
    :param other_certs_to_pull:
        List labels of other certificates to pull from the PKCS#11 device.
        Defaults to the empty tuple. If ``None``, pull *all* certificates.
    :param bulk_fetch:
        Boolean indicating the fetching strategy.
        If ``True``, fetch all certs and filter the unneeded ones.
        If ``False``, fetch the requested certs one by one.
        Default value is ``True``, unless ``other_certs_to_pull`` has one or
        fewer elements, in which case it is always treated as ``False``.
    :param use_raw_mechanism:
        Use the 'raw' equivalent of the selected signature mechanism. This is
        useful when working with tokens that do not support a hash-then-sign
        mode of operation.

        .. note::
            This functionality is only available for ECDSA at this time.
            Support for other signature schemes will be added on an as-needed
            basis.
    :param signature_mechanism:
        Specify the signature mechanism explicitly. By default, pyHanko will
        attempt to infer it from the key's properties.
    :param base_sign_kwargs:
        Directly specify some of the keyword arguments passed to the
        :meth:`pkcs11.SignMixin.sign` method.

        This can be used as an escape hatch to enable functionality
        (advanced auth, additional signing mechanisms) that are not
        directly supported by pyHanko.
    """

    default_cert_query_params: dict[Attribute, Any]
    default_key_query_params: dict[Attribute, Any]

    def __init__(
        self,
        pkcs11_session: Session,
        cert_label: str | None = None,
        signing_cert: x509.Certificate | None = None,
        ca_chain=None,
        key_label: str | None = None,
        prefer_pss=False,
        embed_roots=True,
        other_certs_to_pull=(),
        bulk_fetch=True,
        key_id: bytes | None = None,
        cert_id: bytes | None = None,
        use_raw_mechanism=False,
        signature_mechanism: algos.SignedDigestAlgorithm | None = None,
        base_sign_kwargs: dict[str, Any] | None = None,
    ):
        """
        Initialise a PKCS11 signer.
        """
        self.cert_label = coalesce(
            cert_label, key_label if not cert_id else None
        )
        self.key_id = coalesce(key_id, cert_id if not key_label else None)
        self.cert_id = coalesce(cert_id, key_id if not cert_label else None)
        self.key_label = coalesce(key_label, cert_label if not key_id else None)
        self.pkcs11_session = pkcs11_session
        self.other_certs = other_certs_to_pull
        self._other_certs_loaded = False
        if other_certs_to_pull is not None and len(other_certs_to_pull) <= 1:
            self.bulk_fetch = False
        else:
            self.bulk_fetch = bulk_fetch
        self.use_raw_mechanism = use_raw_mechanism
        self._key_handle: SignMixin | None = None
        self._loaded = False
        self._sign_kwargs = base_sign_kwargs or {}
        self.__loading_event = None
        super().__init__(
            prefer_pss=prefer_pss,
            embed_roots=embed_roots,
            signing_cert=signing_cert,
            signature_mechanism=signature_mechanism,
        )
        if ca_chain:
            self._cert_registry.register_multiple(ca_chain)
        if signing_cert is not None:
            self._cert_registry.register(signing_cert)
        self.default_cert_query_params = {
            Attribute.CLASS: ObjectClass.CERTIFICATE
        }
        self.default_key_query_params = {
            Attribute.CLASS: ObjectClass.PRIVATE_KEY,
            Attribute.SIGN: True,
        }

    def _init_cert_registry(self):
        # it's conceivable that one might want to load this separately from
        # the key data, so we allow for that.
        if not self._other_certs_loaded:
            certs = self._load_other_certs()
            self._cert_registry.register_multiple(certs)
            self._other_certs_loaded = True
        return self._cert_registry

    @property
    def cert_registry(self) -> CertificateStore:
        # apparently mypy doesn't like it when I write
        # cert_registry = property(_init_cert_registry)
        return self._init_cert_registry()

    @property
    def signing_cert(self):
        self._load_objects()
        return self._signing_cert

    # noinspection PyUnusedLocal
    def sign_kwargs(self, data_to_sign: bytes) -> dict[str, Any]:
        """
        Supply keyword arguments to pass to the :meth:`pkcs11.SignMixin.sign`
        method.

        This can be used as an escape hatch to enable functionality
        (advanced auth, additional signing mechanisms) that are not
        directly supported by pyHanko.

        :param data_to_sign:
            The data being signed.

            .. note::
                In typical cases this parameter can be ignored;
                pyHanko exposes it to subclasses to allow for
                just-in-time credentials that depend on
                the content being signed.

        :return:
            A dictionary of keyword arguments.
        """
        return self._sign_kwargs

    def _select_pkcs11_signing_params(
        self, digest_algorithm: str, sign_kwargs: dict[str, Any]
    ) -> PKCS11SignatureOperationSpec:
        digest_algorithm = digest_algorithm.lower()
        return select_pkcs11_signing_params(
            self.get_signature_mechanism_for_digest(digest_algorithm),
            digest_algorithm,
            use_raw_mechanism=self.use_raw_mechanism,
            sign_kwargs=sign_kwargs,
        )

    def estimate_raw_signature_size_bytes(self) -> int:
        algo = self.signing_cert.public_key.algorithm
        return {
            'mldsa44': 2420,
            'mldsa65': 3309,
            'mldsa87': 4627,
        }.get(algo, 512)

    async def async_sign_raw(
        self, data: bytes, digest_algorithm: str, dry_run=False
    ) -> bytes:
        if dry_run:
            return self.estimate_raw_signature_size_bytes() * b'0'

        await self.ensure_objects_loaded()
        from pkcs11 import SignMixin

        assert self._key_handle is not None
        kh: SignMixin = self._key_handle
        spec = self._select_pkcs11_signing_params(
            digest_algorithm, sign_kwargs=self.sign_kwargs(data)
        )

        if spec.pre_sign_transform is not None:
            data = spec.pre_sign_transform(data)

        def _perform_signature():
            signature = kh.sign(data, **spec.sign_kwargs)
            if spec.post_sign_transform is not None:
                signature = spec.post_sign_transform(signature)
            return signature

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, _perform_signature)

    def _load_other_certs(self) -> set[x509.Certificate]:
        return set(self.__pull())

    def bulk_cert_fetch_params(self) -> dict[Attribute, Any]:
        """
        Search parameters to bulk fetch certificates from the token.

        .. note:: Can be overridden in subclasses.
        """
        return dict(self.default_cert_query_params)

    def single_cert_fetch_params(
        self,
        label: str | None,
        cert_id: bytes | None,
    ) -> dict[Attribute, Any]:
        """
        Search parameters to fetch a single certificate from the token.

        .. note:: Can be overridden in subclasses.
        """

        query_params: dict[Attribute, Any] = dict(
            self.default_cert_query_params
        )
        if label is not None:
            query_params[Attribute.LABEL] = label
        if cert_id is not None:
            query_params[Attribute.ID] = cert_id
        return query_params

    def signing_key_fetch_params(self) -> dict[Attribute, Any]:
        """
        Search parameters to fetch the signing key from the token.

        .. note:: Can be overridden in subclasses.
        """
        query_params: dict[Attribute, Any] = dict(self.default_key_query_params)
        if self.key_label is not None:
            query_params[Attribute.LABEL] = self.key_label
        if self.key_id is not None:
            query_params[Attribute.ID] = self.key_id

        return query_params

    def _pull_single_cert(
        self,
        label: str | None,
        cert_id: bytes | None,
    ):
        query_params = self.single_cert_fetch_params(label, cert_id)
        q = self.pkcs11_session.get_objects(query_params)

        # need to run through the full iterator to make sure the operation
        # terminates
        results = list(q)
        if len(results) == 1:
            cert_obj = results[0]
            return x509.Certificate.load(cert_obj[Attribute.VALUE])
        else:
            err = _format_pull_err_msg(
                "certificate",
                no_results=not results,
                label=label,
                id_value=cert_id,
            )
            raise PKCS11Error(err)

    def _pull_signing_key_handle(self):
        query_params = self.signing_key_fetch_params()
        q = self.pkcs11_session.get_objects(query_params)

        results = list(q)
        if len(results) == 1:
            return results[0]
        else:
            err = _format_pull_err_msg(
                "private key",
                no_results=not results,
                label=self.key_label,
                id_value=self.key_id,
            )
            raise PKCS11Error(err)

    def __pull(self):
        other_cert_labels = self.other_certs
        if other_cert_labels is not None and len(other_cert_labels) == 0:
            # if there's nothing to fetch, bail.
            # Recall: None -> fetch everything, so we check the length
            # explicitly
            return
        if other_cert_labels is None or self.bulk_fetch:
            # first, query all certs
            q = self.pkcs11_session.get_objects(self.bulk_cert_fetch_params())
            logger.debug("Pulling all certificates from PKCS#11 token...")
            for cert_obj in q:
                label = cert_obj[Attribute.LABEL]
                if other_cert_labels is None or label in other_cert_labels:
                    # LGTM believes that we're logging sensitive info here, but
                    # I politely disagree: this is just the PKCS#11 label of the
                    # certificate
                    msg = f"Found certificate with label '{label}' on token."
                    logger.debug(msg)  # lgtm
                    yield x509.Certificate.load(cert_obj[Attribute.VALUE])
        else:
            # fetch certs one by one
            for label in other_cert_labels:
                # LGTM believes that we're logging sensitive info here, but
                # I politely disagree: this is just the PKCS#11 label of the
                # certificate
                msg = (
                    f"Pulling certificate with label '{label}' from "
                    f"PKCS#11 token..."
                )
                logger.debug(msg)  # lgtm
                yield self._pull_single_cert(label=label, cert_id=None)

    async def ensure_objects_loaded(self):
        """
        Async method that, when awaited, ensures that objects
        (relevant certificates, key handles, ...) are loaded.

        This coroutine is guaranteed to be called & awaited in :meth:`sign_raw`,
        but some property implementations may cause object loading to be
        triggered synchronously (for backwards compatibility reasons).
        This blocks the event loop the first time it happens.

        To avoid this behaviour, asynchronous code should ideally perform
        `await signer.ensure_objects_loaded()` after instantiating the signer.

        .. note::
            The asynchronous context manager on :class:`PKCS11SigningContext`
            takes care of that automatically.
        """

        if self._loaded:
            return
        if self.__loading_event is None:
            self.__loading_event = event = asyncio.Event()
            loop = asyncio.get_running_loop()
            await loop.run_in_executor(None, self._load_objects)
            event.set()
        else:  # pragma: nocover
            # some other coroutine is dealing with fetching already,
            # just wait for that one to finish
            await self.__loading_event.wait()

    def _load_objects(self):
        if self._loaded:
            return

        self._init_cert_registry()
        if self._signing_cert is None:
            self._signing_cert = self._pull_single_cert(
                label=self.cert_label, cert_id=self.cert_id
            )

        kh = self._pull_signing_key_handle()
        self._key_handle = kh

        self._loaded = True


class PKCS11SigningContext:
    """Context manager for PKCS#11 configurations."""

    def __init__(
        self, config: PKCS11SignatureConfig, user_pin: str | None = None
    ):
        self.config = config
        self._session: Session | None = None
        self._user_pin = user_pin

    def _handle_pin(self):
        pin = self._user_pin or self.config.user_pin
        # mode 'PROMPT' is irrelevant for library usage.
        if pin is not None:
            pin = str(pin)
        elif self.config.prompt_pin == PKCS11PinEntryMode.DEFER:
            pin = PROTECTED_AUTH
        else:
            pin = None
        return pin

    def _instantiate(self) -> PKCS11Signer:
        config = self.config
        pin = self._handle_pin()

        try:
            self._session = session = open_pkcs11_session(
                config.module_path,
                slot_no=config.slot_no,
                token_criteria=config.token_criteria,
                user_pin=pin,
            )
        except pkcs11.PKCS11Error as ex:
            raise SigningError(
                f"PKCS#11 error while opening session to {config.module_path}: [{type(ex).__name__}] {ex}"
            ) from ex
        base_signing_kwargs = {}
        if config.signing_pin_mode == PKCS11SigningPinEntryMode.REAUTHENTICATE:
            base_signing_kwargs["pin"] = config.signing_pin or pin
        signer = PKCS11Signer(
            session,
            config.cert_label,
            ca_chain=config.other_certs,
            key_label=config.key_label,
            prefer_pss=config.prefer_pss,
            use_raw_mechanism=config.raw_mechanism,
            other_certs_to_pull=config.other_certs_to_pull,
            bulk_fetch=config.bulk_fetch,
            key_id=config.key_id,
            cert_id=config.cert_id,
            signing_cert=config.signing_certificate,
            signature_mechanism=config.signature_mechanism,
            base_sign_kwargs=base_signing_kwargs,
        )
        if config.only_resident_certs:
            signer.default_cert_query_params[Attribute.TOKEN] = True
        return signer

    def __enter__(self):
        return self._instantiate()

    async def __aenter__(self):
        loop = asyncio.get_running_loop()
        signer = await loop.run_in_executor(None, self._instantiate)
        await signer.ensure_objects_loaded()
        return signer

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._session.close()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self._session.close()
