vpn/app/services/crypto.py

"""WireGuard-compatible key generation and AmneziaWG config rendering."""

from __future__ import annotations

import base64
import ipaddress
import secrets
from dataclasses import dataclass

from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat, PublicFormat


@dataclass
class KeyPair:
    private_key: str
    public_key: str


def generate_keypair() -> KeyPair:
    private = X25519PrivateKey.generate()
    private_bytes = private.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
    public_bytes = private.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
    return KeyPair(
        private_key=base64.b64encode(private_bytes).decode("ascii"),
        public_key=base64.b64encode(public_bytes).decode("ascii"),
    )


def generate_preshared_key() -> str:
    return base64.b64encode(secrets.token_bytes(32)).decode("ascii")


def next_client_ip(subnet: str, used_addresses: list[str]) -> str:
    network = ipaddress.ip_network(subnet, strict=False)
    used = set()
    for addr in used_addresses:
        ip = addr.split("/")[0]
        used.add(ip)

    # .1 is reserved for server
    for host in network.hosts():
        if str(host).endswith(".0") or str(host).endswith(".255"):
            continue
        if str(host) == str(network.network_address + 1):
            continue
        if str(host) not in used:
            return f"{host}/32"
    raise RuntimeError("No free IP addresses left in subnet")


def server_address(subnet: str) -> str:
    network = ipaddress.ip_network(subnet, strict=False)
    return f"{network.network_address + 1}/{network.prefixlen}"


def render_client_config(
    *,
    protocol: str,
    client_private_key: str,
    client_address: str,
    dns: str,
    server_public_key: str,
    preshared_key: str,
    endpoint_host: str,
    endpoint_port: int,
    allowed_ips: str,
    jc: int | None = None,
    jmin: int | None = None,
    jmax: int | None = None,
    s1: int | None = None,
    s2: int | None = None,
    h1: int | None = None,
    h2: int | None = None,
    h3: int | None = None,
    h4: int | None = None,
) -> str:
    lines = [
        "[Interface]",
        f"PrivateKey = {client_private_key}",
        f"Address = {client_address}",
        f"DNS = {dns}",
    ]

    if protocol == "awg2":
        # AmneziaWG 2.0 obfuscation parameters
        lines.extend(
            [
                f"Jc = {jc if jc is not None else 4}",
                f"Jmin = {jmin if jmin is not None else 40}",
                f"Jmax = {jmax if jmax is not None else 70}",
                f"S1 = {s1 if s1 is not None else 0}",
                f"S2 = {s2 if s2 is not None else 0}",
                f"H1 = {h1 if h1 is not None else 1}",
                f"H2 = {h2 if h2 is not None else 2}",
                f"H3 = {h3 if h3 is not None else 3}",
                f"H4 = {h4 if h4 is not None else 4}",
            ]
        )

    lines.extend(
        [
            "",
            "[Peer]",
            f"PublicKey = {server_public_key}",
            f"PresharedKey = {preshared_key}",
            f"Endpoint = {endpoint_host}:{endpoint_port}",
            f"AllowedIPs = {allowed_ips}",
            "PersistentKeepalive = 25",
            "",
        ]
    )
    return "\n".join(lines)


def render_server_config(
    *,
    protocol: str,
    interface_name: str,
    private_key: str,
    address: str,
    listen_port: int,
    peers: list[dict],
    jc: int | None = None,
    jmin: int | None = None,
    jmax: int | None = None,
    s1: int | None = None,
    s2: int | None = None,
    h1: int | None = None,
    h2: int | None = None,
    h3: int | None = None,
    h4: int | None = None,
) -> str:
    lines = [
        f"# Auto-generated by VpnPanel for {interface_name}",
        "[Interface]",
        f"PrivateKey = {private_key}",
        f"Address = {address}",
        f"ListenPort = {listen_port}",
        "SaveConfig = false",
    ]

    if protocol == "awg2":
        lines.extend(
            [
                f"Jc = {jc if jc is not None else 4}",
                f"Jmin = {jmin if jmin is not None else 40}",
                f"Jmax = {jmax if jmax is not None else 70}",
                f"S1 = {s1 if s1 is not None else 0}",
                f"S2 = {s2 if s2 is not None else 0}",
                f"H1 = {h1 if h1 is not None else 1}",
                f"H2 = {h2 if h2 is not None else 2}",
                f"H3 = {h3 if h3 is not None else 3}",
                f"H4 = {h4 if h4 is not None else 4}",
            ]
        )

    for peer in peers:
        lines.extend(
            [
                "",
                f"# {peer['name']}",
                "[Peer]",
                f"PublicKey = {peer['public_key']}",
                f"PresharedKey = {peer['preshared_key']}",
                f"AllowedIPs = {peer['address']}",
            ]
        )

    lines.append("")
    return "\n".join(lines)