dcrouter/ts/vpn/classes.vpn-manager.ts

import * as plugins from '../plugins.js';
import { logger } from '../logger.js';
import type { StorageManager } from '../storage/classes.storagemanager.js';

const STORAGE_PREFIX_KEYS = '/vpn/server-keys';
const STORAGE_PREFIX_CLIENTS = '/vpn/clients/';

export interface IVpnManagerConfig {
  /** VPN subnet CIDR (default: '10.8.0.0/24') */
  subnet?: string;
  /** WireGuard UDP listen port (default: 51820) */
  wgListenPort?: number;
  /** DNS servers pushed to VPN clients */
  dns?: string[];
  /** Server endpoint hostname for client configs (e.g. 'vpn.example.com') */
  serverEndpoint?: string;
  /** Override forwarding mode. Default: auto-detect (tun if root, socket otherwise) */
  forwardingMode?: 'tun' | 'socket';
}

interface IPersistedServerKeys {
  noisePrivateKey: string;
  noisePublicKey: string;
  wgPrivateKey: string;
  wgPublicKey: string;
}

interface IPersistedClient {
  clientId: string;
  enabled: boolean;
  tags?: string[];
  description?: string;
  assignedIp?: string;
  noisePublicKey: string;
  wgPublicKey: string;
  createdAt: number;
  updatedAt: number;
  expiresAt?: string;
}

/**
 * Manages the SmartVPN server lifecycle and VPN client CRUD.
 * Persists server keys and client registrations via StorageManager.
 */
export class VpnManager {
  private storageManager: StorageManager;
  private config: IVpnManagerConfig;
  private vpnServer?: plugins.smartvpn.VpnServer;
  private clients: Map<string, IPersistedClient> = new Map();
  private serverKeys?: IPersistedServerKeys;
  private _forwardingMode: 'tun' | 'socket';

  constructor(storageManager: StorageManager, config: IVpnManagerConfig) {
    this.storageManager = storageManager;
    this.config = config;
    // Auto-detect forwarding mode: tun if root, socket otherwise
    this._forwardingMode = config.forwardingMode
      ?? (process.getuid?.() === 0 ? 'tun' : 'socket');
  }

  /** The effective forwarding mode (tun or socket). */
  public get forwardingMode(): 'tun' | 'socket' {
    return this._forwardingMode;
  }

  /** The VPN subnet CIDR. */
  public getSubnet(): string {
    return this.config.subnet || '10.8.0.0/24';
  }

  /** Whether the VPN server is running. */
  public get running(): boolean {
    return this.vpnServer?.running ?? false;
  }

  /**
   * Start the VPN server.
   * Loads or generates server keys, loads persisted clients, starts VpnServer.
   */
  public async start(): Promise<void> {
    // Load or generate server keys
    this.serverKeys = await this.loadOrGenerateServerKeys();

    // Load persisted clients
    await this.loadPersistedClients();

    // Build client entries for the daemon
    const clientEntries: plugins.smartvpn.IClientEntry[] = [];
    for (const client of this.clients.values()) {
      clientEntries.push({
        clientId: client.clientId,
        publicKey: client.noisePublicKey,
        wgPublicKey: client.wgPublicKey,
        enabled: client.enabled,
        tags: client.tags,
        description: client.description,
        assignedIp: client.assignedIp,
        expiresAt: client.expiresAt,
      });
    }

    const subnet = this.getSubnet();
    const wgListenPort = this.config.wgListenPort ?? 51820;

    // Create and start VpnServer
    this.vpnServer = new plugins.smartvpn.VpnServer({
      transport: { transport: 'stdio' },
    });

    const serverConfig: plugins.smartvpn.IVpnServerConfig = {
      listenAddr: '0.0.0.0:0', // WS listener not strictly needed but required field
      privateKey: this.serverKeys.noisePrivateKey,
      publicKey: this.serverKeys.noisePublicKey,
      subnet,
      dns: this.config.dns,
      forwardingMode: this._forwardingMode,
      transportMode: 'all',
      wgPrivateKey: this.serverKeys.wgPrivateKey,
      wgListenPort,
      clients: clientEntries,
      socketForwardProxyProtocol: this._forwardingMode === 'socket',
    };

    await this.vpnServer.start(serverConfig);
    logger.log('info', `VPN server started: mode=${this._forwardingMode}, subnet=${subnet}, wg=:${wgListenPort}, clients=${this.clients.size}`);
  }

  /**
   * Stop the VPN server.
   */
  public async stop(): Promise<void> {
    if (this.vpnServer) {
      try {
        await this.vpnServer.stopServer();
      } catch {
        // Ignore stop errors
      }
      this.vpnServer.stop();
      this.vpnServer = undefined;
    }
    logger.log('info', 'VPN server stopped');
  }

  // ── Client CRUD ────────────────────────────────────────────────────────

  /**
   * Create a new VPN client. Returns the config bundle (secrets only shown once).
   */
  public async createClient(opts: {
    clientId: string;
    tags?: string[];
    description?: string;
  }): Promise<plugins.smartvpn.IClientConfigBundle> {
    if (!this.vpnServer) {
      throw new Error('VPN server not running');
    }

    const bundle = await this.vpnServer.createClient({
      clientId: opts.clientId,
      tags: opts.tags,
      description: opts.description,
    });

    // Update WireGuard config endpoint if serverEndpoint is configured
    if (this.config.serverEndpoint && bundle.wireguardConfig) {
      const wgPort = this.config.wgListenPort ?? 51820;
      bundle.wireguardConfig = bundle.wireguardConfig.replace(
        /Endpoint\s*=\s*.+/,
        `Endpoint = ${this.config.serverEndpoint}:${wgPort}`,
      );
    }

    // Persist client entry (without private keys)
    const persisted: IPersistedClient = {
      clientId: bundle.entry.clientId,
      enabled: bundle.entry.enabled ?? true,
      tags: bundle.entry.tags,
      description: bundle.entry.description,
      assignedIp: bundle.entry.assignedIp,
      noisePublicKey: bundle.entry.publicKey,
      wgPublicKey: bundle.entry.wgPublicKey || '',
      createdAt: Date.now(),
      updatedAt: Date.now(),
      expiresAt: bundle.entry.expiresAt,
    };
    this.clients.set(persisted.clientId, persisted);
    await this.persistClient(persisted);

    return bundle;
  }

  /**
   * Remove a VPN client.
   */
  public async removeClient(clientId: string): Promise<void> {
    if (!this.vpnServer) {
      throw new Error('VPN server not running');
    }
    await this.vpnServer.removeClient(clientId);
    this.clients.delete(clientId);
    await this.storageManager.delete(`${STORAGE_PREFIX_CLIENTS}${clientId}`);
  }

  /**
   * List all registered clients (without secrets).
   */
  public listClients(): IPersistedClient[] {
    return [...this.clients.values()];
  }

  /**
   * Enable a client.
   */
  public async enableClient(clientId: string): Promise<void> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    await this.vpnServer.enableClient(clientId);
    const client = this.clients.get(clientId);
    if (client) {
      client.enabled = true;
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }
  }

  /**
   * Disable a client.
   */
  public async disableClient(clientId: string): Promise<void> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    await this.vpnServer.disableClient(clientId);
    const client = this.clients.get(clientId);
    if (client) {
      client.enabled = false;
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }
  }

  /**
   * Rotate a client's keys. Returns the new config bundle.
   */
  public async rotateClientKey(clientId: string): Promise<plugins.smartvpn.IClientConfigBundle> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    const bundle = await this.vpnServer.rotateClientKey(clientId);

    // Update endpoint in WireGuard config
    if (this.config.serverEndpoint && bundle.wireguardConfig) {
      const wgPort = this.config.wgListenPort ?? 51820;
      bundle.wireguardConfig = bundle.wireguardConfig.replace(
        /Endpoint\s*=\s*.+/,
        `Endpoint = ${this.config.serverEndpoint}:${wgPort}`,
      );
    }

    // Update persisted entry with new public keys
    const client = this.clients.get(clientId);
    if (client) {
      client.noisePublicKey = bundle.entry.publicKey;
      client.wgPublicKey = bundle.entry.wgPublicKey || '';
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }

    return bundle;
  }

  /**
   * Export a client config (without secrets).
   */
  public async exportClientConfig(clientId: string, format: 'smartvpn' | 'wireguard'): Promise<string> {
    if (!this.vpnServer) throw new Error('VPN server not running');
    let config = await this.vpnServer.exportClientConfig(clientId, format);

    // Update endpoint in WireGuard config
    if (format === 'wireguard' && this.config.serverEndpoint) {
      const wgPort = this.config.wgListenPort ?? 51820;
      config = config.replace(
        /Endpoint\s*=\s*.+/,
        `Endpoint = ${this.config.serverEndpoint}:${wgPort}`,
      );
    }

    return config;
  }

  // ── Status and telemetry ───────────────────────────────────────────────

  /**
   * Get server status.
   */
  public async getStatus(): Promise<plugins.smartvpn.IVpnStatus | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getStatus();
  }

  /**
   * Get server statistics.
   */
  public async getStatistics(): Promise<plugins.smartvpn.IVpnServerStatistics | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getStatistics();
  }

  /**
   * List currently connected clients.
   */
  public async getConnectedClients(): Promise<plugins.smartvpn.IVpnClientInfo[]> {
    if (!this.vpnServer) return [];
    return this.vpnServer.listClients();
  }

  /**
   * Get telemetry for a specific client.
   */
  public async getClientTelemetry(clientId: string): Promise<plugins.smartvpn.IVpnClientTelemetry | null> {
    if (!this.vpnServer) return null;
    return this.vpnServer.getClientTelemetry(clientId);
  }

  /**
   * Get server public keys (for display/info).
   */
  public getServerPublicKeys(): { noisePublicKey: string; wgPublicKey: string } | null {
    if (!this.serverKeys) return null;
    return {
      noisePublicKey: this.serverKeys.noisePublicKey,
      wgPublicKey: this.serverKeys.wgPublicKey,
    };
  }

  // ── Private helpers ────────────────────────────────────────────────────

  private async loadOrGenerateServerKeys(): Promise<IPersistedServerKeys> {
    const stored = await this.storageManager.getJSON<IPersistedServerKeys>(STORAGE_PREFIX_KEYS);
    if (stored?.noisePrivateKey && stored?.wgPrivateKey) {
      logger.log('info', 'Loaded VPN server keys from storage');
      return stored;
    }

    // Generate new keys via the daemon
    const tempServer = new plugins.smartvpn.VpnServer({
      transport: { transport: 'stdio' },
    });
    await tempServer.start();

    const noiseKeys = await tempServer.generateKeypair();
    const wgKeys = await tempServer.generateWgKeypair();
    tempServer.stop();

    const keys: IPersistedServerKeys = {
      noisePrivateKey: noiseKeys.privateKey,
      noisePublicKey: noiseKeys.publicKey,
      wgPrivateKey: wgKeys.privateKey,
      wgPublicKey: wgKeys.publicKey,
    };

    await this.storageManager.setJSON(STORAGE_PREFIX_KEYS, keys);
    logger.log('info', 'Generated and persisted new VPN server keys');
    return keys;
  }

  private async loadPersistedClients(): Promise<void> {
    const keys = await this.storageManager.list(STORAGE_PREFIX_CLIENTS);
    for (const key of keys) {
      const client = await this.storageManager.getJSON<IPersistedClient>(key);
      if (client) {
        this.clients.set(client.clientId, client);
      }
    }
    if (this.clients.size > 0) {
      logger.log('info', `Loaded ${this.clients.size} persisted VPN client(s)`);
    }
  }

  private async persistClient(client: IPersistedClient): Promise<void> {
    await this.storageManager.setJSON(`${STORAGE_PREFIX_CLIENTS}${client.clientId}`, client);
  }
}