ucloud: Switch to SSLContext.

Signed-off-by: iabdalkader <[email protected]>

Author: iabdalkader

src/arduino_iot_cloud/umqtt.py

@@ -27,6 +27,7 @@
 import select
 import logging
 import arduino_iot_cloud.ussl as ssl
+import sys
 
 
 class MQTTException(Exception):
@@ -92,17 +93,14 @@ def connect(self, clean_session=True, timeout=5.0):
             self.sock.close()
             self.sock = None
 
-        try:
-            self.sock = socket.socket()
-            self.sock.settimeout(timeout)
-            self.sock = ssl.wrap_socket(self.sock, self.ssl_params)
-            self.sock.connect(addr)
-        except Exception:
-            self.sock.close()
-            self.sock = socket.socket()
-            self.sock.settimeout(timeout)
+        self.sock = socket.socket()
+        self.sock.settimeout(timeout)
+        if sys.implementation.name == "micropython":
             self.sock.connect(addr)
             self.sock = ssl.wrap_socket(self.sock, self.ssl_params)
+        else:
+            self.sock = ssl.wrap_socket(self.sock, self.ssl_params)
+            self.sock.connect(addr)
 
         premsg = bytearray(b"\x10\0\0\0\0\0")
         msg = bytearray(b"\x04MQTT\x04\x02\0\0")

src/arduino_iot_cloud/ussl.py

@@ -6,7 +6,6 @@
 #
 # SSL module with m2crypto backend for HSM support.
 
-import sys
 import ssl
 
 pkcs11 = None
@@ -16,62 +15,69 @@
 _MODULE_PATH = "/usr/lib/softhsm/libsofthsm2.so"
 
 
-def wrap_socket(
-    sock,
-    ssl_params={},
-):
-    if any(k not in ssl_params for k in ("keyfile", "certfile", "pin")):
-        # Use Micro/CPython's SSL
-        if sys.implementation.name == "micropython":
-            # Load key, cert and CA from DER files, and pass them as binary blobs.
-            mpargs = {"keyfile": "key", "certfile": "cert", "ca_certs": "cadata"}
-            for k, v in mpargs.items():
-                if k in ssl_params and "der" in ssl_params[k]:
-                    with open(ssl_params.pop(k), "rb") as f:
-                        ssl_params[v] = f.read()
-        return ssl.wrap_socket(sock, **ssl_params)
-
-    # Use M2Crypto to load key and cert from HSM.
-    from M2Crypto import m2, SSL, Engine
-
-    global pkcs11
-    if pkcs11 is None:
-        pkcs11 = Engine.load_dynamic_engine(
-            "pkcs11", ssl_params.get("engine_path", _ENGINE_PATH)
-        )
-        pkcs11.ctrl_cmd_string(
-            "MODULE_PATH", ssl_params.get("module_path", _MODULE_PATH)
-        )
-        pkcs11.ctrl_cmd_string("PIN", ssl_params["pin"])
-        pkcs11.init()
-
-    # Create and configure SSL context
-    ctx = SSL.Context("tls")
-    ctx.set_default_verify_paths()
-    ctx.set_allow_unknown_ca(False)
-
+def wrap_socket(sock, ssl_params={}):
+    keyfile = ssl_params.get("keyfile", None)
+    certfile = ssl_params.get("certfile", None)
+    cafile = ssl_params.get("cafile", None)
+    cadata = ssl_params.get("cadata", None)
     ciphers = ssl_params.get("ciphers", None)
-    if ciphers is not None:
-        ctx.set_cipher_list(ciphers)
+    verify = ssl_params.get("verify_mode", ssl.CERT_NONE)
+    hostname = ssl_params.get("server_hostname", None)
+    use_hsm = ssl_params.get("use_hsm", False)
 
-    ca_certs = ssl_params.get("ca_certs", None)
-    if ca_certs is not None:
-        if ctx.load_verify_locations(ca_certs) != 1:
-            raise Exception("Failed to load CA certs")
-
-    cert_reqs = ssl_params.get("cert_reqs", ssl.CERT_NONE)
-    if cert_reqs == ssl.CERT_NONE:
-        cert_reqs = SSL.verify_none
+    if not use_hsm:
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        if hasattr(ctx, "set_default_verify_paths"):
+            ctx.set_default_verify_paths()
+        if verify != ssl.CERT_REQUIRED:
+            ctx.check_hostname = False
+        ctx.verify_mode = verify
+        if keyfile is not None and certfile is not None:
+            ctx.load_cert_chain(certfile, keyfile)
+        if ciphers is not None:
+            ctx.set_ciphers(ciphers)
+        if cafile is not None or cadata is not None:
+            ctx.load_verify_locations(cafile, cadata)
+        return ctx.wrap_socket(sock, server_hostname=hostname)
     else:
-        cert_reqs = SSL.verify_peer
-    ctx.set_verify(cert_reqs, depth=9)
+        # Use M2Crypto to load key and cert from HSM.
+        from M2Crypto import m2, SSL, Engine
+
+        global pkcs11
+        if pkcs11 is None:
+            pkcs11 = Engine.load_dynamic_engine(
+                "pkcs11", ssl_params.get("engine_path", _ENGINE_PATH)
+            )
+            pkcs11.ctrl_cmd_string(
+                "MODULE_PATH", ssl_params.get("module_path", _MODULE_PATH)
+            )
+            if "pin" in ssl_params:
+                pkcs11.ctrl_cmd_string("PIN", ssl_params["pin"])
+            pkcs11.init()
+
+        # Create and configure SSL context
+        ctx = SSL.Context("tls")
+        ctx.set_default_verify_paths()
+        ctx.set_allow_unknown_ca(False)
+        if verify == ssl.CERT_NONE:
+            ctx.set_verify(SSL.verify_none, depth=9)
+        else:
+            ctx.set_verify(SSL.verify_peer | SSL.verify_fail_if_no_peer_cert, depth=9)
+        if cafile is not None:
+            if ctx.load_verify_locations(cafile) != 1:
+                raise Exception("Failed to load CA certs")
+        if ciphers is not None:
+            ctx.set_cipher_list(ciphers)
 
-    # Set key/cert
-    key = pkcs11.load_private_key(ssl_params["keyfile"])
-    m2.ssl_ctx_use_pkey_privkey(ctx.ctx, key.pkey)
+        key = pkcs11.load_private_key(keyfile)
+        m2.ssl_ctx_use_pkey_privkey(ctx.ctx, key.pkey)
 
-    cert = pkcs11.load_certificate(ssl_params["certfile"])
-    m2.ssl_ctx_use_x509(ctx.ctx, cert.x509)
+        cert = pkcs11.load_certificate(certfile)
+        m2.ssl_ctx_use_x509(ctx.ctx, cert.x509)
 
-    SSL.Connection.postConnectionCheck = None
-    return SSL.Connection(ctx, sock=sock)
+        sslobj = SSL.Connection(ctx, sock=sock)
+        if verify == ssl.CERT_NONE:
+            sslobj.clientPostConnectionCheck = None
+        elif hostname is not None:
+            sslobj.set1_host(hostname)
+        return sslobj