"""TLS-рукопожатие: за что платит первое соединение и что возвращает второе.

ЗАЧЕМ ЭТОТ ФАЙЛ. Про TLS на собеседовании обычно пересказывают одну фразу:
«возобновление сессии экономит круг». Она верна для TLS 1.2 и неверна для
TLS 1.3, где полное рукопожатие и так укладывается в один круг. Проверить это
словами нельзя — нужно увидеть время в кругах, а не в миллисекундах.

КАК ЗДЕСЬ ИЗМЕРЯЮТСЯ КРУГИ. Настоящий сервер в интернете для этого не годится:
его задержка неизвестна и непостоянна, а в этой среде трафик вдобавок идёт
через посредника. Поэтому и сервер, и звено связи здесь свои. Сервер стоит на
loopback, а между ним и клиентом — ретранслятор, который задерживает каждую
переданную порцию ровно на DELAY_MS в одну сторону. Круг известен точно, и
время рукопожатия читается в кругах: 40 мс при RTT 40 мс — один круг, 80 мс —
два.

Тот же ретранслятор считает байты в обе стороны, и это второй предмет замера:
чем именно дешевле возобновлённое рукопожатие, если круг оно не экономит.

ЧТО ЗДЕСЬ ИЗМЕРЯЕТСЯ:
  1. Сколько кругов стоит полное рукопожатие в TLS 1.3 и в TLS 1.2.
  2. Что меняет возобновление сессии в каждой из версий.
  3. Сколько байт занимает рукопожатие полное и возобновлённое.
  4. Цена рукопожатия против цены соединения без шифрования.

ЧЕГО ЗДЕСЬ НЕТ. Ни 0-RTT, ни настоящих сертификатов, ни проверки цепочки:
клиент здесь не проверяет сертификат вовсе. Это сделано намеренно — предмет
замера в числе кругов и байт, а не в стоимости проверки подписи, которая
зависит от типа ключа и длины цепочки.

ПОЧЕМУ ПОДПИСИ ПО-АНГЛИЙСКИ. Урок существует в двух языках и цитирует запись
прогона дословно обеими версиями.

ЗАПУСК: python3 bench/tls/handshake.py
Вывод: runs/handshake.txt
"""

import os
import socket
import ssl
import subprocess
import sys
import tempfile
import threading
import time

DELAY_MS = 20  # задержка звена в одну сторону; круг вдвое больше
RTT_MS = DELAY_MS * 2


def show(title: str) -> None:
    print()
    print(title)
    print("-" * len(title))


def row(label: str, value: object) -> None:
    print(f"  {label:<46} {value}")


def make_cert(directory: str) -> tuple[str, str]:
    """Самоподписанный ключ и сертификат: RSA-2048, ровно один в цепочке."""
    key = os.path.join(directory, "key.pem")
    crt = os.path.join(directory, "cert.pem")
    subprocess.run(
        [
            "openssl", "req", "-x509", "-newkey", "rsa:2048",
            "-keyout", key, "-out", crt, "-days", "1", "-nodes",
            "-subj", "/CN=localhost",
        ],
        check=True,
        capture_output=True,
    )
    return key, crt


class Server:
    """TLS-сервер на loopback с заданной максимальной версией протокола."""

    def __init__(self, key: str, crt: str, version: ssl.TLSVersion) -> None:
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        context.load_cert_chain(crt, key)
        context.maximum_version = version
        self.context = context
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", 0))
        self.sock.listen(64)
        self.port = self.sock.getsockname()[1]
        threading.Thread(target=self._loop, daemon=True).start()

    def _loop(self) -> None:
        while True:
            try:
                conn, _ = self.sock.accept()
            except OSError:
                return
            threading.Thread(target=self._serve, args=(conn,), daemon=True).start()

    def _serve(self, conn: socket.socket) -> None:
        try:
            tls = self.context.wrap_socket(conn, server_side=True)
            tls.recv(16)
            tls.sendall(b"ok")
            tls.close()
        except OSError:
            pass


class DelayedLink:
    """Ретранслятор, который задерживает каждую порцию и считает байты.

    ПОЧЕМУ НЕ tc netem. Настройка сетевого стека требует прав, которых у
    замера может не быть, и меняет машину целиком. Ретранслятор в пользователь-
    ском коде даёт ту же наблюдаемую величину — известный круг, — не трогая
    ничего за пределами процесса.

    ПОЧЕМУ ЗАДЕРЖКА НА КАЖДУЮ ПОРЦИЮ, А НЕ НА ПАКЕТ. Мы работаем с потоком, а
    не с пакетами: каждый recv возвращает то, что сторона отправила одним
    заходом. Для рукопожатия этого достаточно — его порции и есть его полёты.
    """

    def __init__(self, upstream_port: int) -> None:
        self.upstream_port = upstream_port
        self.to_server = 0
        self.to_client = 0
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", 0))
        self.sock.listen(64)
        self.port = self.sock.getsockname()[1]
        threading.Thread(target=self._loop, daemon=True).start()

    def reset(self) -> None:
        self.to_server = 0
        self.to_client = 0

    def _pump(self, src: socket.socket, dst: socket.socket, upward: bool) -> None:
        try:
            while True:
                data = src.recv(65536)
                if not data:
                    break
                if upward:
                    self.to_server += len(data)
                else:
                    self.to_client += len(data)
                time.sleep(DELAY_MS / 1000)
                dst.sendall(data)
        except OSError:
            pass
        finally:
            try:
                dst.shutdown(socket.SHUT_WR)
            except OSError:
                pass

    def _loop(self) -> None:
        while True:
            try:
                client, _ = self.sock.accept()
            except OSError:
                return

            def handle(client: socket.socket = client) -> None:
                upstream = socket.create_connection(("127.0.0.1", self.upstream_port))
                threading.Thread(
                    target=self._pump, args=(client, upstream, True), daemon=True
                ).start()
                threading.Thread(
                    target=self._pump, args=(upstream, client, False), daemon=True
                ).start()

            threading.Thread(target=handle, daemon=True).start()


CLIENT = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
CLIENT.check_hostname = False
CLIENT.verify_mode = ssl.CERT_NONE


def dial(port: int, session: ssl.SSLSession | None = None) -> dict[str, object]:
    """Одно соединение целиком: TCP, рукопожатие, запрос и ответ.

    ПОЧЕМУ СЕССИЯ БЕРЁТСЯ ПОСЛЕ ОБМЕНА ДАННЫМИ. В TLS 1.3 билет на
    возобновление присылается уже ПОСЛЕ рукопожатия, отдельным сообщением.
    Если снять `session` сразу, билета в ней ещё нет, и следующее соединение
    молча окажется полным — с `session_reused` равным False.
    """
    started = time.perf_counter()
    raw = socket.create_connection(("127.0.0.1", port), timeout=10)
    connected = time.perf_counter()
    tls = CLIENT.wrap_socket(raw, server_hostname="localhost", session=session)
    shaken = time.perf_counter()
    tls.sendall(b"x")
    tls.recv(16)
    answered = time.perf_counter()
    result = {
        "connect_ms": (connected - started) * 1000,
        "handshake_ms": (shaken - connected) * 1000,
        "request_ms": (answered - shaken) * 1000,
        "reused": tls.session_reused,
        "version": tls.version(),
        "session": tls.session,
    }
    tls.close()
    return result


def plain_dial(port: int) -> dict[str, float]:
    """То же самое без шифрования — чтобы было с чем сравнивать."""
    started = time.perf_counter()
    raw = socket.create_connection(("127.0.0.1", port), timeout=10)
    connected = time.perf_counter()
    raw.sendall(b"x")
    raw.recv(16)
    answered = time.perf_counter()
    raw.close()
    return {
        "connect_ms": (connected - started) * 1000,
        "request_ms": (answered - connected) * 1000,
    }


def rounds(ms: float) -> str:
    """Время в кругах — то, ради чего звено и задерживает трафик."""
    return f"{ms / RTT_MS:.1f}"


def main() -> None:
    print(f"Python {sys.version.split()[0]} · {ssl.OPENSSL_VERSION}")
    print(f"link delay {DELAY_MS} ms one way, so one round trip is {RTT_MS} ms")
    print("self-signed RSA-2048 certificate, one in the chain, not verified")

    directory = tempfile.mkdtemp()
    key, crt = make_cert(directory)

    measured = {}
    for label, version in (
        ("TLS 1.3", ssl.TLSVersion.TLSv1_3),
        ("TLS 1.2", ssl.TLSVersion.TLSv1_2),
    ):
        server = Server(key, crt, version)
        link = DelayedLink(server.port)
        time.sleep(0.2)

        link.reset()
        full = dial(link.port)
        full["to_server"] = link.to_server
        full["to_client"] = link.to_client

        link.reset()
        resumed = dial(link.port, full["session"])
        resumed["to_server"] = link.to_server
        resumed["to_client"] = link.to_client

        measured[label] = (full, resumed, server)

    show("1. A FULL HANDSHAKE COSTS ROUND TRIPS, AND THE VERSION DECIDES HOW MANY")
    print(f"  {'version':>10} {'handshake, ms':>15} {'in round trips':>16} {'reused':>8}")
    for label in ("TLS 1.3", "TLS 1.2"):
        full, _, _ = measured[label]
        print(
            f"  {label:>10} {full['handshake_ms']:>15.1f} {rounds(full['handshake_ms']):>16}"
            f" {str(full['reused']):>8}"
        )
    print()
    print("  The link is the same in both rows and so is the certificate. What")
    print("  differs is the protocol: 1.3 finishes in one round trip, 1.2 needs")
    print("  two. Half the received wisdom about TLS is really about 1.2.")

    show("2. RESUMPTION SAVES A ROUND TRIP ONLY WHERE THERE WAS ONE TO SAVE")
    print(f"  {'version':>10} {'full, ms':>10} {'resumed, ms':>13} {'full':>8} {'resumed':>9} {'reused':>8}")
    for label in ("TLS 1.3", "TLS 1.2"):
        full, resumed, _ = measured[label]
        print(
            f"  {label:>10} {full['handshake_ms']:>10.1f} {resumed['handshake_ms']:>13.1f}"
            f" {rounds(full['handshake_ms']):>8} {rounds(resumed['handshake_ms']):>9}"
            f" {str(resumed['reused']):>8}"
        )
    print()
    print("  In 1.2 resumption removes a round trip: two become one. In 1.3")
    print("  there was nothing to remove - the full handshake already took one,")
    print("  and the resumed one takes the same. Both really were resumed:")
    print("  the reused column says so.")

    show("3. WHAT RESUMPTION ACTUALLY SAVES: THE CERTIFICATE")
    print(f"  {'version':>10} {'handshake':>12} {'client -> server':>18} {'server -> client':>18}")
    for label in ("TLS 1.3", "TLS 1.2"):
        full, resumed, _ = measured[label]
        print(f"  {label:>10} {'full':>12} {full['to_server']:>18} {full['to_client']:>18}")
        print(f"  {label:>10} {'resumed':>12} {resumed['to_server']:>18} {resumed['to_client']:>18}")
    print()
    print("  The bytes from the server are where the certificate lives. A")
    print("  resumed handshake does not send it, and that is the saving that")
    print("  survives in both versions - unlike the round trip.")

    show("4. WHAT ENCRYPTION COSTS AGAINST NO ENCRYPTION AT ALL")
    plain_server = socket.socket()
    plain_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    plain_server.bind(("127.0.0.1", 0))
    plain_server.listen(64)

    def plain_loop() -> None:
        while True:
            try:
                conn, _ = plain_server.accept()
            except OSError:
                return

            def handle(conn: socket.socket = conn) -> None:
                try:
                    conn.recv(16)
                    conn.sendall(b"ok")
                finally:
                    conn.close()

            threading.Thread(target=handle, daemon=True).start()

    threading.Thread(target=plain_loop, daemon=True).start()
    plain_link = DelayedLink(plain_server.getsockname()[1])
    time.sleep(0.2)

    plain = plain_dial(plain_link.port)
    full13, resumed13, _ = measured["TLS 1.3"]
    print(f"  {'connection':>22} {'connect, ms':>13} {'handshake, ms':>15} {'exchange, ms':>14}")
    print(
        f"  {'plain TCP':>22} {plain['connect_ms']:>13.1f} {'-':>15}"
        f" {plain['request_ms']:>14.1f}"
    )
    print(
        f"  {'TLS 1.3, full':>22} {full13['connect_ms']:>13.1f}"
        f" {full13['handshake_ms']:>15.1f} {full13['request_ms']:>14.1f}"
    )
    print(
        f"  {'TLS 1.3, resumed':>22} {resumed13['connect_ms']:>13.1f}"
        f" {resumed13['handshake_ms']:>15.1f} {resumed13['request_ms']:>14.1f}"
    )
    print()
    row("the handshake alone, round trips", rounds(full13["handshake_ms"]))
    print()
    print("  The handshake column is the whole price of encryption at this")
    print("  layer: one round trip, paid once per connection rather than once")
    print("  per request. A connection that is kept and reused pays it once")
    print("  and never again.")
    print()
    print("  The exchange column is larger under TLS than over plain TCP, and")
    print("  that is not the encryption: in 1.3 the server sends session")
    print("  tickets right after the handshake, and the link delays that")
    print("  flight too. It is an artefact of the measurement, not a cost of")
    print("  sending a request.")


if __name__ == "__main__":
    main()