Use ssl context wrap socket in Python3.12+ (#338)

* Use ssl context wrap socket in Python3.12+ * Remove redundant test modes --------- Co-authored-by: Bret Ambrose <[email protected]>
aws · Feb 14, 2024 · 800c16e · 800c16e
1 parent f78e330
commit 800c16e
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 59 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -36,11 +36,10 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        test-type: [ MutualAuth, MutualAuthT , Websocket, ALPN, ALPNT]
-        python-version: [ '2.x', '3.x' ]
-        #[MutualAuth, Websocket, ALPN]
+        test-type: [ MutualAuth, Websocket, ALPN ]
+        python-version: [ '3.7', '3.12' ]
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
       - uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
@@ -49,4 +48,5 @@ jobs:
           pip install pytest
           pip install mock
           pip install boto3
+          python --version
           ./test-integration/run/run.sh ${{ matrix.test-type }} 1000 100 7
diff --git a/AWSIoTPythonSDK/core/greengrass/discovery/providers.py b/AWSIoTPythonSDK/core/greengrass/discovery/providers.py
@@ -261,17 +261,28 @@ def _create_ssl_connection(self, sock):
             ssl_sock = ssl_context.wrap_socket(sock, server_hostname=self._host, do_handshake_on_connect=False)
             ssl_sock.do_handshake()
         else:
-            ssl_sock = ssl.wrap_socket(sock,
-                                       certfile=self._cert_path,
-                                       keyfile=self._key_path,
-                                       ca_certs=self._ca_path,
-                                       cert_reqs=ssl.CERT_REQUIRED,
-                                       ssl_version=ssl_protocol_version)
+            # To keep the SSL Context update minimal, only apply forced ssl context to python3.12+
+            force_ssl_context = sys.version_info[0] > 3 or (sys.version_info[0] == 3 and sys.version_info[1] >= 12)
+            if force_ssl_context:
+                ssl_context = ssl.SSLContext(ssl_protocol_version)
+                ssl_context.load_cert_chain(self._cert_path, self._key_path)
+                ssl_context.load_verify_locations(self._ca_path)
+                ssl_context.verify_mode = ssl.CERT_REQUIRED
+
+                ssl_sock = ssl_context.wrap_socket(sock)
+            else:
+                ssl_sock = ssl.wrap_socket(sock,
+                                           certfile=self._cert_path,
+                                           keyfile=self._key_path,
+                                           ca_certs=self._ca_path,
+                                           cert_reqs=ssl.CERT_REQUIRED,
+                                           ssl_version=ssl_protocol_version)
 
         self._logger.debug("Matching host name...")
         if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 2):
             self._tls_match_hostname(ssl_sock)
-        else:
+        elif sys.version_info[0] == 3 and sys.version_info[1] < 7:
+            # host name verification is handled internally in Python3.7+
             ssl.match_hostname(ssl_sock.getpeercert(), self._host)
 
         return ssl_sock

diff --git a/AWSIoTPythonSDK/core/protocol/paho/client.py b/AWSIoTPythonSDK/core/protocol/paho/client.py
@@ -793,11 +793,22 @@ def reconnect(self):
 
         verify_hostname = self._tls_insecure is False  # Decide whether we need to verify hostname
 
+        # To keep the SSL Context update minimal, only apply forced ssl context to python3.12+
+        force_ssl_context = sys.version_info[0] > 3 or (sys.version_info[0] == 3 and sys.version_info[1] >= 12)
+
         if self._tls_ca_certs is not None:
             if self._useSecuredWebsocket:
                 # Never assign to ._ssl before wss handshake is finished
                 # Non-None value for ._ssl will allow ops before wss-MQTT connection is established
-                rawSSL = ssl.wrap_socket(sock, ca_certs=self._tls_ca_certs, cert_reqs=ssl.CERT_REQUIRED)  # Add server certificate verification
+                if force_ssl_context:
+                    ssl_context = ssl.SSLContext()
+                    ssl_context.load_verify_locations(self._tls_ca_certs)
+                    ssl_context.verify_mode = ssl.CERT_REQUIRED
+
+                    rawSSL = ssl_context.wrap_socket(sock)
+                else:
+                    rawSSL = ssl.wrap_socket(sock, ca_certs=self._tls_ca_certs, cert_reqs=ssl.CERT_REQUIRED)  # Add server certificate verification
+
                 rawSSL.setblocking(0)  # Non-blocking socket
                 self._ssl = SecuredWebSocketCore(rawSSL, self._host, self._port, self._AWSAccessKeyIDCustomConfig, self._AWSSecretAccessKeyCustomConfig, self._AWSSessionTokenCustomConfig)  # Override the _ssl socket
                 # self._ssl.enableDebug()
@@ -816,19 +827,30 @@ def reconnect(self):
                 verify_hostname = False  # Since check_hostname in SSLContext is already set to True, no need to verify it again
                 self._ssl.do_handshake()
             else:
-                self._ssl = ssl.wrap_socket(
-                    sock,
-                    certfile=self._tls_certfile,
-                    keyfile=self._tls_keyfile,
-                    ca_certs=self._tls_ca_certs,
-                    cert_reqs=self._tls_cert_reqs,
-                    ssl_version=self._tls_version,
-                    ciphers=self._tls_ciphers)
+                if force_ssl_context:
+                    ssl_context = ssl.SSLContext(self._tls_version)
+                    ssl_context.load_cert_chain(self._tls_certfile, self._tls_keyfile)
+                    ssl_context.load_verify_locations(self._tls_ca_certs)
+                    ssl_context.verify_mode = self._tls_cert_reqs
+                    if self._tls_ciphers is not None:
+                        ssl_context.set_ciphers(self._tls_ciphers)
+
+                    self._ssl = ssl_context.wrap_socket(sock)
+                else:
+                    self._ssl = ssl.wrap_socket(
+                        sock,
+                        certfile=self._tls_certfile,
+                        keyfile=self._tls_keyfile,
+                        ca_certs=self._tls_ca_certs,
+                        cert_reqs=self._tls_cert_reqs,
+                        ssl_version=self._tls_version,
+                        ciphers=self._tls_ciphers)
 
             if verify_hostname:
                 if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 5):  # No IP host match before 3.5.x
                     self._tls_match_hostname()
-                else:
+                elif sys.version_info[0] == 3 and sys.version_info[1] < 7:
+                    # host name verification is handled internally in Python3.7+
                     ssl.match_hostname(self._ssl.getpeercert(), self._host)
 
         self._sock = sock

diff --git a/test-integration/run/run.sh b/test-integration/run/run.sh
@@ -35,16 +35,13 @@ USAGE="usage: run.sh <testMode> <NumberOfMQTTMessages> <LengthOfShadowRandomStri
 
 AWSMutualAuth_TodWorker_private_key="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestPrivateKey-vNUQU8"
 AWSMutualAuth_TodWorker_certificate="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestCertificate-vTRwjE"
-AWSMutualAuth_Desktop_private_key="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestDesktopPrivateKey-DdC7nv"
-AWSMutualAuth_Desktop_certificate="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestDesktopCertificate-IA4xbj"
 
 AWSGGDiscovery_TodWorker_private_key="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestGGDiscoveryPrivateKey-YHQI1F"
 AWSGGDiscovery_TodWorker_certificate="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestGGDiscoveryCertificate-TwlAcS"
 
 AWSSecretForWebsocket_TodWorker_KeyId="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestWebsocketAccessKeyId-1YdB9z"
 AWSSecretForWebsocket_TodWorker_SecretKey="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestWebsocketSecretAccessKey-MKTSaV"
-AWSSecretForWebsocket_Desktop_KeyId="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestWebsocketAccessKeyId-1YdB9z"
-AWSSecretForWebsocket_Desktop_SecretKey="arn:aws:secretsmanager:us-east-1:123124136734:secret:V1IotSdkIntegrationTestWebsocketSecretAccessKey-MKTSaV"
+
 
 SDKLocation="./AWSIoTPythonSDK"
 RetrieveAWSKeys="./test-integration/Tools/retrieve-key.py"
@@ -72,55 +69,41 @@ else
     TestMode=""
     echo "[STEP] Retrieve credentials from AWS"
     echo "***************************************************"
-    if [ "$1"x == "MutualAuth"x -o "$1"x == "MutualAuthT"x ]; then
+    if [ "$1"x == "MutualAuth"x ]; then
         AWSSetName_privatekey=${AWSMutualAuth_TodWorker_private_key}
-    	AWSSetName_certificate=${AWSMutualAuth_TodWorker_certificate}
-    	AWSDRSName_privatekey=${AWSGGDiscovery_TodWorker_private_key}
+    	  AWSSetName_certificate=${AWSMutualAuth_TodWorker_certificate}
+    	  AWSDRSName_privatekey=${AWSGGDiscovery_TodWorker_private_key}
         AWSDRSName_certificate=${AWSGGDiscovery_TodWorker_certificate}
         TestMode="MutualAuth"
-        if [ "$1"x == "MutualAuthT"x ]; then
-            AWSSetName_privatekey=${AWSMutualAuth_Desktop_private_key}
-    	    AWSSetName_certificate=${AWSMutualAuth_Desktop_certificate}
-        fi
-    	python ${RetrieveAWSKeys} ${AWSSetName_certificate} > ${CREDENTIAL_DIR}certificate.pem.crt
-    	python ${RetrieveAWSKeys} ${AWSSetName_privatekey} > ${CREDENTIAL_DIR}privateKey.pem.key
+    	  python ${RetrieveAWSKeys} ${AWSSetName_certificate} > ${CREDENTIAL_DIR}certificate.pem.crt
+    	  python ${RetrieveAWSKeys} ${AWSSetName_privatekey} > ${CREDENTIAL_DIR}privateKey.pem.key
         curl -s "${CA_CERT_URL}" > ${CA_CERT_PATH}
-        echo -e "URL retrieved certificate data:\n$(cat ${CA_CERT_PATH})\n"
-    	python ${RetrieveAWSKeys} ${AWSDRSName_certificate} > ${CREDENTIAL_DIR}certificate_drs.pem.crt
-    	python ${RetrieveAWSKeys} ${AWSDRSName_privatekey} > ${CREDENTIAL_DIR}privateKey_drs.pem.key
-    elif [ "$1"x == "Websocket"x -o "$1"x == "WebsocketT"x ]; then
-    	ACCESS_KEY_ID_ARN=$(python ${RetrieveAWSKeys} ${AWSSecretForWebsocket_TodWorker_KeyId})
+        echo -e "URL retrieved certificate data\n"
+    	  python ${RetrieveAWSKeys} ${AWSDRSName_certificate} > ${CREDENTIAL_DIR}certificate_drs.pem.crt
+    	  python ${RetrieveAWSKeys} ${AWSDRSName_privatekey} > ${CREDENTIAL_DIR}privateKey_drs.pem.key
+    elif [ "$1"x == "Websocket"x ]; then
+    	  ACCESS_KEY_ID_ARN=$(python ${RetrieveAWSKeys} ${AWSSecretForWebsocket_TodWorker_KeyId})
         ACCESS_SECRET_KEY_ARN=$(python ${RetrieveAWSKeys} ${AWSSecretForWebsocket_TodWorker_SecretKey})
         TestMode="Websocket"
-        if [ "$1"x == "WebsocketT"x ]; then
-            ACCESS_KEY_ID_ARN=$(python ${RetrieveAWSKeys} ${AWSSecretForWebsocket_Desktop_KeyId})
-            ACCESS_SECRET_KEY_ARN=$(python ${RetrieveAWSKeys} ${AWSSecretForWebsocket_Desktop_SecretKey})
-        fi
-        echo ${ACCESS_KEY_ID_ARN}
-        echo ${ACCESS_SECRET_KEY_ARN}
         export AWS_ACCESS_KEY_ID=${ACCESS_KEY_ID_ARN}
         export AWS_SECRET_ACCESS_KEY=${ACCESS_SECRET_KEY_ARN}
         curl -s "${CA_CERT_URL}" > ${CA_CERT_PATH}
-        echo -e "URL retrieved certificate data:\n$(cat ${CA_CERT_PATH})\n"
-    elif [ "$1"x == "ALPN"x -o "$1"x == "ALPNT"x ]; then
+        echo -e "URL retrieved certificate data\n"
+    elif [ "$1"x == "ALPN"x ]; then
         AWSSetName_privatekey=${AWSMutualAuth_TodWorker_private_key}
-    	AWSSetName_certificate=${AWSMutualAuth_TodWorker_certificate}
-    	AWSDRSName_privatekey=${AWSGGDiscovery_TodWorker_private_key}
+    	  AWSSetName_certificate=${AWSMutualAuth_TodWorker_certificate}
+    	  AWSDRSName_privatekey=${AWSGGDiscovery_TodWorker_private_key}
         AWSDRSName_certificate=${AWSGGDiscovery_TodWorker_certificate}
         TestMode="ALPN"
-        if [ "$1"x == "ALPNT"x ]; then
-            AWSSetName_privatekey=${AWSMutualAuth_Desktop_private_key}
-    	    AWSSetName_certificate=${AWSMutualAuth_Desktop_certificate}
-        fi
         python ${RetrieveAWSKeys} ${AWSSetName_certificate} > ${CREDENTIAL_DIR}certificate.pem.crt
-    	python ${RetrieveAWSKeys} ${AWSSetName_privatekey} > ${CREDENTIAL_DIR}privateKey.pem.key
+    	  python ${RetrieveAWSKeys} ${AWSSetName_privatekey} > ${CREDENTIAL_DIR}privateKey.pem.key
         curl -s "${CA_CERT_URL}" > ${CA_CERT_PATH}
-        echo -e "URL retrieved certificate data:\n$(cat ${CA_CERT_PATH})\n"
-    	python ${RetrieveAWSKeys} ${AWSDRSName_certificate} > ${CREDENTIAL_DIR}certificate_drs.pem.crt
-    	python ${RetrieveAWSKeys} ${AWSDRSName_privatekey} > ${CREDENTIAL_DIR}privateKey_drs.pem.key
+        echo -e "URL retrieved certificate data\n"
+    	  python ${RetrieveAWSKeys} ${AWSDRSName_certificate} > ${CREDENTIAL_DIR}certificate_drs.pem.crt
+    	  python ${RetrieveAWSKeys} ${AWSDRSName_privatekey} > ${CREDENTIAL_DIR}privateKey_drs.pem.key
     else
-    	echo "Mode not supported"
-    	exit 1
+    	  echo "Mode not supported"
+    	  exit 1
     fi
 # Obtain ZIP package and unzip it locally
     echo ${TestMode}