WGSecure/app/core/wireguard.py

import base64
import os
import socket
import time
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from app.core.config import Config
from app.utils.platform_utils import (
    get_config_dir,
    get_wg_config_dir,
    is_windows,
    run_command,
    run_privileged,
    wg_available,
)


def generate_keypair() -> tuple[str, str]:
    """Retourne (private_key_b64, public_key_b64) au format WireGuard."""
    private = X25519PrivateKey.generate()
    private_bytes = private.private_bytes_raw()
    public_bytes = private.public_key().public_bytes_raw()
    return (
        base64.b64encode(private_bytes).decode(),
        base64.b64encode(public_bytes).decode(),
    )


def generate_preshared_key() -> str:
    return base64.b64encode(os.urandom(32)).decode()


def build_client_config(cfg: Config) -> str:
    wg = cfg.wg
    lines = [
        "[Interface]",
        f"PrivateKey = {wg['client_private_key']}",
        f"Address = {wg['client_address']}",
        f"DNS = {wg['dns']}",
        "",
        "[Peer]",
        f"PublicKey = {wg['server_public_key']}",
        f"AllowedIPs = {wg['allowed_ips']}",
        f"Endpoint = {wg['server_endpoint']}:{wg['server_port']}",
        f"PersistentKeepalive = {wg['keepalive']}",
    ]
    return "\n".join(lines) + "\n"


def get_client_config_path(cfg: Config) -> str:
    name = cfg.wg.get("interface_name", "wgs0")
    if is_windows():
        config_dir = get_config_dir()
    else:
        config_dir = get_wg_config_dir()
    return os.path.join(config_dir, f"{name}.conf")


def write_client_config(cfg: Config) -> tuple[bool, str]:
    content = build_client_config(cfg)
    path = get_client_config_path(cfg)
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as f:
            f.write(content)
        if not is_windows():
            os.chmod(path, 0o600)
        return True, path
    except PermissionError:
        # Écriture via sudo sur Linux
        config_dir_path = get_config_dir()
        tmp = os.path.join(config_dir_path, "wgs_tmp.conf")
        with open(tmp, "w") as f:
            f.write(content)
        code, _, err = run_privileged(["cp", tmp, path])
        os.unlink(tmp)
        if code == 0:
            run_privileged(["chmod", "600", path])
            return True, path
        return False, err


def is_connected(cfg: Config) -> bool:
    name = cfg.wg.get("interface_name", "wgs0")
    if is_windows():
        code, out, _ = run_command(["sc", "query", f"WireGuardTunnel${name}"])
        return code == 0 and "RUNNING" in out
    code, out, _ = run_command(["wg", "show", name])
    return code == 0 and bool(out)


def connect(cfg: Config) -> tuple[bool, str]:
    if not cfg.configured:
        return False, "WireGuard non configuré. Ouvrez le panneau Admin."

    ok, result = write_client_config(cfg)
    if not ok:
        return False, f"Impossible d'écrire la config : {result}"

    name = cfg.wg.get("interface_name", "wgs0")

    if is_windows():
        code, _, err = run_command(["wireguard", "/installtunnel", result])
        return (code == 0), (err or "Connecté")
    else:
        code, _, err = run_privileged(["wg-quick", "up", result])
        if code == 0:
            return True, "Tunnel WireGuard activé"
        return False, err or "Erreur lors de la connexion"


def disconnect(cfg: Config) -> tuple[bool, str]:
    name = cfg.wg.get("interface_name", "wgs0")

    if is_windows():
        code, _, err = run_command(["wireguard", "/uninstalltunnel", name])
        return (code == 0), (err or "Déconnecté")
    else:
        path = get_client_config_path(cfg)
        if not os.path.exists(path):
            path = name
        code, _, err = run_privileged(["wg-quick", "down", path])
        if code == 0:
            return True, "Tunnel WireGuard désactivé"
        return False, err or "Erreur lors de la déconnexion"


def test_connection(cfg: Config, timeout: int = 5) -> tuple[bool, str]:
    endpoint = cfg.wg.get("server_endpoint", "")
    port = int(cfg.wg.get("server_port", 51820))
    if not endpoint:
        return False, "Aucun serveur configuré"
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(timeout)
        start = time.monotonic()
        sock.connect((endpoint, port))
        sock.close()
        latency = int((time.monotonic() - start) * 1000)
        return True, f"Serveur joignable ({latency} ms)"
    except socket.timeout:
        return False, "Timeout — serveur inaccessible"
    except OSError as e:
        return False, str(e)


def get_status_info(cfg: Config) -> dict:
    name = cfg.wg.get("interface_name", "wgs0")
    info = {
        "connected": False,
        "interface": name,
        "peer": "",
        "rx_bytes": 0,
        "tx_bytes": 0,
        "last_handshake": "",
    }
    if is_windows():
        info["connected"] = is_connected(cfg)
        return info

    code, out, _ = run_command(["wg", "show", name])
    if code != 0 or not out:
        return info

    info["connected"] = True
    for line in out.splitlines():
        line = line.strip()
        if line.startswith("peer:"):
            info["peer"] = line.split(":", 1)[1].strip()[:16] + "…"
        elif line.startswith("transfer:"):
            parts = line.split(":", 1)[1].strip().split(",")
            try:
                rx = parts[0].strip().split()[0]
                tx = parts[1].strip().split()[0]
                info["rx_bytes"] = rx
                info["tx_bytes"] = tx
            except Exception:
                pass
        elif line.startswith("latest handshake:"):
            info["last_handshake"] = line.split(":", 1)[1].strip()

    return info