Add possibility to use ssl_context extra for SMTP and IMAP connections

airflow/providers/imap/CHANGELOG.rst

@@ -38,6 +38,8 @@ Setting it to "none" brings back the "none" setting that was used in previous ve
 but it is not recommended due to security reasons and this setting disables validation
 of certificates and allows MITM attacks.
 
+You can also override "ssl_context" per-connection by setting "ssl_context" in the connection extra.
+
 3.2.2
 .....
 

airflow/providers/imap/hooks/imap.py

@@ -84,7 +84,11 @@ def _build_client(self, conn: Connection) -> imaplib.IMAP4_SSL | imaplib.IMAP4:
         if use_ssl:
             from airflow.configuration import conf
 
-            ssl_context_string = conf.get("imap", "SSL_CONTEXT", fallback=None)
+            extra_ssl_context = conn.extra_dejson.get("ssl_context", None)
+            if extra_ssl_context:
+                ssl_context_string = extra_ssl_context
+            else:
+                ssl_context_string = conf.get("imap", "SSL_CONTEXT", fallback=None)
             if ssl_context_string is None:
                 ssl_context_string = conf.get("email", "SSL_CONTEXT", fallback=None)
             if ssl_context_string is None:

airflow/providers/smtp/CHANGELOG.rst

@@ -39,6 +39,8 @@ Setting it to "none" brings back the "none" setting that was used in previous ve
 but it is not recommended due to security reasons ad this setting disables validation
 of certificates and allows MITM attacks.
 
+You can also override "ssl_context" per-connection by setting "ssl_context" in the connection extra.
+
 1.2.0
 .....
 

airflow/providers/smtp/hooks/smtp.py

@@ -112,7 +112,11 @@ def _build_client(self) -> smtplib.SMTP_SSL | smtplib.SMTP:
         if self.use_ssl:
             from airflow.configuration import conf
 
-            ssl_context_string = conf.get("smtp_provider", "SSL_CONTEXT", fallback=None)
+            extra_ssl_context = self.conn.extra_dejson.get("ssl_context", None)
+            if extra_ssl_context:
+                ssl_context_string = extra_ssl_context
+            else:
+                ssl_context_string = conf.get("smtp_provider", "SSL_CONTEXT", fallback=None)
             if ssl_context_string is None:
                 ssl_context_string = conf.get("email", "SSL_CONTEXT", fallback=None)
             if ssl_context_string is None:

docs/apache-airflow-providers-imap/connections/imap.rst

@@ -55,6 +55,9 @@ Extra (optional)
     Specify the extra parameters (as json dictionary)
 
     * ``use_ssl``: If set to false, then a non-ssl connection is being used. Default is true. Also note that changing the ssl option also influences the default port being used.
+    * ``ssl_context``: Can be "default" or "none". Only valid when "use_ssl" is used. The "default" context provides a balance between security and compatibility, "none" is not recommended
+      as it disables validation of certificates and allow MITM attacks and is only needed in case your certificates are wrongly configured in your system. If not specified, defaults are taken from the
+      "imap", "ssl_context" configuration with the fallback to "email". "ssl_context" configuration. If none of it is specified, "default" is used.
 
 When specifying the connection in environment variable you should specify
 it using URI syntax.

docs/apache-airflow-providers-smtp/connections/smtp.rst

@@ -59,6 +59,9 @@ Extra (optional)
     * ``timeout``: The SMTP connection creation timeout in seconds. Default is 30.
     * ``disable_tls``: By default the SMTP connection is created in TLS mode. Set to false to disable tls mode.
     * ``retry_limit``: How many attempts to connect to the server before raising an exception. Default is 5.
+    * ``ssl_context``: Can be "default" or "none". Only valid when SSL is used. The "default" context provides a balance between security and compatibility, "none" is not recommended
+      as it disables validation of certificates and allow MITM attacks, and is only needed in case your certificates are wrongly configured in your system. If not specified, defaults are taken from the
+      "smtp_provider", "ssl_context" configuration with the fallback to "email". "ssl_context" configuration. If none of it is specified, "default" is used.
 
 When specifying the connection in environment variable you should specify
 it using URI syntax.

tests/providers/imap/hooks/test_imap.py

@@ -114,6 +114,33 @@ def test_connect_and_disconnect_imap_ssl_context_none(self, create_default_conte
         mock_conn.login.assert_called_once_with("imap_user", "imap_password")
         assert mock_conn.logout.call_count == 1
 
+    @patch(imaplib_string)
+    @patch("ssl.create_default_context")
+    def test_connect_and_disconnect_imap_ssl_context_from_extra(self, create_default_context, mock_imaplib):
+        mock_conn = _create_fake_imap(mock_imaplib)
+        db.merge_conn(
+            Connection(
+                conn_id="imap_ssl_context_from_extra",
+                conn_type="imap",
+                host="imap_server_address",
+                login="imap_user",
+                password="imap_password",
+                port=1993,
+                extra=json.dumps(dict(use_ssl=True, ssl_context="default")),
+            )
+        )
+
+        with conf_vars({("imap", "ssl_context"): "none"}):
+            with ImapHook(imap_conn_id="imap_ssl_context_from_extra"):
+                pass
+
+        assert create_default_context.called
+        mock_imaplib.IMAP4_SSL.assert_called_once_with(
+            "imap_server_address", 1993, ssl_context=create_default_context.return_value
+        )
+        mock_conn.login.assert_called_once_with("imap_user", "imap_password")
+        assert mock_conn.logout.call_count == 1
+
     @patch(imaplib_string)
     @patch("ssl.create_default_context")
     def test_connect_and_disconnect_imap_ssl_context_default(self, create_default_context, mock_imaplib):

tests/providers/smtp/hooks/test_smtp.py

@@ -230,6 +230,30 @@ def test_send_mime_ssl_none_email_context(self, create_default_context, mock_smt
         assert not create_default_context.called
         mock_smtp_ssl.assert_called_once_with(host="smtp_server_address", port=465, timeout=30, context=None)
 
+    @patch("smtplib.SMTP_SSL")
+    @patch("smtplib.SMTP")
+    @patch("ssl.create_default_context")
+    def test_send_mime_ssl_extra_context(self, create_default_context, mock_smtp, mock_smtp_ssl):
+        mock_smtp_ssl.return_value = Mock()
+        conn = Connection(
+            conn_id="smtp_ssl_extra",
+            conn_type="smtp",
+            host="smtp_server_address",
+            login=None,
+            password="None",
+            port=465,
+            extra=json.dumps(dict(ssl_context="none", from_email="from")),
+        )
+        db.merge_conn(conn)
+        with conf_vars({("smtp", "smtp_ssl"): "True", ("smtp_provider", "ssl_context"): "default"}):
+            with SmtpHook(smtp_conn_id="smtp_ssl_extra") as smtp_hook:
+                smtp_hook.send_email_smtp(
+                    to="to", subject="subject", html_content="content", from_email="from"
+                )
+        assert not mock_smtp.called
+        assert not create_default_context.called
+        mock_smtp_ssl.assert_called_once_with(host="smtp_server_address", port=465, timeout=30, context=None)
+
     @patch("smtplib.SMTP_SSL")
     @patch("smtplib.SMTP")
     @patch("ssl.create_default_context")