vpn/app/services/crypto.py

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

from __future__ import annotations

import base64
import ipaddress
import random
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


@dataclass
class AwgParams:
    jc: int
    jmin: int
    jmax: int
    s1: int
    s2: int
    s3: int
    s4: int
    h1: int
    h2: int
    h3: int
    h4: int


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 generate_awg_params() -> AwgParams:
    """Real AmneziaWG-style obfuscation params (not 1/2/3/4 placeholders)."""
    jc = random.randint(3, 10)
    jmin = random.randint(10, 40)
    jmax = random.randint(jmin + 10, jmin + 60)
    s1 = random.randint(15, 70)
    s2 = random.randint(15, 70)
    s3 = random.randint(15, 70)
    s4 = random.randint(15, 70)
    # Distinct large magic headers
    headers = random.sample(range(100_000_000, 2_000_000_000), 4)
    return AwgParams(
        jc=jc,
        jmin=jmin,
        jmax=jmax,
        s1=s1,
        s2=s2,
        s3=s3,
        s4=s4,
        h1=headers[0],
        h2=headers[1],
        h3=headers[2],
        h4=headers[3],
    )


def next_client_ip(subnet: str, used_addresses: list[str]) -> str:
    network = ipaddress.ip_network(subnet, strict=False)
    used = {addr.split("/")[0] for addr in used_addresses}
    for host in network.hosts():
        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,
    mtu: int = 1420,
    jc: int | None = None,
    jmin: int | None = None,
    jmax: int | None = None,
    s1: int | None = None,
    s2: int | None = None,
    s3: int | None = None,
    s4: int | None = None,
    h1: int | None = None,
    h2: int | None = None,
    h3: int | None = None,
    h4: int | None = None,
) -> str:
    addr = client_address if "/" in client_address else f"{client_address}/32"
    lines = [
        "[Interface]",
        f"PrivateKey = {client_private_key}",
        f"Address = {addr}",
        f"DNS = {dns}",
        f"MTU = {mtu}",
    ]

    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 15}",
                f"S2 = {s2 if s2 is not None else 18}",
                f"S3 = {s3 if s3 is not None else 20}",
                f"S4 = {s4 if s4 is not None else 23}",
                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,
    s3: int | None = None,
    s4: 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 15}",
                f"S2 = {s2 if s2 is not None else 18}",
                f"S3 = {s3 if s3 is not None else 20}",
                f"S4 = {s4 if s4 is not None else 23}",
                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:
        peer_addr = peer["address"]
        if not peer_addr.endswith("/32"):
            peer_addr = f"{peer_addr.split('/')[0]}/32"
        lines.extend(
            [
                "",
                f"# {peer['name']}",
                "[Peer]",
                f"PublicKey = {peer['public_key']}",
                f"PresharedKey = {peer['preshared_key']}",
                f"AllowedIPs = {peer_addr}",
            ]
        )

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