"""Sonde de handshake WireGuard, sans monter le tunnel.

Le diagnostic s'arrêtait au constat « le port UDP ne rejette rien ». C'est le
maximum qu'une sonde aveugle puisse dire : un serveur WireGuard ignore
silencieusement tout paquet non authentifié, si bien qu'un port fermé par un
pare-feu, un serveur écoutant sur un autre port et un serveur en bonne santé
produisent exactement la même absence de réponse. Quand le tunnel n'était pas
monté, le rapport concluait donc « chaîne validée, reste à monter le tunnel »
sur une configuration qui n'avait aucune chance de négocier.

Ce module construit un vrai message d'initiation (Noise_IKpsk2_25519_
ChaChaPoly_BLAKE2s, protocole WireGuard v1) et interprète la réponse. Le
serveur ne répond que si notre clé publique statique lui est connue et si sa
propre clé publique est celle que nous croyons : une réponse prouve donc la
configuration, et son déchiffrement dit si une clé pré-partagée est en jeu.

Aucun état n'est conservé : la session négociée est jetée aussitôt. La sonde
n'ouvre pas de tunnel et ne perturbe pas une session existante — le serveur
traite l'initiation comme celle d'un pair qui se reconnecte.
"""

import base64
import hashlib
import hmac
import os
import socket
import struct
import time

from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey, X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305

_CONSTRUCTION = b"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
_IDENTIFIER = b"WireGuard v1 zx2c4 Jason@zx2c4.com"
_LABEL_MAC1 = b"mac1----"

_MSG_INITIATION = 1
_MSG_RESPONSE = 2
_MSG_COOKIE = 3

_INITIATION_LEN = 148
_RESPONSE_LEN = 92
_ZERO_PSK = b"\x00" * 32


def _hash(*parts: bytes) -> bytes:
    h = hashlib.blake2s(digest_size=32)
    for p in parts:
        h.update(p)
    return h.digest()


def _mac(key: bytes, data: bytes) -> bytes:
    return hashlib.blake2s(data, digest_size=16, key=key).digest()


def _hmac(key: bytes, data: bytes) -> bytes:
    return hmac.new(key, data, hashlib.blake2s).digest()


def _kdf(key: bytes, data: bytes, n: int) -> list[bytes]:
    """KDF de Noise : dérive `n` clés de 32 octets."""
    tau0 = _hmac(key, data)
    out = [_hmac(tau0, b"\x01")]
    for i in range(2, n + 1):
        out.append(_hmac(tau0, out[-1] + bytes([i])))
    return out


def _aead_encrypt(key: bytes, counter: int, plain: bytes, ad: bytes) -> bytes:
    nonce = b"\x00" * 4 + struct.pack("<Q", counter)
    return ChaCha20Poly1305(key).encrypt(nonce, plain, ad)


def _aead_decrypt(key: bytes, counter: int, cipher: bytes, ad: bytes) -> bytes:
    nonce = b"\x00" * 4 + struct.pack("<Q", counter)
    return ChaCha20Poly1305(key).decrypt(nonce, cipher, ad)


def _tai64n() -> bytes:
    """Horodatage TAI64N, tel que WireGuard l'attend (12 octets)."""
    now = time.time()
    secs = int(now)
    nanos = int((now - secs) * 1e9)
    return struct.pack(">QI", 0x400000000000000A + secs, nanos)


class _Initiation:
    """Message 1 construit, et l'état de Noise nécessaire à lire le message 2."""

    def __init__(self, static_priv: X25519PrivateKey, server_pub: X25519PublicKey):
        self._static_priv = static_priv
        server_pub_raw = server_pub.public_bytes_raw()

        c = _hash(_CONSTRUCTION)
        h = _hash(c, _IDENTIFIER)
        h = _hash(h, server_pub_raw)

        self._eph_priv = X25519PrivateKey.generate()
        eph_pub = self._eph_priv.public_key().public_bytes_raw()

        c = _kdf(c, eph_pub, 1)[0]
        h = _hash(h, eph_pub)

        c, k = _kdf(c, self._eph_priv.exchange(server_pub), 2)
        enc_static = _aead_encrypt(
            k, 0, static_priv.public_key().public_bytes_raw(), h)
        h = _hash(h, enc_static)

        c, k = _kdf(c, static_priv.exchange(server_pub), 2)
        enc_ts = _aead_encrypt(k, 0, _tai64n(), h)
        h = _hash(h, enc_ts)

        self._c, self._h = c, h
        self.sender_index = int.from_bytes(os.urandom(4), "little")

        body = (struct.pack("<I", _MSG_INITIATION)
                + struct.pack("<I", self.sender_index)
                + eph_pub + enc_static + enc_ts)
        mac1 = _mac(_hash(_LABEL_MAC1, server_pub_raw), body)
        # mac2 reste nul tant que le serveur n'a pas envoyé de cookie : il ne
        # l'exige que sous charge.
        self.packet = body + mac1 + b"\x00" * 16
        assert len(self.packet) == _INITIATION_LEN

    def consume_response(self, packet: bytes, psk: bytes) -> bool:
        """Le message 2 se déchiffre-t-il avec cette clé pré-partagée ?"""
        eph_pub_r = packet[12:44]
        enc_empty = packet[44:60]

        c, h = self._c, self._h
        c = _kdf(c, eph_pub_r, 1)[0]
        h = _hash(h, eph_pub_r)

        peer_eph = X25519PublicKey.from_public_bytes(eph_pub_r)
        c = _kdf(c, self._eph_priv.exchange(peer_eph), 1)[0]
        c = _kdf(c, self._static_priv.exchange(peer_eph), 1)[0]

        c, tau, k = _kdf(c, psk, 3)
        h = _hash(h, tau)
        try:
            _aead_decrypt(k, 0, enc_empty, h)
            return True
        except Exception:
            return False


#: Issues possibles de la sonde.
OK_NO_PSK = "ok"                  # handshake complet, sans clé pré-partagée
OK_PSK = "ok_psk"                 # handshake complet avec la PSK configurée
PSK_REQUIRED = "psk_required"     # le serveur répond, mais la PSK ne convient pas
NO_REPLY = "no_reply"             # aucune réponse : clé, port ou filtrage
SRC_PORT_FILTERED = "src_port"    # accepté seulement si le port source = port serveur
COOKIE = "cookie"                 # le serveur exige un cookie (charge/anti-DoS)
BAD_CONFIG = "bad_config"         # clés illisibles côté client
NET_ERROR = "net_error"           # résolution ou socket en échec


def probe_handshake(server_endpoint: str, server_port: int,
                    server_public_key: str, client_private_key: str,
                    preshared_key: str = "", timeout: float = 5.0,
                    attempts: int = 2) -> tuple[str, str]:
    """Tente un handshake réel et retourne (issue, message).

    Ne monte aucun tunnel, n'écrit aucun fichier et ne demande aucun
    privilège : un simple datagramme UDP sortant.
    """
    try:
        server_pub = X25519PublicKey.from_public_bytes(
            base64.b64decode(server_public_key))
        static_priv = X25519PrivateKey.from_private_bytes(
            base64.b64decode(client_private_key))
    except Exception as e:
        return BAD_CONFIG, f"Clés illisibles : {e}"

    psk = _ZERO_PSK
    if preshared_key.strip():
        try:
            psk = base64.b64decode(preshared_key.strip())
            if len(psk) != 32:
                return BAD_CONFIG, "Clé pré-partagée : 32 octets attendus"
        except Exception:
            return BAD_CONFIG, "Clé pré-partagée : base64 invalide"

    try:
        infos = socket.getaddrinfo(server_endpoint, server_port,
                                   type=socket.SOCK_DGRAM)
    except socket.gaierror as e:
        return NET_ERROR, f"« {server_endpoint} » non résolu : {e.strerror or e}"
    family, socktype, proto, _, addr = infos[0]

    def _attempt(source_port: int | None) -> tuple[str, str] | None:
        """Une tentative. None = pas de réponse dans le délai imparti."""
        init = _Initiation(static_priv, server_pub)
        sock = socket.socket(family, socktype, proto)
        sock.settimeout(timeout)
        try:
            if source_port:
                sock.bind(("", source_port))
            start = time.monotonic()
            sock.sendto(init.packet, addr)
            while True:
                data, _src = sock.recvfrom(1024)
                ms = int((time.monotonic() - start) * 1000)
                if len(data) < 4:
                    continue
                kind = struct.unpack("<I", data[:4])[0]

                if kind == _MSG_COOKIE:
                    return COOKIE, ("Le serveur répond par un cookie : il est "
                                    "sous charge ou en protection anti-déni de "
                                    "service. La configuration n'est pas en "
                                    "cause, réessayez.")
                if kind != _MSG_RESPONSE or len(data) != _RESPONSE_LEN:
                    continue
                if struct.unpack("<I", data[8:12])[0] != init.sender_index:
                    continue  # réponse destinée à une autre session

                # Le serveur a répondu : il connaît notre clé publique et sa
                # propre clé est bien celle que nous utilisons. Reste à savoir
                # quelle clé pré-partagée il attend.
                if init.consume_response(data, psk):
                    if psk == _ZERO_PSK:
                        return OK_NO_PSK, (
                            f"Handshake complet en {ms} ms, sans clé "
                            f"pré-partagée. Le serveur accepte ce pair.")
                    return OK_PSK, (
                        f"Handshake complet en {ms} ms, clé pré-partagée "
                        f"acceptée.")

                other = _ZERO_PSK if psk != _ZERO_PSK else None
                if other is not None and init.consume_response(data, other):
                    return PSK_REQUIRED, (
                        "Le serveur accepte ce pair mais n'attend aucune clé "
                        "pré-partagée : décochez l'option, la clé configurée "
                        "fait échouer la négociation.")
                return PSK_REQUIRED, (
                    "Le serveur accepte ce pair mais attend une clé "
                    "pré-partagée différente de celle configurée (ou une clé "
                    "alors qu'aucune n'est renseignée).")
        except socket.timeout:
            return None
        except ConnectionRefusedError:
            return NO_REPLY, (f"Port UDP {server_port} fermé sur {addr[0]} : "
                              f"rien n'écoute à cette adresse.")
        except OSError as e:
            return NET_ERROR, str(e)
        finally:
            sock.close()

    for _ in range(max(1, attempts)):
        result = _attempt(None)
        if result is not None:
            return result

    # Silence complet. Avant de conclure, rejouer une fois depuis le port du
    # serveur : certains pare-feux — une règle FortiGate dont l'objet Service
    # fixe le port source, typiquement `udp-portrange 51820:51820` — n'acceptent
    # le trafic que si le port source égale le port de destination. Un client
    # WireGuard émettant depuis un port éphémère, il est rejeté avant même
    # d'atteindre le serveur, sans le moindre message.
    try:
        result = _attempt(server_port)
    except OSError:
        result = None
    if result is not None and result[0] in (OK_NO_PSK, OK_PSK, PSK_REQUIRED):
        return SRC_PORT_FILTERED, (
            f"Le serveur répond lorsque le port source vaut {server_port}, mais "
            f"reste muet depuis un port éphémère. Un équipement en chemin filtre "
            f"sur le port source : sur FortiGate, c'est l'objet Service dont le "
            f"champ « source port » est renseigné (`udp-portrange "
            f"{server_port}:{server_port}`) — videz-le. Un client WireGuard émet "
            f"depuis un port aléatoire et ne peut pas satisfaire cette règle.")

    return NO_REPLY, (
        f"Aucune réponse du serveur après {attempts} tentatives. Le message "
        f"d'initiation est pourtant valide, donc : le serveur n'écoute pas sur "
        f"le port {server_port}, il ne connaît pas notre clé publique, la clé "
        f"publique du serveur configurée ici n'est pas la sienne, ou un "
        f"pare-feu bloque l'UDP en chemin.")