"""
TLS probing utility for the Port Agent.

This module performs a full TCP connection and TLS handshake against a local
port, returning raw negotiated parameters such as TLS version, cipher suite,
ALPN protocol, and the peer certificate in DER form. It is intentionally
lightweight and avoids interpretation: the caller decides what constitutes a
vulnerability or misconfiguration.

The probe enforces a strict “local‑only” guardrail by resolving the target
address through `resolve_local_address()`. This prevents the agent from scanning
remote hosts and keeps the tool aligned with the project’s safety boundaries.

Each stage of the connection—DNS resolution, TCP connect, TLS wrapping, and
handshake—is individually timed and error‑coded. Failures return structured
diagnostics (`error` and `errno`) to help the agent reason about what went wrong
and to support developer feedback when the agent suggests improvements.

The result object (`TlsProbeResult`) is a simple dataclass designed for
agent‑friendly consumption: no exceptions, no side effects, just structured
signals about what happened during the handshake.
"""

import socket
import ssl
import time
from dataclasses import dataclass
from typing import Optional
from .utility import resolve_local_address


@dataclass
class TlsProbeResult:
    success: bool
    tls_version: Optional[str]
    cipher_suite: Optional[str]
    alpn_protocol: Optional[str]
    certificate_der: Optional[bytes]
    tcp_connect_ms: float
    tls_handshake_ms: float
    resolved_address: Optional[str]
    error: Optional[str]          # NEW
    errno: Optional[int]          # NEW

def tls_probe(
    local_address_mode: str,
    port: int,
    timeout: float = 3.0,
    min_version: ssl.TLSVersion = ssl.TLSVersion.TLSv1,
    max_version: ssl.TLSVersion = ssl.TLSVersion.MAXIMUM_SUPPORTED,
) -> TlsProbeResult:
    """
    Perform a TLS handshake and extract negotiated parameters.
    Returns raw negotiated values without interpretation.
    """

    # --- Local-host guardrail ---
    local_address = resolve_local_address(local_address_mode)
    if not local_address:
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=0.0,
            tls_handshake_ms=0.0,
            resolved_address=None,
            error="INVALID_HOST_NON_LOCAL",
            errno=-1,
        )

    # --- Resolve address ---
    try:
        resolved = socket.getaddrinfo(local_address, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
        family, socktype, proto, canonname, sockaddr = resolved[0]
        resolved_addr = str(sockaddr[0])
    except Exception as e:
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=0.0,
            tls_handshake_ms=0.0,
            resolved_address=None,
            error=f"RESOLUTION_FAILED: {e}",
            errno=-2,
        )

    # --- TCP socket ---
    s = socket.socket(family, socket.SOCK_STREAM)
    s.settimeout(timeout)

    tcp_start = time.perf_counter()
    try:
        err = s.connect_ex((local_address, port))
    except Exception as e:
        s.close()
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=0.0,
            tls_handshake_ms=0.0,
            resolved_address=resolved_addr,
            error=f"CONNECT_EXCEPTION: {e}",
            errno=-3,
        )
    tcp_end = time.perf_counter()

    tcp_connect_ms = (tcp_end - tcp_start) * 1000.0

    # --- TCP failed ---
    if err != 0:
        s.close()
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=tcp_connect_ms,
            tls_handshake_ms=0.0,
            resolved_address=resolved_addr,
            error="TCP_CONNECT_FAILED",
            errno=err,
        )

    # --- TLS context ---
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    ctx.minimum_version = min_version
    ctx.maximum_version = max_version

    # --- Wrap socket ---
    try:
        tls_sock = ctx.wrap_socket(s, server_hostname=local_address)
    except Exception as e:
        s.close()
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=tcp_connect_ms,
            tls_handshake_ms=0.0,
            resolved_address=resolved_addr,
            error=f"WRAP_SOCKET_FAILED: {e}",
            errno=-4,
        )

    # --- TLS handshake ---
    hs_start = time.perf_counter()
    handshake_error = None  # ensure always defined

    try:
        tls_sock.do_handshake()
        hs_success = True
    except Exception as e:
        hs_success = False
        handshake_error = e

    hs_end = time.perf_counter()
    tls_handshake_ms = (hs_end - hs_start) * 1000.0


    if not hs_success:
        tls_sock.close()
        return TlsProbeResult(
            success=False,
            tls_version=None,
            cipher_suite=None,
            alpn_protocol=None,
            certificate_der=None,
            tcp_connect_ms=tcp_connect_ms,
            tls_handshake_ms=tls_handshake_ms,
            resolved_address=resolved_addr,
            error=f"HANDSHAKE_FAILED: {handshake_error}",
            errno=-5,
        )

    # --- Extract negotiated parameters ---
    try:
        tls_version = tls_sock.version()
    except Exception:
        tls_version = None

    try:
        cipher_info = tls_sock.cipher()
        cipher_suite = cipher_info[0] if cipher_info else None
    except Exception:
        cipher_suite = None

    try:
        alpn_protocol = tls_sock.selected_alpn_protocol()
    except Exception:
        alpn_protocol = None

    try:
        cert = tls_sock.getpeercert(binary_form=True)
    except Exception:
        cert = None

    tls_sock.close()

    return TlsProbeResult(
        success=True,
        tls_version=tls_version,
        cipher_suite=cipher_suite,
        alpn_protocol=alpn_protocol,
        certificate_der=cert,
        tcp_connect_ms=tcp_connect_ms,
        tls_handshake_ms=tls_handshake_ms,
        resolved_address=resolved_addr,
        error=None,
        errno=0,
    )


if __name__ == "__main__":
    host = "127.0.0.1"
    port = 443

    print(f"Running tls_probe({host}, {port})...")
    result = tls_probe(host, port)

    print("\n=== TLS Probe Result ===")
    print(f"Resolved Address : {result.resolved_address}")
    print(f"Success          : {result.success}")
    print(f"TLS Version      : {result.tls_version}")
    print(f"Cipher Suite     : {result.cipher_suite}")
    print(f"ALPN Protocol    : {result.alpn_protocol}")
    print(f"Cert (bytes)     : {len(result.certificate_der) if result.certificate_der else None}")
    print(f"TCP Connect (ms) : {result.tcp_connect_ms:.2f}")
    print(f"TLS Handshake(ms): {result.tls_handshake_ms:.2f}")
    print(f"Error            : {result.error}")
    print(f"error #          : {result.errno}")
