from __future__ import annotations
import asyncio
import struct

from ..protocol import constants as C
from ..protocol import wire


async def probe(host: str, port: int, sni: str | None, timeout: float = 6.0) -> tuple[bool, str | None]:
    """CVE-2014-0224 CCS Injection probe.

    Send ChangeCipherSpec immediately after ClientHello. A vulnerable server
    will accept and not send an Unexpected Message alert.
    """
    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
    except Exception as e:
        return False, f"connect: {e}"
    try:
        exts = (
            (wire.ext_server_name(sni) if sni and not wire._is_ip_literal(sni) else b"")
            + wire.ext_ec_point_formats()
            + wire.ext_supported_groups([0x001d, 0x0017, 0x0018])
            + wire.ext_signature_algorithms([0x0403, 0x0401])
            + wire.ext_renegotiation_info_empty()
        )
        ch = wire.build_client_hello(
            record_version=C.TLS_1_0,
            client_hello_version=C.TLS_1_2,
            hostname=sni,
            cipher_suites=[0xc013, 0xc014, 0x002f, 0x0035, 0x000a],
            extensions=exts,
        )
        writer.write(ch)
        # Send early ChangeCipherSpec
        ccs_record = bytes([C.CT_CHANGE_CIPHER_SPEC]) + struct.pack("!H", C.TLS_1_2) + struct.pack("!H", 1) + b"\x01"
        writer.write(ccs_record)
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        try:
            data = await asyncio.wait_for(reader.read(4096), timeout=timeout)
        except Exception:
            data = b""
        if not data:
            return False, "no response (likely dropped connection)"
        alerts = _find_alerts(data)
        # Any fatal alert (level=2) indicates the server rejected the early CCS: NOT vulnerable.
        if any(lv == 2 for lv, _desc in alerts):
            return False, "server sent fatal alert (rejected early CCS)"
        # Explicit unexpected_message alert (10): NOT vulnerable.
        if any(desc == 10 for _lv, desc in alerts):
            return False, "server rejected with unexpected_message alert"
        # Otherwise: no definitive signal (server may have ignored CCS or handshake still pending).
        # Conservative default: not vulnerable. True detection requires a complete handshake exchange.
        return False, "no definitive signal of CCS acceptance"
    except Exception as e:
        return False, f"{e.__class__.__name__}: {e}"
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.0)
        except Exception:
            pass


def _find_alerts(data: bytes) -> list[tuple[int, int]]:
    alerts: list[tuple[int, int]] = []
    i = 0
    while i + 5 <= len(data):
        ct = data[i]
        rec_len = (data[i + 3] << 8) | data[i + 4]
        body = data[i + 5:i + 5 + rec_len]
        if ct == C.CT_ALERT and len(body) >= 2:
            alerts.append((body[0], body[1]))
        i += 5 + rec_len
        if rec_len == 0:
            break
    return alerts
