[minor] Adding mqtt method to IOTX509RefreshableSession (#80)

* mqtt method

* adding logging

* Docs update

* readme

* Adding missing dependency.

* readme changes update

* final readme update

Author: michaelthomasletts

README.md

@@ -99,6 +99,7 @@
   - **STS**
   - **IoT Core** 
     - X.509 certificates w/ role aliases over mTLS (PEM files and PKCS#11)
+    - MQTT actions are available!
   - Custom authentication methods
 - Natively supports all parameters supported by `boto3.session.Session`
 - [Tested](https://github.com/michaelthomasletts/boto3-refresh-session/tree/main/tests), [documented](https://michaelthomasletts.github.io/boto3-refresh-session/index.html), and [published to PyPI](https://pypi.org/project/boto3-refresh-session/)
@@ -317,6 +318,8 @@ pip install boto3-refresh-session
 
 ## ⚠️ Changes
 
+Browse through the various changes to `boto3-refresh-session` over time.
+
 #### 😥 v3.0.0
 
 **The changes introduced by v3.0.0 will not impact ~99% of users** who generally interact with `boto3-refresh-session` by only `RefreshableSession`, *which is the intended usage for this package after all.* 
@@ -331,4 +334,8 @@ The `ecs` module has been dropped. For additional details and rationale, please
 
 #### 😛 v5.0.0
 
-Support for IoT Core via X.509 certificate-based authentication (over HTTPS) is now available!
+Support for IoT Core via X.509 certificate-based authentication (over HTTPS) is now available!
+
+#### ➕ v5.1.0
+
+MQTT support added for IoT Core via X.509 certificate-based authentication.

boto3_refresh_session/methods/iot/x509.py

@@ -2,32 +2,42 @@
 
 import json
 import re
+from atexit import register
 from pathlib import Path
-from typing import cast
+from tempfile import NamedTemporaryFile
+from typing import cast, get_args
 from urllib.parse import ParseResult, urlparse
 
+from awscrt import auth, io
 from awscrt.exceptions import AwsCrtError
 from awscrt.http import HttpClientConnection, HttpRequest
 from awscrt.io import (
     ClientBootstrap,
     ClientTlsContext,
     DefaultHostResolver,
     EventLoopGroup,
+    LogLevel,
     Pkcs11Lib,
     TlsConnectionOptions,
     TlsContextOptions,
+    init_logging,
 )
+from awscrt.mqtt import Connection
+from awsiot import mqtt_connection_builder
 
 from ...exceptions import BRSError, BRSWarning
 from ...utils import (
     PKCS11,
     AWSCRTResponse,
     Identity,
     TemporaryCredentials,
+    Transport,
     refreshable_session,
 )
 from .core import BaseIoTRefreshableSession
 
+_TEMP_PATHS: list[str] = []
+
 
 @refreshable_session
 class IOTX509RefreshableSession(
@@ -71,6 +81,9 @@ class IOTX509RefreshableSession(
         The duration for which the temporary credentials are valid, in
         seconds. Cannot exceed the value declared in the IAM policy.
         Default is None.
+    awscrt_log_level : awscrt.LogLevel | None, optional
+        The logging level for the AWS CRT library, e.g.
+        ``awscrt.LogLevel.INFO``. Default is None.
 
     Notes
     -----
@@ -90,34 +103,34 @@ def __init__(
         verify_peer: bool = True,
         timeout: float | int | None = None,
         duration_seconds: int | None = None,
+        awscrt_log_level: LogLevel | None = None,
         **kwargs,
     ):
         # initializing BRSSession
         super().__init__(refresh_method="iot-x509", **kwargs)
 
+        # logging
+        if awscrt_log_level:
+            init_logging(log_level=awscrt_log_level, file_name="stdout")
+
         # initializing public attributes
         self.endpoint = self._normalize_iot_credential_endpoint(
             endpoint=endpoint
         )
         self.role_alias = role_alias
-        self.certificate = certificate
+        self.certificate = self._read_maybe_path_to_bytes(
+            certificate, fallback=None, name="certificate"
+        )
         self.thing_name = thing_name
-        self.private_key = private_key
-        self.pkcs11 = pkcs11
-        self.ca = ca
+        self.private_key = self._read_maybe_path_to_bytes(
+            private_key, fallback=None, name="private_key"
+        )
+        self.pkcs11 = self._validate_pkcs11(pkcs11) if pkcs11 else None
+        self.ca = self._read_maybe_path_to_bytes(ca, fallback=None, name="ca")
         self.verify_peer = verify_peer
         self.timeout = 10.0 if timeout is None else timeout
         self.duration_seconds = duration_seconds
 
-        # loading X.509 certificate if presented as a string, which
-        # is presumed to be the file path.
-        # if presented as bytes then self.certificate is presumed to be
-        # the actual certificate itself
-        if self.certificate and isinstance(self.certificate, str):
-            self.certificate = (
-                Path(self.certificate).expanduser().resolve().read_bytes()
-            )
-
         # either private_key or pkcs11 must be provided
         if self.private_key is None and self.pkcs11 is None:
             raise BRSError(
@@ -130,22 +143,6 @@ def __init__(
                 "Only one of 'private_key' or 'pkcs11' can be provided."
             )
 
-        # if the provided private_key is bytes then it's presumed to be
-        # the actual private key. but if it's string then it's presumed
-        # to be the file path
-        if self.private_key and isinstance(self.private_key, str):
-            self.private_key = (
-                Path(self.private_key).expanduser().resolve().read_bytes()
-            )
-
-        # verifying PKCS#11 dict
-        if self.pkcs11:
-            self.pkcs11 = self._validate_pkcs11(pkcs11=self.pkcs11)
-
-        # ca is like many other attributes in that str implies file location
-        if self.ca and isinstance(self.ca, str):
-            self.ca = Path(self.ca).expanduser().resolve().read_bytes()
-
     def _get_credentials(self) -> TemporaryCredentials:
         url = urlparse(
             f"https://{self.endpoint}/role-aliases/{self.role_alias}"
@@ -334,3 +331,208 @@ def _validate_pkcs11(pkcs11: PKCS11) -> PKCS11:
         pkcs11.setdefault("token_label", None)
         pkcs11.setdefault("private_key_label", None)
         return pkcs11
+
+    @staticmethod
+    def _read_maybe_path_to_bytes(
+        v: str | bytes | None, fallback: bytes | None, name: str
+    ) -> bytes | None:
+        match v:
+            case None:
+                return fallback
+            case bytes():
+                return v
+            case str() as p if Path(p).expanduser().resolve().is_file():
+                return Path(p).expanduser().resolve().read_bytes()
+            case _:
+                raise BRSError(f"Invalid {name} provided.")
+
+    @staticmethod
+    def _bytes_to_tempfile(b: bytes, suffix: str = ".pem") -> str:
+        f = NamedTemporaryFile("wb", suffix=suffix, delete=False)
+        f.write(b)
+        f.flush()
+        f.close()
+        _TEMP_PATHS.append(f.name)
+        return f.name
+
+    @staticmethod
+    @register
+    def _cleanup_tempfiles():
+        for p in _TEMP_PATHS:
+            try:
+                Path(p).unlink(missing_ok=True)
+            except Exception:
+                ...
+
+    def mqtt(
+        self,
+        *,
+        endpoint: str,
+        client_id: str,
+        transport: Transport = "x509",
+        certificate: str | bytes | None = None,
+        private_key: str | bytes | None = None,
+        ca: str | bytes | None = None,
+        pkcs11: PKCS11 | None = None,
+        region: str | None = None,
+        keep_alive_secs: int = 60,
+        clean_start: bool = True,
+        port: int | None = None,
+        use_alpn: bool = False,
+    ) -> Connection:
+        """Establishes an MQTT connection using the specified parameters.
+
+        Parameters
+        ----------
+        endpoint: str
+            The MQTT endpoint to connect to.
+        client_id: str
+            The client ID to use for the MQTT connection.
+        transport: Transport
+            The transport protocol to use (e.g., "x509" or "ws").
+        certificate: str | bytes | None
+            The client certificate to use for the connection. Defaults to the
+            session certificate.
+        private_key: str | bytes | None
+            The private key to use for the connection. Defaults to the
+            session private key.
+        ca: str | bytes | None
+            The CA certificate to use for the connection. Defaults to the
+            session CA certificate.
+        pkcs11: PKCS11 | None
+            PKCS#11 configuration for hardware-backed keys. Defaults to the
+            session PKCS#11 configuration.
+        region: str | None
+            The AWS region to use for the connection. Defaults to the
+            session region.
+        keep_alive_secs: int
+            The keep-alive interval for the MQTT connection. Default is 60
+            seconds.
+        clean_start: bool
+            Whether to start a clean session. Default is True.
+        port: int | None
+            The port to use for the MQTT connection. Default is 8883 if not
+            using ALPN, otherwise 443.
+        use_alpn: bool
+            Whether to use ALPN for the connection. Default is False.
+
+        Returns
+        -------
+        awscrt.mqtt.Connection
+            The established MQTT connection.
+        """
+
+        # Validate transport
+        if transport not in list(get_args(Transport)):
+            raise BRSError("Transport must be 'x509' or 'ws'")
+
+        # Region default (WS only)
+        if region is None:
+            region = self.region_name
+
+        # Normalize inputs to bytes using session defaults
+        cert_bytes = self._read_maybe_path_to_bytes(
+            certificate, getattr(self, "certificate", None), "certificate"
+        )
+        key_bytes = self._read_maybe_path_to_bytes(
+            private_key, getattr(self, "private_key", None), "private_key"
+        )
+        ca_bytes = self._read_maybe_path_to_bytes(
+            ca, getattr(self, "ca", None), "ca"
+        )
+
+        # Validate PKCS#11
+        match pkcs11:
+            case None:
+                pkcs11 = getattr(self, "pkcs11", None)
+            case dict():
+                pkcs11 = self._validate_pkcs11(pkcs11)
+            case _:
+                raise BRSError("Invalid PKCS#11 configuration provided.")
+
+        # X.509 invariants
+        if transport == "x509":
+            has_key = key_bytes is not None
+            has_hsm = pkcs11 is not None
+            if not has_key and not has_hsm:
+                raise BRSError(
+                    "For transport='x509', provide either 'private_key' "
+                    "(bytes/path) or 'pkcs11'."
+                )
+            if has_key and has_hsm:
+                raise BRSError(
+                    "Provide only one of 'private_key' or 'pkcs11' for "
+                    "transport='x509'."
+                )
+            if cert_bytes is None:
+                raise BRSError("Certificate is required for transport='x509'")
+
+        # CRT bootstrap
+        event_loop = io.EventLoopGroup(1)
+        host_resolver = io.DefaultHostResolver(event_loop)
+        bootstrap = io.ClientBootstrap(event_loop, host_resolver)
+
+        # Build connection
+        if transport == "x509":
+            if pkcs11 is not None:
+                # Cert must be a filepath for PKCS#11 builder → write temp
+                cert_path = self._bytes_to_tempfile(
+                    cast(bytes, cert_bytes), ".crt"
+                )
+                ca_path = (
+                    self._bytes_to_tempfile(ca_bytes, ".pem")
+                    if ca_bytes
+                    else None
+                )
+
+                return mqtt_connection_builder.mtls_with_pkcs11(
+                    endpoint=endpoint,
+                    client_bootstrap=bootstrap,
+                    pkcs11_lib=Pkcs11Lib(file=pkcs11["pkcs11_lib"]),
+                    user_pin=pkcs11.get("user_pin"),
+                    slot_id=pkcs11.get("slot_id"),
+                    token_label=pkcs11.get("token_label"),
+                    private_key_object=pkcs11.get("private_key_label"),
+                    cert_filepath=cert_path,
+                    ca_filepath=ca_path,
+                    client_id=client_id,
+                    clean_session=clean_start,
+                    keep_alive_secs=keep_alive_secs,
+                    port=port or (443 if use_alpn else 8883),
+                    alpn_list=["x-amzn-mqtt-ca"] if use_alpn else None,
+                )
+            else:
+                # pure mTLS with in-memory cert/key/CA
+                return mqtt_connection_builder.mtls_from_bytes(
+                    endpoint=endpoint,
+                    cert_bytes=cert_bytes,
+                    pri_key_bytes=key_bytes,
+                    ca_bytes=ca_bytes,
+                    client_bootstrap=bootstrap,
+                    client_id=client_id,
+                    clean_session=clean_start,
+                    keep_alive_secs=keep_alive_secs,
+                    port=port or (443 if use_alpn else 8883),
+                    alpn_list=["x-amzn-mqtt-ca"] if use_alpn else None,
+                )
+
+        else:  # transport == "ws"
+            # WebSockets + SigV4
+            creds_provider = auth.AwsCredentialsProvider.new_delegate(
+                self._credentials
+            )
+            ca_path = (
+                self._bytes_to_tempfile(ca_bytes, ".pem") if ca_bytes else None
+            )
+
+            return mqtt_connection_builder.websockets_with_default_aws_signing(
+                endpoint=endpoint,
+                client_bootstrap=bootstrap,
+                region=region,
+                credentials_provider=creds_provider,
+                client_id=client_id,
+                clean_session=clean_start,
+                keep_alive_secs=keep_alive_secs,
+                ca_filepath=ca_path,
+                port=port or 443,
+            )

boto3_refresh_session/utils/typing.py

@@ -13,6 +13,7 @@
     "STSClientParams",
     "TemporaryCredentials",
     "RefreshableTemporaryCredentials",
+    "Transport",
 ]
 
 from datetime import datetime
@@ -57,6 +58,9 @@
 #: Type alias for values returned by get_identity
 Identity: TypeAlias = dict[str, Any]
 
+#: Type alias for acceptable transports
+Transport: TypeAlias = Literal["x509", "ws"]
+
 
 class TemporaryCredentials(TypedDict):
     """Temporary IAM credentials."""