Server latency tool

README.md

@@ -14,6 +14,7 @@ Scan the ports of your server or get the DNS records with this small scripts.
 - **WHOIS Lookup**: Perform a WHOIS lookup for a given domain.
 - **TraceRoute**: Trace the route to a given IP address or hostname.
 - **SSL/TLS Certificate Check**: Inspect a remote server's certificate for subject, issuer, and expiry.
+- **Website Latency Check**: Break down request latency into DNS lookup, TCP connect, TLS handshake, time to first byte, and content download.
 
 ## Installation
 

ctrl_node/common/modules/latency.py

@@ -0,0 +1,136 @@
+import socket
+import ssl
+import time
+
+import questionary
+
+from ctrl_node.common.modules.module_utils import rate_ttfb
+from ctrl_node.common.utils import askforhost
+
+
+class LatencyCheck:
+    """
+    Class that measures the latency of connecting to a website, broken
+    down by phase: DNS lookup, TCP connect, TLS handshake, time to first
+    byte, and content download.
+    """
+
+    def __init__(self, host=None, port=443, use_tls=True, path="/"):
+        """
+        Initializes the class with the host to measure.
+
+        Arguments:
+            host {str}: Optional host to check. If not provided, it will default to None.
+            port {int}: Port to connect to. Defaults to 443.
+            use_tls {bool}: Whether to connect over TLS. Defaults to True.
+            path {str}: Path to request. Defaults to "/".
+        """
+        self.host = host or askforhost()
+        self.port = port
+        self.use_tls = use_tls
+        self.path = path
+        self.timings = {}
+        self.ip_address = None
+
+    def measure_latency(self):
+        """
+        Connects to the host and records the time taken by each phase of
+        the request: DNS lookup, TCP connect, TLS handshake, time to
+        first byte, and content download.
+        """
+        self.timings = {}
+        try:
+            address_info = self.time_phase(
+                "DNS lookup", lambda: socket.getaddrinfo(self.host, self.port)
+            )
+            self.ip_address = address_info[0][4][0]
+
+            sock = self.time_phase(
+                "TCP connect",
+                lambda: socket.create_connection(
+                    (self.ip_address, self.port), timeout=10
+                ),
+            )
+
+            if self.use_tls:
+                context = ssl.create_default_context()
+                sock = self.time_phase(
+                    "TLS handshake",
+                    lambda: context.wrap_socket(sock, server_hostname=self.host),
+                )
+
+            self.request_content(sock)
+            sock.close()
+        except (socket.error, ssl.SSLError, socket.timeout) as error:
+            print(f"Could not measure latency for {self.host}:{self.port} - {error}")
+
+    def time_phase(self, name: str, action):
+        """
+        Runs the given action and stores how long it took under the
+        given phase name.
+
+        Arguments:
+            name {str}: Name of the phase being timed.
+            action {Callable}: Zero-argument callable to run and time.
+
+        Returns:
+            The return value of the action.
+        """
+        start = time.perf_counter()
+        result = action()
+        self.timings[name] = time.perf_counter() - start
+        return result
+
+    def request_content(self, sock):
+        """
+        Sends an HTTP GET request over the given socket and measures the
+        time to first byte and the time to download the rest of the
+        content.
+
+        Arguments:
+            sock {socket.socket}: Connected (and possibly TLS-wrapped) socket.
+        """
+        request = (
+            f"GET {self.path} HTTP/1.1\r\n"
+            f"Host: {self.host}\r\n"
+            "Connection: close\r\n"
+            "\r\n"
+        ).encode()
+
+        start = time.perf_counter()
+        sock.sendall(request)
+        sock.recv(4096)
+        self.timings["Time to first byte"] = time.perf_counter() - start
+
+        download_start = time.perf_counter()
+        while sock.recv(4096):
+            pass
+        self.timings["Content download"] = time.perf_counter() - download_start
+
+    def display_results(self):
+        """
+        Displays the measured timings for each phase, each phase's share
+        of the total time, and a rating of the server's responsiveness.
+        """
+        if not self.timings:
+            print("No latency data available.")
+        else:
+            total = sum(self.timings.values())
+            header = f"Latency results for {self.host}"
+            if self.ip_address:
+                header += f" ({self.ip_address})"
+
+            print("-" * 60)
+            print(header)
+            print("-" * 60)
+            for phase, duration in self.timings.items():
+                share = (duration / total * 100) if total else 0
+                print(f"{phase}: \t {duration * 1000:.2f} ms \t ({share:.1f}%)")
+            print("-" * 60)
+            print(f"Total: \t {total * 1000:.2f} ms")
+
+            ttfb = self.timings.get("Time to first byte")
+            if ttfb is not None:
+                print(f"Server response rating: {rate_ttfb(ttfb)} (TTFB-based)")
+
+        questionary.text("Press Enter to return to the menu...").ask()

ctrl_node/common/modules/module_utils.py

@@ -14,3 +14,21 @@ def show_port_open_or_close(result: int, port: int, silence: bool = False):
     else:
         if not silence:
             print(f"Port {port}: 	 Close")
+
+
+def rate_ttfb(ttfb: float) -> str:
+    """
+    Rates a time to first byte measurement using the Core Web Vitals
+    TTFB thresholds: good is at or under 0.8s, poor is over 1.8s.
+
+    Arguments:
+        ttfb {float}: Time to first byte, in seconds.
+
+    Returns:
+        str: "Good", "Needs Improvement", or "Poor".
+    """
+    if ttfb <= 0.8:
+        return "Good"
+    if ttfb <= 1.8:
+        return "Needs Improvement"
+    return "Poor"

ctrl_node/ctrl_node.py

@@ -4,6 +4,7 @@
 
 from ctrl_node.common.modules.dns import DnsScan
 from ctrl_node.common.modules.host_to_ip import HostToIp
+from ctrl_node.common.modules.latency import LatencyCheck
 from ctrl_node.common.modules.ports import ScanPorts
 from ctrl_node.common.modules.ssl_check import SslCheck
 from ctrl_node.common.modules.traceroute import Traceroute
@@ -40,6 +41,7 @@ def show_menu(self):
                 "5 - Traceroute",
                 "6 - Show application version",
                 "7 - SSL Certificate Check",
+                "8 - Website Latency Check",
                 "0 - Exit",
             ],
         ).ask()
@@ -76,6 +78,10 @@ def execute_menu(self, option: str):
             ssl_check = SslCheck()
             ssl_check.fetch_certificate()
             ssl_check.display_certificate_info()
+        elif option == "8 - Website Latency Check":
+            latency_check = LatencyCheck()
+            latency_check.measure_latency()
+            latency_check.display_results()
         elif option == "0 - Exit":
             sys.exit()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "CTRL_Node"
-version = "0.0.2"
+version = "0.0.3"
 description = "Various tools for managing or checking servers"
 authors = ["NotTheRealWallyx"]
 keywords = ["python", "scripts", "server", "tools", "dns", "records"]
@@ -26,6 +26,9 @@ isort = "^6.1.0"
 [tool.poetry.scripts]
 ctrl-node = "ctrl_node.ctrl_node:main"
 
+[tool.isort]
+profile = "black"
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"

tests/common/modules/test_latency.py

@@ -0,0 +1,105 @@
+import socket
+import unittest
+from unittest.mock import MagicMock, patch
+
+from ctrl_node.common.modules.latency import LatencyCheck
+
+
+class TestLatencyCheck(unittest.TestCase):
+    @patch("ctrl_node.common.modules.latency.askforhost", return_value="example.com")
+    def test_initialization(self, mock_askforhost):
+        instance = LatencyCheck()
+        self.assertEqual(instance.host, "example.com")
+        self.assertEqual(instance.port, 443)
+        self.assertTrue(instance.use_tls)
+        self.assertEqual(instance.path, "/")
+        mock_askforhost.assert_called_once()
+
+    @patch("ctrl_node.common.modules.latency.ssl.create_default_context")
+    @patch("ctrl_node.common.modules.latency.socket.create_connection")
+    @patch("ctrl_node.common.modules.latency.socket.getaddrinfo")
+    def test_measure_latency_success(
+        self, mock_getaddrinfo, mock_create_connection, mock_context_factory
+    ):
+        mock_getaddrinfo.return_value = [
+            (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
+        ]
+        mock_create_connection.return_value = MagicMock()
+
+        mock_ssock = MagicMock()
+        mock_ssock.recv.side_effect = [b"HTTP/1.1 200 OK", b""]
+        mock_context_factory.return_value.wrap_socket.return_value = mock_ssock
+
+        instance = LatencyCheck(host="example.com")
+        instance.measure_latency()
+
+        self.assertIn("DNS lookup", instance.timings)
+        self.assertIn("TCP connect", instance.timings)
+        self.assertIn("TLS handshake", instance.timings)
+        self.assertIn("Time to first byte", instance.timings)
+        self.assertIn("Content download", instance.timings)
+        self.assertEqual(instance.ip_address, "93.184.216.34")
+        mock_ssock.sendall.assert_called_once()
+        mock_ssock.close.assert_called_once()
+
+    @patch(
+        "ctrl_node.common.modules.latency.socket.getaddrinfo",
+        side_effect=socket.gaierror("Name resolution failed"),
+    )
+    @patch("builtins.print")
+    def test_measure_latency_dns_error(self, mock_print, mock_getaddrinfo):
+        instance = LatencyCheck(host="example.com")
+        instance.measure_latency()
+
+        self.assertEqual(instance.timings, {})
+        mock_print.assert_called_once_with(
+            "Could not measure latency for example.com:443 - Name resolution failed"
+        )
+
+    @patch("ctrl_node.common.modules.latency.questionary.text")
+    @patch("builtins.print")
+    def test_display_results_with_data(self, mock_print, mock_questionary_text):
+        mock_questionary_text.return_value.ask.return_value = None
+        instance = LatencyCheck(host="example.com")
+        instance.timings = {
+            "DNS lookup": 0.01,
+            "TCP connect": 0.02,
+            "Time to first byte": 0.05,
+        }
+
+        instance.display_results()
+
+        mock_print.assert_any_call("Latency results for example.com")
+        mock_print.assert_any_call("Server response rating: Good (TTFB-based)")
+        mock_questionary_text.assert_called_once_with(
+            "Press Enter to return to the menu..."
+        )
+
+    @patch("ctrl_node.common.modules.latency.questionary.text")
+    @patch("builtins.print")
+    def test_display_results_includes_ip_and_percentage(
+        self, mock_print, mock_questionary_text
+    ):
+        mock_questionary_text.return_value.ask.return_value = None
+        instance = LatencyCheck(host="example.com")
+        instance.ip_address = "93.184.216.34"
+        instance.timings = {"DNS lookup": 0.01, "TCP connect": 0.03}
+
+        instance.display_results()
+
+        mock_print.assert_any_call("Latency results for example.com (93.184.216.34)")
+        mock_print.assert_any_call("DNS lookup: \t 10.00 ms \t (25.0%)")
+
+    @patch("ctrl_node.common.modules.latency.questionary.text")
+    @patch("builtins.print")
+    def test_display_results_without_data(self, mock_print, mock_questionary_text):
+        mock_questionary_text.return_value.ask.return_value = None
+        instance = LatencyCheck(host="example.com")
+
+        instance.display_results()
+
+        mock_print.assert_any_call("No latency data available.")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/common/modules/test_module_utils.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 
-from ctrl_node.common.modules.module_utils import show_port_open_or_close
+from ctrl_node.common.modules.module_utils import rate_ttfb, show_port_open_or_close
 
 
 class TestUtils(unittest.TestCase):
@@ -16,6 +16,15 @@ def test_show_port_open_or_close_closed(self, mock_print):
         show_port_open_or_close(1, 80)
         mock_print.assert_called_once_with("Port 80: \t Close")
 
+    def test_rate_ttfb_good(self):
+        self.assertEqual(rate_ttfb(0.8), "Good")
+
+    def test_rate_ttfb_needs_improvement(self):
+        self.assertEqual(rate_ttfb(1.8), "Needs Improvement")
+
+    def test_rate_ttfb_poor(self):
+        self.assertEqual(rate_ttfb(1.9), "Poor")
+
 
 if __name__ == "__main__":
     unittest.main()