diff --git a/README.md b/README.md index 641e2b2..6669aff 100644 --- a/README.md +++ b/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 diff --git a/ctrl_node/common/modules/latency.py b/ctrl_node/common/modules/latency.py new file mode 100644 index 0000000..36639fd --- /dev/null +++ b/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() diff --git a/ctrl_node/common/modules/module_utils.py b/ctrl_node/common/modules/module_utils.py index b09707c..138eb5f 100644 --- a/ctrl_node/common/modules/module_utils.py +++ b/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" diff --git a/ctrl_node/ctrl_node.py b/ctrl_node/ctrl_node.py index b5d3d97..428a41b 100644 --- a/ctrl_node/ctrl_node.py +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index fd713e0..e6cea93 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/common/modules/test_latency.py b/tests/common/modules/test_latency.py new file mode 100644 index 0000000..46041df --- /dev/null +++ b/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() diff --git a/tests/common/modules/test_module_utils.py b/tests/common/modules/test_module_utils.py index 3d5ff90..1dc0d23 100644 --- a/tests/common/modules/test_module_utils.py +++ b/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()