From db583bed592e7706b07e6ca08069ff08229b3263 Mon Sep 17 00:00:00 2001
From: Sasha Romijn <github@mxsasha.eu>
Date: Tue, 2 Apr 2024 09:33:58 +0200
Subject: [PATCH] Add support for custom HTTP user agent

---
 sslyze/connection_helpers/http_request_generator.py | 8 ++++++--
 sslyze/plugins/early_data_plugin.py                 | 7 ++++++-
 sslyze/plugins/http_headers_plugin.py               | 8 ++++++--
 sslyze/server_setting.py                            | 3 +++
 4 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/sslyze/connection_helpers/http_request_generator.py b/sslyze/connection_helpers/http_request_generator.py
index e283aa8f..9d698e19 100644
--- a/sslyze/connection_helpers/http_request_generator.py
+++ b/sslyze/connection_helpers/http_request_generator.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from sslyze import __version__
 
 
@@ -16,5 +18,7 @@ class HttpRequestGenerator:
     )
 
     @classmethod
-    def get_request(cls, host: str, path: str = "/") -> bytes:
-        return cls.HTTP_GET_FORMAT.format(host=host, path=path, user_agent=cls.DEFAULT_USER_AGENT).encode("utf-8")
+    def get_request(cls, host: str, path: str = "/", user_agent: Optional[str] = None) -> bytes:
+        if not user_agent:
+            user_agent = cls.DEFAULT_USER_AGENT
+        return cls.HTTP_GET_FORMAT.format(host=host, path=path, user_agent=user_agent).encode("utf-8")
diff --git a/sslyze/plugins/early_data_plugin.py b/sslyze/plugins/early_data_plugin.py
index b79420dc..f03719e0 100644
--- a/sslyze/plugins/early_data_plugin.py
+++ b/sslyze/plugins/early_data_plugin.py
@@ -89,7 +89,12 @@ def _test_early_data_support(server_info: ServerConnectivityInfo) -> bool:
         # Perform an SSL handshake and keep the session
         ssl_connection.connect()
         # Send and receive data for the TLS session to be created
-        ssl_connection.ssl_client.write(HttpRequestGenerator.get_request(host=server_info.server_location.hostname))
+        ssl_connection.ssl_client.write(
+            HttpRequestGenerator.get_request(
+                host=server_info.server_location.hostname,
+                user_agent=server_info.network_configuration.http_user_agent,
+            )
+        )
         ssl_connection.ssl_client.read(2048)
         session = ssl_connection.ssl_client.get_session()
     except ServerRejectedTlsHandshake:
diff --git a/sslyze/plugins/http_headers_plugin.py b/sslyze/plugins/http_headers_plugin.py
index bc57a9c2..0754cff0 100755
--- a/sslyze/plugins/http_headers_plugin.py
+++ b/sslyze/plugins/http_headers_plugin.py
@@ -200,7 +200,9 @@ def _retrieve_and_analyze_http_response(server_info: ServerConnectivityInfo) ->
             # Send an HTTP GET request to the server
             ssl_connection.ssl_client.write(
                 HttpRequestGenerator.get_request(
-                    host=server_info.network_configuration.tls_server_name_indication, path=next_location_path
+                    host=server_info.network_configuration.tls_server_name_indication,
+                    path=next_location_path,
+                    user_agent=server_info.network_configuration.http_user_agent,
                 )
             )
             http_response = HttpResponseParser.parse_from_ssl_connection(ssl_connection.ssl_client)
@@ -225,7 +227,9 @@ def _retrieve_and_analyze_http_response(server_info: ServerConnectivityInfo) ->
 
     # Prepare the results
     initial_http_request = HttpRequestGenerator.get_request(
-        host=server_info.network_configuration.tls_server_name_indication, path="/"
+        host=server_info.network_configuration.tls_server_name_indication,
+        path="/",
+        user_agent=server_info.network_configuration.http_user_agent,
     ).decode("ascii")
 
     if http_error_trace:
diff --git a/sslyze/server_setting.py b/sslyze/server_setting.py
index 0b1c61de..e46c6812 100644
--- a/sslyze/server_setting.py
+++ b/sslyze/server_setting.py
@@ -173,6 +173,8 @@ class ServerNetworkConfiguration:
         xmpp_to_hostname: The hostname to set within the `to` attribute of the XMPP stream. If not supplied, the
             server's hostname will be used. Should only be set if the supplied `tls_wrapped_protocol` is an
             XMPP protocol.
+        http_user_agent: The User-Agent to send in HTTP requests. If not supplied, a default Chrome-like
+            is used that includes the sslyze version.
         network_timeout: The timeout (in seconds) to be used when attempting to establish a connection to the
             server.
         network_max_retries: The number of retries SSLyze will perform when attempting to establish a connection
@@ -184,6 +186,7 @@ class ServerNetworkConfiguration:
     tls_client_auth_credentials: Optional[ClientAuthenticationCredentials] = None
 
     xmpp_to_hostname: Optional[str] = None
+    http_user_agent: Optional[str] = None
 
     network_timeout: int = 5
     network_max_retries: int = 3