from __future__ import annotations import asyncio from datetime import datetime, timezone from pathlib import Path from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload from app.config import settings from app.models import Protocol, VpnClient, VpnServer from app.services.crypto import ( generate_keypair, generate_preshared_key, next_client_ip, render_client_config, render_server_config, server_address, ) from app.services.ssh import probe_tcp, verify_ssh async def ensure_default_servers(session: AsyncSession) -> None: existing = (await session.execute(select(VpnServer))).scalars().all() by_protocol = {s.protocol: s for s in existing} if Protocol.WIREGUARD.value not in by_protocol: keys = generate_keypair() private = settings.server_private_key or keys.private_key public = settings.server_public_key or keys.public_key if settings.server_private_key and settings.server_public_key: private, public = settings.server_private_key, settings.server_public_key elif settings.server_private_key: private = settings.server_private_key # derive not available without raw parse; keep generated public if only private set session.add( VpnServer( name="WireGuard", protocol=Protocol.WIREGUARD.value, public_host=settings.public_host, public_port=settings.public_port, interface_name=settings.wg_interface, subnet=settings.wg_subnet, dns=settings.wg_dns, server_private_key=private, server_public_key=public, ) ) if Protocol.AWG2.value not in by_protocol: keys = generate_keypair() private = settings.awg_server_private_key or keys.private_key public = settings.awg_server_public_key or keys.public_key session.add( VpnServer( name="AmneziaWG 2.0", protocol=Protocol.AWG2.value, public_host=settings.public_host, public_port=settings.awg_public_port, interface_name=settings.awg_interface, subnet=settings.awg_subnet, dns=settings.wg_dns, server_private_key=private, server_public_key=public, jc=4, jmin=40, jmax=70, s1=0, s2=0, h1=1, h2=2, h3=3, h4=4, ) ) await session.commit() def _defaults_for_protocol(protocol: str) -> dict: if protocol == Protocol.AWG2.value: return { "interface_name": settings.awg_interface, "subnet": settings.awg_subnet, "public_port": settings.awg_public_port, "jc": 4, "jmin": 40, "jmax": 70, "s1": 0, "s2": 0, "h1": 1, "h2": 2, "h3": 3, "h4": 4, } return { "interface_name": settings.wg_interface, "subnet": settings.wg_subnet, "public_port": settings.public_port, "jc": None, "jmin": None, "jmax": None, "s1": None, "s2": None, "h1": None, "h2": None, "h3": None, "h4": None, } async def create_server( session: AsyncSession, *, name: str, protocol: str, ssh_host: str, ssh_username: str, ssh_port: int = 22, ssh_password: str | None = None, ssh_private_key: str | None = None, public_host: str | None = None, public_port: int | None = None, dns: str = "1.1.1.1", ) -> VpnServer: protocol = protocol.strip().lower() if protocol not in {Protocol.WIREGUARD.value, Protocol.AWG2.value}: raise ValueError("Unsupported protocol") password = (ssh_password or "").strip() or None private_key = (ssh_private_key or "").strip() or None if not password and not private_key: raise ValueError("Укажите пароль SSH или приватный ключ") host = ssh_host.strip() if not host: raise ValueError("Укажите SSH host") username = ssh_username.strip() if not username: raise ValueError("Укажите SSH логин") # Like Amnezia-Web-Panel: verify SSH before saving try: server_info = await asyncio.to_thread( verify_ssh, host, ssh_port or 22, username, password, private_key ) except Exception as exc: # noqa: BLE001 raise ValueError(f"Подключение не удалось: {exc}") from exc defaults = _defaults_for_protocol(protocol) keys = generate_keypair() # Unique interface per server id is assigned after insert; use temp then update existing_count = await session.scalar(select(func.count()).select_from(VpnServer)) or 0 iface_base = "awg" if protocol == Protocol.AWG2.value else "wg" interface_name = f"{iface_base}{existing_count}" # Unique subnet per server: 10.{8+n}.0.0/24 or 10.{9+n}.0.0/24 octet = (8 if protocol == Protocol.WIREGUARD.value else 9) + int(existing_count) if octet > 250: raise ValueError("Слишком много серверов для авто-подсети") subnet = f"10.{octet}.0.0/24" now = datetime.now(timezone.utc) server = VpnServer( name=name.strip() or host, protocol=protocol, public_host=(public_host or host).strip(), public_port=public_port or int(defaults["public_port"]), interface_name=interface_name, subnet=subnet, dns=(dns or "1.1.1.1").strip(), server_private_key=keys.private_key, server_public_key=keys.public_key, ssh_host=host, ssh_port=ssh_port or 22, ssh_username=username, ssh_password=password, ssh_private_key=private_key, ssh_server_info=server_info, ssh_reachable=True, ssh_last_checked=now, jc=defaults["jc"], jmin=defaults["jmin"], jmax=defaults["jmax"], s1=defaults["s1"], s2=defaults["s2"], h1=defaults["h1"], h2=defaults["h2"], h3=defaults["h3"], h4=defaults["h4"], is_enabled=True, ) session.add(server) await session.commit() await session.refresh(server) await sync_server_config(session, server.id) return server async def check_server_ssh(session: AsyncSession, server_id: int) -> VpnServer: server = await session.get(VpnServer, server_id) if not server: raise ValueError("Server not found") if not server.ssh_host or not server.ssh_username: raise ValueError("SSH данные не заданы") if not server.ssh_password and not server.ssh_private_key: raise ValueError("Нет пароля и ключа SSH") now = datetime.now(timezone.utc) try: info = await asyncio.to_thread( verify_ssh, server.ssh_host, server.ssh_port or 22, server.ssh_username, server.ssh_password, server.ssh_private_key, ) server.ssh_server_info = info server.ssh_reachable = True server.ssh_last_checked = now except Exception as exc: # noqa: BLE001 server.ssh_reachable = False server.ssh_last_checked = now await session.commit() raise ValueError(f"Подключение не удалось: {exc}") from exc await session.commit() await session.refresh(server) return server async def probe_servers(servers: list[VpnServer]) -> dict[int, dict]: """TCP ping SSH ports for UI online indicators.""" async def one(s: VpnServer) -> tuple[int, dict]: if not s.ssh_host: return s.id, {"online": False, "latency_ms": None} result = await asyncio.to_thread(probe_tcp, s.ssh_host, s.ssh_port or 22) return s.id, {"online": result.online, "latency_ms": result.latency_ms} pairs = await asyncio.gather(*(one(s) for s in servers)) return dict(pairs) async def update_server_settings( session: AsyncSession, server_id: int, *, name: str | None = None, public_host: str, public_port: int, dns: str = "1.1.1.1", ssh_host: str | None = None, ssh_port: int = 22, ssh_username: str | None = None, ssh_password: str | None = None, ssh_private_key: str | None = None, clear_ssh_password: bool = False, clear_ssh_private_key: bool = False, ) -> VpnServer: server = await session.get(VpnServer, server_id) if not server: raise ValueError("Server not found") if name is not None and name.strip(): server.name = name.strip() server.public_host = public_host.strip() server.public_port = public_port server.dns = dns.strip() or "1.1.1.1" if ssh_host is not None: server.ssh_host = ssh_host.strip() or None server.ssh_port = ssh_port or 22 if ssh_username is not None: server.ssh_username = ssh_username.strip() or None if clear_ssh_password: server.ssh_password = None elif ssh_password is not None and ssh_password.strip(): server.ssh_password = ssh_password.strip() if clear_ssh_private_key: server.ssh_private_key = None elif ssh_private_key is not None and ssh_private_key.strip(): server.ssh_private_key = ssh_private_key.strip() await session.commit() await session.refresh(server) await sync_server_config(session, server_id) return server async def create_client( session: AsyncSession, *, server_id: int, name: str, notes: str | None = None, allowed_ips: str = "0.0.0.0/0, ::/0", ) -> VpnClient: server = await session.get(VpnServer, server_id) if not server: raise ValueError("Server not found") used = ( await session.execute(select(VpnClient.address).where(VpnClient.server_id == server_id)) ).scalars().all() address = next_client_ip(server.subnet, list(used)) keys = generate_keypair() client = VpnClient( server_id=server_id, name=name.strip(), private_key=keys.private_key, public_key=keys.public_key, preshared_key=generate_preshared_key(), address=address, allowed_ips=allowed_ips, notes=notes, is_enabled=True, ) session.add(client) await session.commit() await session.refresh(client) await sync_server_config(session, server_id) return client async def set_client_enabled(session: AsyncSession, client_id: int, enabled: bool) -> VpnClient: client = await session.get(VpnClient, client_id) if not client: raise ValueError("Client not found") client.is_enabled = enabled await session.commit() await session.refresh(client) await sync_server_config(session, client.server_id) return client async def delete_client(session: AsyncSession, client_id: int) -> None: client = await session.get(VpnClient, client_id) if not client: raise ValueError("Client not found") server_id = client.server_id await session.delete(client) await session.commit() await sync_server_config(session, server_id) async def get_client_config(session: AsyncSession, client_id: int) -> str: result = await session.execute( select(VpnClient).options(selectinload(VpnClient.server)).where(VpnClient.id == client_id) ) client = result.scalar_one_or_none() if not client: raise ValueError("Client not found") server = client.server return render_client_config( protocol=server.protocol, client_private_key=client.private_key, client_address=client.address, dns=server.dns, server_public_key=server.server_public_key, preshared_key=client.preshared_key, endpoint_host=server.public_host, endpoint_port=server.public_port, allowed_ips=client.allowed_ips, jc=server.jc, jmin=server.jmin, jmax=server.jmax, s1=server.s1, s2=server.s2, h1=server.h1, h2=server.h2, h3=server.h3, h4=server.h4, ) async def sync_server_config(session: AsyncSession, server_id: int) -> Path: result = await session.execute( select(VpnServer).options(selectinload(VpnServer.clients)).where(VpnServer.id == server_id) ) server = result.scalar_one() peers = [ { "name": c.name, "public_key": c.public_key, "preshared_key": c.preshared_key, "address": c.address if c.address.endswith("/32") else f"{c.address.split('/')[0]}/32", } for c in server.clients if c.is_enabled ] content = render_server_config( protocol=server.protocol, interface_name=server.interface_name, private_key=server.server_private_key, address=server_address(server.subnet), listen_port=server.public_port, peers=peers, jc=server.jc, jmin=server.jmin, jmax=server.jmax, s1=server.s1, s2=server.s2, h1=server.h1, h2=server.h2, h3=server.h3, h4=server.h4, ) base = Path(settings.wg_config_dir if server.protocol == Protocol.WIREGUARD.value else settings.awg_config_dir) base.mkdir(parents=True, exist_ok=True) path = base / f"{server.interface_name}.conf" path.write_text(content, encoding="utf-8") return path