From 3419ae3b393ee95b466ef70ce33ea662395900da Mon Sep 17 00:00:00 2001
From: Andrew Hoener <ahoener@boston-engineering.com>
Date: Thu, 23 May 2024 16:03:46 -0400
Subject: [PATCH] Add a config variable for users to set expected CN when using
 CA verification

---
 docs/configuration.md     | 7 ++++++-
 pytak/client_functions.py | 7 ++++++-
 pytak/constants.py        | 1 +
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index eea3c57..3b2ae76 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -140,6 +140,11 @@ PyTAK can send & receive data over TLS by setting the following configuration pa
     Path to a file containing the CA Trust Store to use for remote certificate verification.
 
 
+* **`PYTAK_TLS_SERVER_EXPECTED_HOSTNAME`** (optional)
+
+  Expected hostname or CN of the connected server. Not used unless verifying hostname.
+
+
 * **`PYTAK_TLS_CLIENT_CIPHERS`** (optional)
     * Default: ``ALL`` 
 
@@ -152,4 +157,4 @@ PyTAK can send & receive data over TLS by setting the following configuration pa
 
 * **`PYTAK_TLS_CLIENT_PASSWORD`** (optional)
 
-    Password for PKCS#12 (.p12) password protected certificates or password protected Private Keys.
\ No newline at end of file
+    Password for PKCS#12 (.p12) password protected certificates or password protected Private Keys.
diff --git a/pytak/client_functions.py b/pytak/client_functions.py
index e62c4d9..0402674 100644
--- a/pytak/client_functions.py
+++ b/pytak/client_functions.py
@@ -227,6 +227,8 @@ async def protocol_factory(  # NOQA pylint: disable=too-many-locals,too-many-bra
         client_cafile = tls_config.get("PYTAK_TLS_CLIENT_CAFILE")
         client_password = tls_config.get("PYTAK_TLS_CLIENT_PASSWORD")
 
+        expected_server_hostname = tls_config.get("PYTAK_TLS_SERVER_EXPECTED_HOSTNAME")
+
         # Default cipher suite: ALL.
         #  Also available in FIPS: DEFAULT_FIPS_CIPHERS
         client_ciphers = tls_config.get("PYTAK_TLS_CLIENT_CIPHERS") or "ALL"
@@ -267,6 +269,7 @@ async def protocol_factory(  # NOQA pylint: disable=too-many-locals,too-many-bra
             warnings.warn(
                 "TLS CN/Hostname Check DISABLED by PYTAK_TLS_DONT_CHECK_HOSTNAME."
             )
+            expected_server_hostname = None
             ssl_ctx.check_hostname = False
 
         # Default to verifying cert:
@@ -277,7 +280,9 @@ async def protocol_factory(  # NOQA pylint: disable=too-many-locals,too-many-bra
             ssl_ctx.verify_mode = ssl.CERT_NONE
 
         try:
-            reader, writer = await asyncio.open_connection(host, port, ssl=ssl_ctx)
+            reader, writer = await asyncio.open_connection(
+                host, port, ssl=ssl_ctx, server_hostname=expected_server_hostname
+            )
         except ssl.SSLCertVerificationError as exc:
             raise SyntaxError(
                 "Consider setting PYTAK_TLS_DONT_CHECK_HOSTNAME=1 ?"
diff --git a/pytak/constants.py b/pytak/constants.py
index a330174..54e5c5d 100644
--- a/pytak/constants.py
+++ b/pytak/constants.py
@@ -78,6 +78,7 @@
     "PYTAK_TLS_DONT_CHECK_HOSTNAME",
     "PYTAK_TLS_DONT_VERIFY",
     "PYTAK_TLS_CLIENT_PASSWORD",
+    "PYTAK_TLS_SERVER_EXPECTED_HOSTNAME",
 ]
 
 DEFAULT_IMPORT_OTHER_CONFIGS: str = "0"