#!/usr/bin/env python3

import argparse
import socket
import ssl
import sys

import h2.connection
import h2.config
import h2.events


def make_authority(host: str, port: int) -> str:
    if port == 443:
        return host
    return f"{host}:{port}"


def run_probe(
    host: str,
    port: int,
    path: str,
    regular_header_count: int,
    insecure: bool,
) -> None:
    context = ssl.create_default_context()
    context.set_alpn_protocols(["h2"])

    if insecure:
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE

    configuration = h2.config.H2Configuration(
        client_side=True,
        header_encoding="utf-8",
    )
    connection = h2.connection.H2Connection(config=configuration)

    with socket.create_connection((host, port), timeout=5) as tcp_socket:
        with context.wrap_socket(
            tcp_socket,
            server_hostname=host,
        ) as tls_socket:
            negotiated_protocol = tls_socket.selected_alpn_protocol()

            if negotiated_protocol != "h2":
                raise RuntimeError(
                    f"The server negotiated {negotiated_protocol!r}, "
                    "not HTTP/2 ('h2')."
                )

            print(f"[+] Negotiated ALPN protocol: {negotiated_protocol}")

            connection.initiate_connection()
            tls_socket.sendall(connection.data_to_send())

            stream_id = connection.get_next_available_stream_id()

            request_headers = [
                (":method", "GET"),
                (":scheme", "https"),
                (":authority", make_authority(host, port)),
                (":path", path),
            ]
            request_headers.extend(
                (f"x-fooooooo-{index:04d}", "a")
                for index in range(regular_header_count)
            )

            print(
                f"[+] Sending one stream with "
                f"{regular_header_count} regular headers and "
                f"4 pseudo-headers"
            )

            connection.send_headers(
                stream_id=stream_id,
                headers=request_headers,
                end_stream=True,
            )
            tls_socket.sendall(connection.data_to_send())

            tls_socket.settimeout(5)

            stream_finished = False

            while not stream_finished:
                try:
                    received = tls_socket.recv(65535)
                except socket.timeout:
                    print("[-] Timed out waiting for a response.")
                    break

                if not received:
                    print("[-] The peer closed the TLS connection.")
                    break

                events = connection.receive_data(received)

                for event in events:
                    if isinstance(event, h2.events.ResponseReceived):
                        print("[+] Response headers:")
                        for name, value in event.headers:
                            print(f"    {name}: {value}")

                    elif isinstance(event, h2.events.DataReceived):
                        print(
                            f"[+] Received {len(event.data)} response bytes"
                        )
                        connection.acknowledge_received_data(
                            event.flow_controlled_length,
                            event.stream_id,
                        )

                    elif isinstance(event, h2.events.StreamEnded):
                        print(f"[+] Stream {event.stream_id} ended normally.")
                        stream_finished = True

                    elif isinstance(event, h2.events.StreamReset):
                        print(
                            f"[+] Stream {event.stream_id} was reset; "
                            f"error code={event.error_code}"
                        )
                        stream_finished = True

                    elif isinstance(event, h2.events.ConnectionTerminated):
                        print(
                            "[+] HTTP/2 connection terminated; "
                            f"error code={event.error_code}, "
                            f"last stream={event.last_stream_id}"
                        )
                        stream_finished = True

                pending_data = connection.data_to_send()
                if pending_data:
                    tls_socket.sendall(pending_data)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Send one bounded HTTP/2 header-count probe."
    )
    parser.add_argument("host", help="Laboratory HTTP/2 server hostname or IP")
    parser.add_argument("--port", type=int, default=443)
    parser.add_argument("--path", default="/")
    parser.add_argument(
        "--headers",
        type=int,
        default=20,
        help="Number of regular test headers.",
    )
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="Disable certificate verification for a self-signed lab server",
    )

    args = parser.parse_args()

    try:
        run_probe(
            host=args.host,
            port=args.port,
            path=args.path,
            regular_header_count=args.headers,
            insecure=args.insecure,
        )
    except (OSError, ssl.SSLError, RuntimeError, ValueError) as exc:
        print(f"[!] Error: {exc}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())