refactor: replace from awsiot-credentialhelper[security]

pyproject.toml

@@ -31,7 +31,7 @@ dynamic = [
 ]
 dependencies = [
   "aiohttp>=3.10.11,<3.14",
-  "awsiot-credentialhelper>=0.6,<1.1",
+  "awscrt>=0.16.9,<0.32",
   "boto3>=1.34.35,<1.43",
   "grpcio>=1.70,<1.79",
   "protobuf>=4.25.8,<6.34",

requirements.txt

@@ -21,29 +21,19 @@ attrs==25.3.0 ; python_full_version < '3.9'
 attrs==25.4.0 ; python_full_version >= '3.9'
     # via aiohttp
 awscrt==0.16.26 ; python_full_version < '3.8.1'
-    # via awsiot-credentialhelper
-awscrt==0.31.1 ; python_full_version >= '3.8.1'
-    # via awsiot-credentialhelper
-awsiot-credentialhelper==0.6.0 ; python_full_version < '3.8.1'
     # via otaclient-iot-logging-server
-awsiot-credentialhelper==1.0.2 ; python_full_version >= '3.8.1'
+awscrt==0.31.1 ; python_full_version >= '3.8.1'
     # via otaclient-iot-logging-server
 boto3==1.37.38 ; python_full_version < '3.9'
-    # via
-    #   awsiot-credentialhelper
-    #   otaclient-iot-logging-server
+    # via otaclient-iot-logging-server
 boto3==1.42.44 ; python_full_version >= '3.9'
-    # via
-    #   awsiot-credentialhelper
-    #   otaclient-iot-logging-server
+    # via otaclient-iot-logging-server
 botocore==1.37.38 ; python_full_version < '3.9'
     # via
-    #   awsiot-credentialhelper
     #   boto3
     #   s3transfer
 botocore==1.42.44 ; python_full_version >= '3.9'
     # via
-    #   awsiot-credentialhelper
     #   boto3
     #   s3transfer
 frozenlist==1.5.0 ; python_full_version < '3.9'
@@ -121,15 +111,13 @@ six==1.17.0
 typing-extensions==4.13.2 ; python_full_version < '3.9'
     # via
     #   annotated-types
-    #   awsiot-credentialhelper
     #   multidict
     #   otaclient-iot-logging-server
     #   pydantic
     #   pydantic-core
 typing-extensions==4.15.0 ; python_full_version >= '3.9'
     # via
     #   aiosignal
-    #   awsiot-credentialhelper
     #   grpcio
     #   multidict
     #   otaclient-iot-logging-server
@@ -140,10 +128,10 @@ typing-inspection==0.4.2 ; python_full_version >= '3.9'
     # via
     #   pydantic
     #   pydantic-settings
-urllib3==1.26.20
-    # via
-    #   awsiot-credentialhelper
-    #   botocore
+urllib3==1.26.20 ; python_full_version < '3.10'
+    # via botocore
+urllib3==2.6.3 ; python_full_version >= '3.10'
+    # via botocore
 yarl==1.15.2 ; python_full_version < '3.9'
     # via aiohttp
 yarl==1.22.0 ; python_full_version >= '3.9'

src/otaclient_iot_logging_server/boto3_session.py

@@ -15,21 +15,37 @@
 
 from __future__ import annotations
 
+import json
+import logging
 import ssl
 import subprocess
+from http import HTTPStatus
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
-from awsiot_credentialhelper.boto3_session import Boto3SessionProvider
-from awsiot_credentialhelper.boto3_session import Pkcs11Config as aws_PKcs11Config
+from awscrt.http import HttpClientConnection, HttpRequest
+from awscrt.io import (
+    ClientBootstrap,
+    ClientTlsContext,
+    DefaultHostResolver,
+    EventLoopGroup,
+    Pkcs11Lib,
+    TlsContextOptions,
+)
 from boto3 import Session
+from botocore.credentials import DeferredRefreshableCredentials
+from botocore.session import get_session as get_botocore_session
 
 from otaclient_iot_logging_server._utils import parse_pkcs11_uri
 from otaclient_iot_logging_server.greengrass_config import (
     IoTSessionConfig,
     PKCS11Config,
 )
 
+logger = logging.getLogger(__name__)
+
+_AWSCRT_TIMEOUT_SEC = 10
+
 #
 # ------ certificate loading helpers ------ #
 #
@@ -71,7 +87,7 @@ def _convert_to_pem(_data: bytes) -> bytes:
 
 def _load_certificate(cert_path: str, pkcs11_cfg: Optional[PKCS11Config]) -> bytes:
     """
-    NOTE: Boto3SessionProvider only takes PEM format cert.
+    NOTE: Only PEM format cert is supported.
     """
     if cert_path.startswith("pkcs11"):
         assert pkcs11_cfg, (
@@ -91,20 +107,200 @@ def _load_certificate(cert_path: str, pkcs11_cfg: Optional[PKCS11Config]) -> byt
     return _convert_to_pem(Path(cert_path).read_bytes())
 
 
+#
+# ------ credential fetching ------ #
+#
+# AWS IoT Core Credential Provider API:
+#   https://docs.aws.amazon.com/iot/latest/developerguide/authorizing-direct-aws.html
+
+
+def _parse_credentials_response(
+    response_status: int,
+    response_body: bytes,
+    credential_url: str,
+) -> Dict[str, Any]:
+    """Parse credential provider response, raising on non-200 status."""
+    if response_status != HTTPStatus.OK:
+        logger.error(
+            "Failed to get credentials from %s: status=%s, body=%s",
+            credential_url,
+            response_status,
+            response_body.decode(),
+        )
+        raise ValueError(
+            f"Error getting credentials from IoT credential provider: "
+            f"status={response_status}"
+        )
+    credentials = json.loads(response_body.decode())["credentials"]
+    return {
+        "access_key": credentials["accessKeyId"],
+        "secret_key": credentials["secretAccessKey"],
+        "token": credentials["sessionToken"],
+        "expiry_time": credentials["expiration"],
+    }
+
+
+def _build_tls_context_from_path(
+    cert_path: str,
+    key_path: str,
+) -> TlsContextOptions:
+    """Build TLS context options using plain certificate/key files."""
+    return TlsContextOptions.create_client_with_mtls_from_path(
+        cert_filepath=cert_path,
+        pk_filepath=key_path,
+    )
+
+
+def _build_tls_context_pkcs11(
+    cert_pem: bytes,
+    pkcs11_cfg: PKCS11Config,
+    private_key_label: Optional[str] = None,
+) -> TlsContextOptions:
+    """Build TLS context options using PKCS#11 for private key operations."""
+    return TlsContextOptions.create_client_with_mtls_pkcs11(
+        pkcs11_lib=Pkcs11Lib(file=pkcs11_cfg.pkcs11_lib),
+        user_pin=pkcs11_cfg.user_pin,
+        slot_id=int(pkcs11_cfg.slot_id),
+        token_label=None,  # type: ignore[arg-type]
+        private_key_label=private_key_label,  # type: ignore[arg-type]
+        cert_file_path=None,  # type: ignore[arg-type]
+        cert_file_contents=cert_pem,
+    )
+
+
+def _fetch_iot_credentials_via_awscrt(
+    endpoint: str,
+    role_alias: str,
+    thing_name: str,
+    tls_ctx_opt: TlsContextOptions,
+) -> Dict[str, Any]:
+    """Fetch IAM credentials from AWS IoT Core Credential Provider via mTLS.
+
+    Uses awscrt for the HTTP request with the given TLS context.
+
+    Args:
+        endpoint: AWS IoT credential provider endpoint FQDN.
+        role_alias: IoT Role Alias name.
+        thing_name: IoT Thing Name.
+        tls_ctx_opt: TLS context options configured for mTLS.
+
+    Returns:
+        Credentials dict with access_key, secret_key, token, expiry_time.
+    """
+    path = f"/role-aliases/{role_alias}/credentials"
+    url = f"https://{endpoint}{path}"
+
+    event_loop_group = EventLoopGroup()
+    host_resolver = DefaultHostResolver(event_loop_group)
+    bootstrap = ClientBootstrap(event_loop_group, host_resolver)
+
+    tls_ctx = ClientTlsContext(tls_ctx_opt)
+    tls_conn_opt = tls_ctx.new_connection_options()
+    tls_conn_opt.set_server_name(endpoint)
+
+    connection = HttpClientConnection.new(
+        host_name=endpoint,
+        port=443,
+        bootstrap=bootstrap,
+        tls_connection_options=tls_conn_opt,
+    ).result(_AWSCRT_TIMEOUT_SEC)
+
+    request = HttpRequest("GET", path)
+    request.headers.add("host", endpoint)
+    request.headers.add("x-amzn-iot-thingname", thing_name)
+
+    response_status_code: int = 0
+    response_body = bytearray()
+
+    def on_response(
+        http_stream: Any, status_code: int, headers: Any, **_kwargs: Any
+    ) -> None:
+        nonlocal response_status_code
+        response_status_code = status_code
+
+    def on_body(http_stream: Any, chunk: bytes, **_kwargs: Any) -> None:
+        response_body.extend(chunk)
+
+    stream = connection.request(request, on_response, on_body)
+    stream.activate()
+    stream.completion_future.result(_AWSCRT_TIMEOUT_SEC)
+
+    return _parse_credentials_response(response_status_code, bytes(response_body), url)
+
+
+def _fetch_iot_credentials(
+    endpoint: str,
+    role_alias: str,
+    thing_name: str,
+    cert_path: str,
+    key_path: str,
+) -> Dict[str, Any]:
+    """Fetch IAM credentials using plain certificate/key files."""
+    tls_ctx_opt = _build_tls_context_from_path(cert_path, key_path)
+    return _fetch_iot_credentials_via_awscrt(
+        endpoint=endpoint,
+        role_alias=role_alias,
+        thing_name=thing_name,
+        tls_ctx_opt=tls_ctx_opt,
+    )
+
+
+def _fetch_iot_credentials_pkcs11(
+    endpoint: str,
+    role_alias: str,
+    thing_name: str,
+    cert_pem: bytes,
+    pkcs11_cfg: PKCS11Config,
+    private_key_label: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Fetch IAM credentials using PKCS#11 for private key operations."""
+    tls_ctx_opt = _build_tls_context_pkcs11(cert_pem, pkcs11_cfg, private_key_label)
+    return _fetch_iot_credentials_via_awscrt(
+        endpoint=endpoint,
+        role_alias=role_alias,
+        thing_name=thing_name,
+        tls_ctx_opt=tls_ctx_opt,
+    )
+
+
 #
 # ------ session creating helpers ------ #
 #
 
 
+def _create_boto3_session(region: str, refresh_func: Any) -> Session:
+    """Create a boto3 Session with auto-refreshing credentials.
+
+    Args:
+        region: AWS region name.
+        refresh_func: Callable that returns credentials dict.
+
+    Returns:
+        boto3 Session with refreshable credentials.
+    """
+    botocore_session = get_botocore_session()
+    botocore_session._credentials = DeferredRefreshableCredentials(  # type: ignore[attr-defined]
+        method="custom-iot-core-credential-provider",
+        refresh_using=refresh_func,
+    )
+    botocore_session.set_config_variable("region", region)
+
+    return Session(botocore_session=botocore_session)
+
+
 def _get_session(config: IoTSessionConfig) -> Session:
     """Get a session that using plain privkey."""
-    return Boto3SessionProvider(
-        endpoint=config.aws_credential_provider_endpoint,
-        role_alias=config.aws_role_alias,
-        certificate=_load_certificate(config.certificate_path, config.pkcs11_config),
-        private_key=config.private_key_path,
-        thing_name=config.thing_name,
-    ).get_session()  # type: ignore
+
+    def _refresh():
+        return _fetch_iot_credentials(
+            endpoint=config.aws_credential_provider_endpoint,
+            role_alias=config.aws_role_alias,
+            thing_name=config.thing_name,
+            cert_path=config.certificate_path,
+            key_path=config.private_key_path,
+        )
+
+    return _create_boto3_session(config.region, _refresh)
 
 
 def _get_session_pkcs11(config: IoTSessionConfig) -> Session:
@@ -113,19 +309,20 @@ def _get_session_pkcs11(config: IoTSessionConfig) -> Session:
         "privkey is provided by pkcs11, but pkcs11_config is not available"
     )
 
+    cert_pem = _load_certificate(config.certificate_path, config.pkcs11_config)
     _parsed_key_uri = parse_pkcs11_uri(config.private_key_path)
-    return Boto3SessionProvider(
-        endpoint=config.aws_credential_provider_endpoint,
-        role_alias=config.aws_role_alias,
-        certificate=_load_certificate(config.certificate_path, config.pkcs11_config),
-        thing_name=config.thing_name,
-        pkcs11=aws_PKcs11Config(
-            pkcs11_lib=pkcs11_cfg.pkcs11_lib,
-            slot_id=int(pkcs11_cfg.slot_id),
-            user_pin=pkcs11_cfg.user_pin,
+
+    def _refresh():
+        return _fetch_iot_credentials_pkcs11(
+            endpoint=config.aws_credential_provider_endpoint,
+            role_alias=config.aws_role_alias,
+            thing_name=config.thing_name,
+            cert_pem=cert_pem,
+            pkcs11_cfg=pkcs11_cfg,
             private_key_label=_parsed_key_uri.get("object"),
-        ),
-    ).get_session()  # type: ignore
+        )
+
+    return _create_boto3_session(config.region, _refresh)
 
 
 # API