http-latency-prober/python/latprobe/probe.py

from __future__ import annotations

import datetime
import socket
import ssl
import time
import urllib.parse
from dataclasses import dataclass, field

# ── verbose detail dataclasses ────────────────────────────────────────────────


@dataclass
class CertInfo:
    cn: str = ""
    sans: list[str] = field(default_factory=list)
    expiry: str = ""       # ISO date "YYYY-MM-DD"
    issuer_cn: str = ""
    verified: bool = False  # True = TLS handshake passed verification


@dataclass
class VerboseDetail:
    resolved_ip: str = ""
    tls_version: str = ""
    tls_cipher: str = ""
    tls_bits: int = 0
    cert: CertInfo | None = None
    headers: dict[str, str] = field(default_factory=dict)  # all parsed response headers


# ── core dataclasses ──────────────────────────────────────────────────────────


@dataclass
class Options:
    timeout: float = 10.0
    verbose: bool = False


@dataclass
class Phase:
    ms: float = 0.0
    present: bool = False


@dataclass
class Result:
    url: str
    dns:      Phase = field(default_factory=Phase)
    connect:  Phase = field(default_factory=Phase)
    tls:      Phase = field(default_factory=Phase)
    ttfb:     Phase = field(default_factory=Phase)
    transfer: Phase = field(default_factory=Phase)
    total:    Phase = field(default_factory=Phase)
    status_code: int = 0
    fail_phase: str = ""
    err: Exception | None = None
    detail: VerboseDetail | None = None  # populated only when opts.verbose=True


# ── internal helpers ──────────────────────────────────────────────────────────


def _p(start: float, end: float) -> Phase:
    return Phase(ms=(end - start) * 1000, present=True)


def _parse_status(data: bytes) -> int:
    eol = data.find(b"\r\n")
    if eol == -1:
        eol = data.find(b"\n")
    if eol == -1:
        return 0
    parts = data[:eol].decode("latin-1", errors="replace").split(None, 2)
    if len(parts) >= 2:
        try:
            return int(parts[1])
        except ValueError:
            pass
    return 0


def _parse_cert_date(s: str) -> str:
    """Convert SSL cert date 'Jun 14 00:00:00 2025 GMT' → '2025-06-14'."""
    if not s:
        return ""
    s = " ".join(s.split())  # collapse double-spaces ("May  5 ..." → "May 5 ...")
    for fmt in ("%b %d %H:%M:%S %Y %Z", "%b %d %H:%M:%S %Y"):
        try:
            return datetime.datetime.strptime(s, fmt).strftime("%Y-%m-%d")
        except ValueError:
            pass
    return ""


def _parse_cert(peer: dict, *, verified: bool) -> CertInfo:
    """Build CertInfo from the dict returned by SSLSocket.getpeercert()."""
    def _attr(rdns, key: str) -> str:
        for rdn in rdns:
            for k, v in rdn:
                if k == key:
                    return v
        return ""

    issuer = peer.get("issuer", ())
    # Prefer organizationName for issuer (more human-readable than CA CN)
    issuer_cn = _attr(issuer, "organizationName") or _attr(issuer, "commonName")

    return CertInfo(
        cn=_attr(peer.get("subject", ()), "commonName"),
        sans=[v for k, v in peer.get("subjectAltName", ()) if k == "DNS"],
        expiry=_parse_cert_date(peer.get("notAfter", "")),
        issuer_cn=issuer_cn,
        verified=verified,
    )


def _parse_response_headers(buf: bytes) -> dict[str, str]:
    """Parse all HTTP response headers from a raw buffer (up to \\r\\n\\r\\n)."""
    header_section = buf.split(b"\r\n\r\n", 1)[0]
    lines = header_section.decode("latin-1", errors="replace").splitlines()
    headers: dict[str, str] = {}
    for line in lines[1:]:  # skip status line
        if ":" in line:
            name, _, value = line.partition(":")
            headers[name.strip()] = value.strip()
    return headers


# ── measure ───────────────────────────────────────────────────────────────────


def measure(raw_url: str, opts: Options | None = None) -> Result:
    """Probe raw_url and return a Result with per-phase timings.

    Partial phases are preserved when the request fails mid-flight.
    Redirects are not followed; bodies are drained so Transfer timing is real.
    When opts.verbose=True, Result.detail is populated with resolved IP,
    TLS metadata, certificate info, and response headers.
    """
    if opts is None:
        opts = Options()
    r = Result(url=raw_url)

    if opts.verbose:
        r.detail = VerboseDetail()

    parsed = urllib.parse.urlparse(raw_url)
    scheme = parsed.scheme.lower()
    if scheme not in ("http", "https"):
        r.fail_phase = "request"
        r.err = ValueError(f"unsupported scheme: {scheme!r}")
        return r

    host = parsed.hostname or ""
    port = parsed.port or (443 if scheme == "https" else 80)
    path = parsed.path or "/"
    if parsed.query:
        path = path + "?" + parsed.query
    use_tls = scheme == "https"

    t_start = time.perf_counter()

    # ── DNS ──────────────────────────────────────────────────────────────────

    t0 = time.perf_counter()
    try:
        infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
    except socket.gaierror as exc:
        r.fail_phase = "dns"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r
    r.dns = _p(t0, time.perf_counter())

    if r.detail is not None:
        r.detail.resolved_ip = str(infos[0][4][0])

    # ── TCP connect ───────────────────────────────────────────────────────────

    addr = infos[0][4]
    family = infos[0][0]
    sock = socket.socket(family, socket.SOCK_STREAM)
    sock.settimeout(opts.timeout)
    t0 = time.perf_counter()
    try:
        sock.connect(addr)
    except socket.timeout as exc:
        sock.close()
        r.connect = _p(t0, time.perf_counter())
        r.fail_phase = "timeout"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r
    except OSError as exc:
        sock.close()
        r.connect = _p(t0, time.perf_counter())
        r.fail_phase = "connect"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r
    r.connect = _p(t0, time.perf_counter())

    # ── TLS handshake (HTTPS only) ────────────────────────────────────────────

    if use_tls:
        ctx = ssl.create_default_context()
        t0 = time.perf_counter()
        try:
            sock = ctx.wrap_socket(sock, server_hostname=host)
        except socket.timeout as exc:
            r.tls = _p(t0, time.perf_counter())
            r.fail_phase = "timeout"
            r.err = exc
            r.total = _p(t_start, time.perf_counter())
            return r
        except (ssl.SSLError, OSError) as exc:
            r.tls = _p(t0, time.perf_counter())
            r.fail_phase = "tls"
            r.err = exc
            r.total = _p(t_start, time.perf_counter())
            return r
        r.tls = _p(t0, time.perf_counter())

        if r.detail is not None:
            r.detail.tls_version = sock.version() or ""
            cipher_name, _, bits = sock.cipher()
            r.detail.tls_cipher = cipher_name or ""
            r.detail.tls_bits   = bits or 0
            peer = sock.getpeercert()
            if peer:
                r.detail.cert = _parse_cert(peer, verified=True)

    # ── Send request ──────────────────────────────────────────────────────────

    request = (
        f"GET {path} HTTP/1.1\r\n"
        f"Host: {host}\r\n"
        f"Connection: close\r\n"
        f"User-Agent: latprobe/1.0\r\n"
        f"\r\n"
    ).encode()
    try:
        sock.sendall(request)
    except OSError as exc:
        sock.close()
        r.fail_phase = "request"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r
    t_wrote = time.perf_counter()

    # ── TTFB — accumulate until end of headers (\r\n\r\n) ────────────────────
    #
    # t_first_byte is stamped on the first recv() that returns data (unchanged
    # semantics vs the old single-recv approach).  We keep reading until the
    # full header section is in buf so that response headers can be parsed when
    # opts.verbose=True.

    try:
        buf = b""
        t_first_byte: float | None = None
        while b"\r\n\r\n" not in buf:
            chunk = sock.recv(4096)
            if not chunk:
                raise OSError("server closed connection before headers complete")
            if t_first_byte is None:
                t_first_byte = time.perf_counter()
            buf += chunk
    except socket.timeout as exc:
        sock.close()
        r.fail_phase = "timeout"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r
    except OSError as exc:
        sock.close()
        r.fail_phase = "transfer"
        r.err = exc
        r.total = _p(t_start, time.perf_counter())
        return r

    t_first_byte = t_first_byte or time.perf_counter()
    r.ttfb = _p(t_wrote, t_first_byte)
    r.status_code = _parse_status(buf)

    if r.detail is not None:
        r.detail.headers = _parse_response_headers(buf)

    # ── Transfer — drain remaining body ───────────────────────────────────────

    try:
        while True:
            chunk = sock.recv(65536)
            if not chunk:
                break
    except OSError:
        pass
    finally:
        sock.close()

    t_end = time.perf_counter()
    r.transfer = _p(t_first_byte, t_end)
    r.total = _p(t_start, t_end)
    return r