add cipher suite option in mqtt builder

Author: xiazhvera

awsiot/mqtt5_client_builder.py

@@ -362,7 +362,8 @@ def _builder(
     return client
 
 
-def mtls_from_path(cert_filepath, pri_key_filepath, **kwargs) -> awscrt.mqtt5.Client:
+def mtls_from_path(cert_filepath, pri_key_filepath, cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
+                   **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an mTLS MQTT5 Client to AWS IoT.
     TLS arguments are passed as filepaths.
@@ -377,10 +378,15 @@ def mtls_from_path(cert_filepath, pri_key_filepath, **kwargs) -> awscrt.mqtt5.Cl
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_from_path(cert_filepath, pri_key_filepath)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
-def mtls_from_bytes(cert_bytes, pri_key_bytes, **kwargs) -> awscrt.mqtt5.Client:
+def mtls_from_bytes(
+        cert_bytes,
+        pri_key_bytes,
+        cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
+        **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an mTLS MQTT5 Client to AWS IoT.
     TLS arguments are passed as in-memory bytes.
@@ -395,6 +401,7 @@ def mtls_from_bytes(cert_bytes, pri_key_bytes, **kwargs) -> awscrt.mqtt5.Client:
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls(cert_bytes, pri_key_bytes)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
@@ -406,6 +413,7 @@ def mtls_with_pkcs11(*,
                      private_key_label: str = None,
                      cert_filepath: str = None,
                      cert_bytes=None,
+                     cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
                      **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an mTLS MQTT connection to AWS IoT,
@@ -451,11 +459,14 @@ def mtls_with_pkcs11(*,
         private_key_label=private_key_label,
         cert_file_path=cert_filepath,
         cert_file_contents=cert_bytes)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
+
 def mtls_with_pkcs12(*,
                      pkcs12_filepath: str,
                      pkcs12_password: str,
+                     cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
                      **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT,
@@ -476,11 +487,13 @@ def mtls_with_pkcs12(*,
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_pkcs12(
         pkcs12_filepath=pkcs12_filepath,
         pkcs12_password=pkcs12_password)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
 def mtls_with_windows_cert_store_path(*,
                                       cert_store_path: str,
+                                      cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
                                       **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an mTLS MQTT5 Client to AWS IoT,
@@ -499,13 +512,15 @@ def mtls_with_windows_cert_store_path(*,
     _check_required_kwargs(**kwargs)
 
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_windows_cert_store_path(cert_store_path)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
 def websockets_with_default_aws_signing(
         region,
         credentials_provider,
         websocket_proxy_options=None,
+        cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an MQTT5 Client over websockets to AWS IoT.
@@ -543,12 +558,17 @@ def _sign_websocket_handshake_request(transform_args, **kwargs):
         except Exception as e:
             transform_args.set_done(e)
 
-    return websockets_with_custom_handshake(_sign_websocket_handshake_request, websocket_proxy_options, **kwargs)
+    return websockets_with_custom_handshake(
+        _sign_websocket_handshake_request,
+        websocket_proxy_options,
+        cipher_suite,
+        **kwargs)
 
 
 def websockets_with_custom_handshake(
         websocket_handshake_transform,
         websocket_proxy_options=None,
+        cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an MQTT5 Client over websockets,
@@ -576,6 +596,7 @@ def websockets_with_custom_handshake(
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions()
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options=tls_ctx_options,
                     use_websockets=True,
                     websocket_handshake_transform=websocket_handshake_transform,
@@ -607,6 +628,7 @@ def direct_with_custom_authorizer(
         auth_password=None,
         auth_token_key_name=None,
         auth_token_value=None,
+        cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an MQTT5 Client using a custom
@@ -673,6 +695,7 @@ def direct_with_custom_authorizer(
 
     tls_ctx_options = awscrt.io.TlsContextOptions()
     tls_ctx_options.alpn_list = ["mqtt"]
+    tls_ctx_options.cipher_pref = cipher_suite
 
     return _builder(tls_ctx_options=tls_ctx_options,
                     use_websockets=False,
@@ -688,6 +711,7 @@ def websockets_with_custom_authorizer(
         websocket_proxy_options=None,
         auth_token_key_name=None,
         auth_token_value=None,
+        cipher_suite=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, configured for an MQTT5 Client using a custom
@@ -757,6 +781,7 @@ def websockets_with_custom_authorizer(
     kwargs["password"] = auth_password
 
     tls_ctx_options = awscrt.io.TlsContextOptions()
+    tls_ctx_options.cipher_pref = cipher_suite
 
     def _sign_websocket_handshake_request(transform_args, **kwargs):
         # transform_args need to know when transform is done
@@ -773,7 +798,7 @@ def _sign_websocket_handshake_request(transform_args, **kwargs):
                     **kwargs)
 
 
-def new_default_builder(**kwargs) -> awscrt.mqtt5.Client:
+def new_default_builder(cipher_suite=awscrt.io.TlsCipherPref.DEFAULT, **kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, without any configuration besides the default TLS context options.
 
@@ -782,6 +807,7 @@ def new_default_builder(**kwargs) -> awscrt.mqtt5.Client:
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions()
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options=tls_ctx_options,
                     use_websockets=False,
                     **kwargs)

awsiot/mqtt_connection_builder.py

@@ -258,7 +258,11 @@ def _builder(
     )
 
 
-def mtls_from_path(cert_filepath, pri_key_filepath, **kwargs) -> awscrt.mqtt.Connection:
+def mtls_from_path(
+        cert_filepath,
+        pri_key_filepath,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
+        **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT.
     TLS arguments are passed as filepaths.
@@ -273,10 +277,15 @@ def mtls_from_path(cert_filepath, pri_key_filepath, **kwargs) -> awscrt.mqtt.Con
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_from_path(cert_filepath, pri_key_filepath)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
-def mtls_from_bytes(cert_bytes, pri_key_bytes, **kwargs) -> awscrt.mqtt.Connection:
+def mtls_from_bytes(
+        cert_bytes,
+        pri_key_bytes,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
+        **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT.
     TLS arguments are passed as in-memory bytes.
@@ -291,6 +300,7 @@ def mtls_from_bytes(cert_bytes, pri_key_bytes, **kwargs) -> awscrt.mqtt.Connecti
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls(cert_bytes, pri_key_bytes)
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
@@ -302,6 +312,7 @@ def mtls_with_pkcs11(*,
                      private_key_label: str = None,
                      cert_filepath: str = None,
                      cert_bytes=None,
+                     cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
                      **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT,
@@ -347,12 +358,15 @@ def mtls_with_pkcs11(*,
         private_key_label=private_key_label,
         cert_file_path=cert_filepath,
         cert_file_contents=cert_bytes)
+    tls_ctx_options.cipher_pref = cipher_suite
 
     return _builder(tls_ctx_options, **kwargs)
 
+
 def mtls_with_pkcs12(*,
                      pkcs12_filepath: str,
                      pkcs12_password: str,
+                     cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
                      **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT,
@@ -373,11 +387,13 @@ def mtls_with_pkcs12(*,
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_pkcs12(
         pkcs12_filepath=pkcs12_filepath,
         pkcs12_password=pkcs12_password)
+    tls_ctx_options.cipher_suite = cipher_suite
     return _builder(tls_ctx_options, **kwargs)
 
 
 def mtls_with_windows_cert_store_path(*,
                                       cert_store_path: str,
+                                      cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
                                       **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an mTLS MQTT connection to AWS IoT,
@@ -396,6 +412,7 @@ def mtls_with_windows_cert_store_path(*,
     _check_required_kwargs(**kwargs)
 
     tls_ctx_options = awscrt.io.TlsContextOptions.create_client_with_mtls_windows_cert_store_path(cert_store_path)
+    tls_ctx_options.cipher_pref = cipher_suite
 
     return _builder(tls_ctx_options, **kwargs)
 
@@ -404,6 +421,7 @@ def websockets_with_default_aws_signing(
         region,
         credentials_provider,
         websocket_proxy_options=None,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection over websockets to AWS IoT.
@@ -441,12 +459,17 @@ def _sign_websocket_handshake_request(transform_args, **kwargs):
         except Exception as e:
             transform_args.set_done(e)
 
-    return websockets_with_custom_handshake(_sign_websocket_handshake_request, websocket_proxy_options, **kwargs)
+    return websockets_with_custom_handshake(
+        _sign_websocket_handshake_request,
+        cipher_suite,
+        websocket_proxy_options,
+        **kwargs)
 
 
 def websockets_with_custom_handshake(
         websocket_handshake_transform,
         websocket_proxy_options=None,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection over websockets,
@@ -474,6 +497,7 @@ def websockets_with_custom_handshake(
     """
     _check_required_kwargs(**kwargs)
     tls_ctx_options = awscrt.io.TlsContextOptions()
+    tls_ctx_options.cipher_pref = cipher_suite
     return _builder(tls_ctx_options=tls_ctx_options,
                     use_websockets=True,
                     websocket_handshake_transform=websocket_handshake_transform,
@@ -505,6 +529,7 @@ def direct_with_custom_authorizer(
         auth_password=None,
         auth_token_key_name=None,
         auth_token_value=None,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection using a custom
@@ -550,8 +575,10 @@ def direct_with_custom_authorizer(
         auth_token_key_name=auth_token_key_name,
         auth_token_value=auth_token_value,
         use_websockets=False,
+        cipher_suite: awscrt.io.TlsCipherPref=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs)
 
+
 def websockets_with_custom_authorizer(
         region=None,
         credentials_provider=None,
@@ -561,6 +588,7 @@ def websockets_with_custom_authorizer(
         auth_password=None,
         auth_token_key_name=None,
         auth_token_value=None,
+        cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection using a custom
@@ -590,7 +618,7 @@ def websockets_with_custom_authorizer(
         auth_authorizer_signature (`str`):  The digital signature of the token value in the `auth_token_value`
             parameter. The signature must be based on the private key associated with the custom authorizer.  The
             signature must be base64 encoded.
-            Required if the custom authorizer has signing enabled.  
+            Required if the custom authorizer has signing enabled.
 
         auth_token_key_name (`str`): Key used to extract the custom authorizer token from MQTT username query-string
             properties.
@@ -612,19 +640,21 @@ def websockets_with_custom_authorizer(
         use_websockets=True,
         websockets_region=region,
         websockets_credentials_provider=credentials_provider,
+        cipher_suite: awscrt.io.TlsCipherPref=awscrt.io.TlsCipherPref.DEFAULT,
         **kwargs)
 
 
 def _with_custom_authorizer(auth_username=None,
-        auth_authorizer_name=None,
-        auth_authorizer_signature=None,
-        auth_password=None,
-        auth_token_key_name=None,
-        auth_token_value=None,
-        use_websockets=False,
-        websockets_credentials_provider=None,
-        websockets_region=None,
-        **kwargs) -> awscrt.mqtt.Connection:
+                            auth_authorizer_name=None,
+                            auth_authorizer_signature=None,
+                            auth_password=None,
+                            auth_token_key_name=None,
+                            auth_token_value=None,
+                            use_websockets=False,
+                            websockets_credentials_provider=None,
+                            websockets_region=None,
+                            cipher_suite: awscrt.io.TlsCipherPref = awscrt.io.TlsCipherPref.DEFAULT,
+                            **kwargs) -> awscrt.mqtt.Connection:
     """
     Helper function that contains the setup needed for custom authorizers
     """
@@ -657,7 +687,8 @@ def _with_custom_authorizer(auth_username=None,
     kwargs["password"] = auth_password
 
     tls_ctx_options = awscrt.io.TlsContextOptions()
-    if use_websockets == False:
+    tls_ctx_options.cipher_pref = cipher_suite
+    if not use_websockets:
         kwargs["port"] = 443
         tls_ctx_options.alpn_list = ["mqtt"]