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

import * as plugins from '../plugins.js';
import { logger } from '../logger.js';
import { VpnServerKeysDoc, VpnClientDoc } from '../db/index.js';

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;
  /** Pre-defined VPN clients created on startup (idempotent — skips already-persisted clients) */
  initialClients?: Array<{
    clientId: string;
    serverDefinedClientTags?: string[];
    description?: string;
  }>;
  /** Called when clients are created/deleted/toggled — triggers route re-application */
  onClientChanged?: () => void;
  /** Destination routing policy override. Default: forceTarget to 127.0.0.1 */
  destinationPolicy?: {
    default: 'forceTarget' | 'block' | 'allow';
    target?: string;
    allowList?: string[];
    blockList?: string[];
  };
  /** Compute per-client AllowedIPs based on the client's server-defined tags.
   *  Called at config generation time (create/export). Returns CIDRs for WireGuard AllowedIPs.
   *  When not set, defaults to [subnet]. */
  getClientAllowedIPs?: (clientTags: string[]) => Promise<string[]>;
  /** Forwarding mode: 'socket' (default, userspace NAT), 'bridge' (L2 bridge to host LAN),
   *  or 'hybrid' (socket default, bridge for clients with useHostIp=true) */
  forwardingMode?: 'socket' | 'bridge' | 'hybrid';
  /** LAN subnet CIDR for bridge mode (e.g., '192.168.1.0/24') */
  bridgeLanSubnet?: string;
  /** Physical network interface for bridge mode (auto-detected if omitted) */
  bridgePhysicalInterface?: string;
  /** Start of VPN client IP range in LAN subnet (host offset, default: 200) */
  bridgeIpRangeStart?: number;
  /** End of VPN client IP range in LAN subnet (host offset, default: 250) */
  bridgeIpRangeEnd?: number;
}

/**
 * Manages the SmartVPN server lifecycle and VPN client CRUD.
 * Persists server keys and client registrations via smartdata document classes.
 */
export class VpnManager {
  private config: IVpnManagerConfig;
  private vpnServer?: plugins.smartvpn.VpnServer;
  private clients: Map<string, VpnClientDoc> = new Map();
  private serverKeys?: VpnServerKeysDoc;

  constructor(config: IVpnManagerConfig) {
    this.config = config;
  }

  /** 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[] = [];
    let anyClientUsesHostIp = false;
    for (const client of this.clients.values()) {
      if (client.useHostIp) {
        anyClientUsesHostIp = true;
      }
      const entry: plugins.smartvpn.IClientEntry = {
        clientId: client.clientId,
        publicKey: client.noisePublicKey,
        wgPublicKey: client.wgPublicKey,
        enabled: client.enabled,
        serverDefinedClientTags: client.serverDefinedClientTags,
        description: client.description,
        assignedIp: client.assignedIp,
        expiresAt: client.expiresAt,
        security: this.buildClientSecurity(client),
      };
      // Pass per-client bridge fields if present (for hybrid/bridge mode)
      if (client.useHostIp !== undefined) (entry as any).useHostIp = client.useHostIp;
      if (client.useDhcp !== undefined) (entry as any).useDhcp = client.useDhcp;
      if (client.staticIp !== undefined) (entry as any).staticIp = client.staticIp;
      if (client.forceVlan !== undefined) (entry as any).forceVlan = client.forceVlan;
      if (client.vlanId !== undefined) (entry as any).vlanId = client.vlanId;
      clientEntries.push(entry);
    }

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

    // Auto-detect hybrid mode: if any persisted client uses host IP and mode is
    // 'socket' (or unset), upgrade to 'hybrid' so the daemon can handle both
    let configuredMode = this.config.forwardingMode ?? 'socket';
    if (anyClientUsesHostIp && configuredMode === 'socket') {
      configuredMode = 'hybrid';
      logger.log('info', 'VPN: Auto-upgrading forwarding mode to hybrid (client with useHostIp detected)');
    }
    const forwardingMode = configuredMode === 'hybrid' ? 'hybrid' : configuredMode;
    const isBridge = forwardingMode === 'bridge';

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

    // Default destination policy: bridge mode allows traffic through directly,
    // socket mode forces traffic to SmartProxy on 127.0.0.1
    const defaultDestinationPolicy: plugins.smartvpn.IDestinationPolicy = isBridge
      ? { default: 'allow' as const }
      : { default: 'forceTarget' as const, target: '127.0.0.1' };

    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: forwardingMode as any,
      transportMode: 'all',
      wgPrivateKey: this.serverKeys.wgPrivateKey,
      wgListenPort,
      clients: clientEntries,
      socketForwardProxyProtocol: !isBridge,
      destinationPolicy: this.config.destinationPolicy ?? defaultDestinationPolicy,
      serverEndpoint: this.config.serverEndpoint
        ? `${this.config.serverEndpoint}:${wgListenPort}`
        : undefined,
      clientAllowedIPs: [subnet],
      // Bridge-specific config
      ...(isBridge ? {
        bridgeLanSubnet: this.config.bridgeLanSubnet,
        bridgePhysicalInterface: this.config.bridgePhysicalInterface,
        bridgeIpRangeStart: this.config.bridgeIpRangeStart,
        bridgeIpRangeEnd: this.config.bridgeIpRangeEnd,
      } : {}),
    };

    await this.vpnServer.start(serverConfig);

    // Create initial clients from config (idempotent — skip already-persisted)
    if (this.config.initialClients) {
      for (const initial of this.config.initialClients) {
        if (!this.clients.has(initial.clientId)) {
          const bundle = await this.createClient({
            clientId: initial.clientId,
            serverDefinedClientTags: initial.serverDefinedClientTags,
            description: initial.description,
          });
          logger.log('info', `VPN: Created initial client '${initial.clientId}' (IP: ${bundle.entry.assignedIp})`);
        }
      }
    }

    logger.log('info', `VPN server started: 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;
    serverDefinedClientTags?: string[];
    description?: string;
    forceDestinationSmartproxy?: boolean;
    destinationAllowList?: string[];
    destinationBlockList?: string[];
    useHostIp?: boolean;
    useDhcp?: boolean;
    staticIp?: string;
    forceVlan?: boolean;
    vlanId?: number;
  }): Promise<plugins.smartvpn.IClientConfigBundle> {
    if (!this.vpnServer) {
      throw new Error('VPN server not running');
    }

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

    // Override AllowedIPs with per-client values based on tag-matched routes
    if (this.config.getClientAllowedIPs && bundle.wireguardConfig) {
      const allowedIPs = await this.config.getClientAllowedIPs(opts.serverDefinedClientTags || []);
      bundle.wireguardConfig = bundle.wireguardConfig.replace(
        /AllowedIPs\s*=\s*.+/,
        `AllowedIPs = ${allowedIPs.join(', ')}`,
      );
    }

    // Persist client entry (including WG private key for export/QR)
    const doc = new VpnClientDoc();
    doc.clientId = bundle.entry.clientId;
    doc.enabled = bundle.entry.enabled ?? true;
    doc.serverDefinedClientTags = bundle.entry.serverDefinedClientTags;
    doc.description = bundle.entry.description;
    doc.assignedIp = bundle.entry.assignedIp;
    doc.noisePublicKey = bundle.entry.publicKey;
    doc.wgPublicKey = bundle.entry.wgPublicKey || '';
    doc.wgPrivateKey = bundle.secrets?.wgPrivateKey
      || bundle.wireguardConfig?.match(/PrivateKey\s*=\s*(.+)/)?.[1]?.trim();
    doc.createdAt = Date.now();
    doc.updatedAt = Date.now();
    doc.expiresAt = bundle.entry.expiresAt;
    if (opts.forceDestinationSmartproxy !== undefined) {
      doc.forceDestinationSmartproxy = opts.forceDestinationSmartproxy;
    }
    if (opts.destinationAllowList !== undefined) {
      doc.destinationAllowList = opts.destinationAllowList;
    }
    if (opts.destinationBlockList !== undefined) {
      doc.destinationBlockList = opts.destinationBlockList;
    }
    if (opts.useHostIp !== undefined) {
      doc.useHostIp = opts.useHostIp;
    }
    if (opts.useDhcp !== undefined) {
      doc.useDhcp = opts.useDhcp;
    }
    if (opts.staticIp !== undefined) {
      doc.staticIp = opts.staticIp;
    }
    if (opts.forceVlan !== undefined) {
      doc.forceVlan = opts.forceVlan;
    }
    if (opts.vlanId !== undefined) {
      doc.vlanId = opts.vlanId;
    }
    this.clients.set(doc.clientId, doc);
    await this.persistClient(doc);

    // Sync per-client security to the running daemon
    const security = this.buildClientSecurity(doc);
    if (security.destinationPolicy) {
      await this.vpnServer!.updateClient(doc.clientId, { security });
    }

    this.config.onClientChanged?.();
    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);
    const doc = this.clients.get(clientId);
    this.clients.delete(clientId);
    if (doc) {
      await doc.delete();
    }
    this.config.onClientChanged?.();
  }

  /**
   * List all registered clients (without secrets).
   */
  public listClients(): VpnClientDoc[] {
    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);
    }
    this.config.onClientChanged?.();
  }

  /**
   * 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);
    }
    this.config.onClientChanged?.();
  }

  /**
   * Update a client's metadata (description, tags) without rotating keys.
   */
  public async updateClient(clientId: string, update: {
    description?: string;
    serverDefinedClientTags?: string[];
    forceDestinationSmartproxy?: boolean;
    destinationAllowList?: string[];
    destinationBlockList?: string[];
    useHostIp?: boolean;
    useDhcp?: boolean;
    staticIp?: string;
    forceVlan?: boolean;
    vlanId?: number;
  }): Promise<void> {
    const client = this.clients.get(clientId);
    if (!client) throw new Error(`Client not found: ${clientId}`);
    if (update.description !== undefined) client.description = update.description;
    if (update.serverDefinedClientTags !== undefined) client.serverDefinedClientTags = update.serverDefinedClientTags;
    if (update.forceDestinationSmartproxy !== undefined) client.forceDestinationSmartproxy = update.forceDestinationSmartproxy;
    if (update.destinationAllowList !== undefined) client.destinationAllowList = update.destinationAllowList;
    if (update.destinationBlockList !== undefined) client.destinationBlockList = update.destinationBlockList;
    if (update.useHostIp !== undefined) client.useHostIp = update.useHostIp;
    if (update.useDhcp !== undefined) client.useDhcp = update.useDhcp;
    if (update.staticIp !== undefined) client.staticIp = update.staticIp;
    if (update.forceVlan !== undefined) client.forceVlan = update.forceVlan;
    if (update.vlanId !== undefined) client.vlanId = update.vlanId;
    client.updatedAt = Date.now();
    await this.persistClient(client);

    // Sync per-client security to the running daemon
    if (this.vpnServer) {
      const security = this.buildClientSecurity(client);
      await this.vpnServer.updateClient(clientId, { security });
    }

    this.config.onClientChanged?.();
  }

  /**
   * 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 persisted entry with new keys (including private key for export/QR)
    const client = this.clients.get(clientId);
    if (client) {
      client.noisePublicKey = bundle.entry.publicKey;
      client.wgPublicKey = bundle.entry.wgPublicKey || '';
      client.wgPrivateKey = bundle.secrets?.wgPrivateKey
        || bundle.wireguardConfig?.match(/PrivateKey\s*=\s*(.+)/)?.[1]?.trim();
      client.updatedAt = Date.now();
      await this.persistClient(client);
    }

    return bundle;
  }

  /**
   * Export a client config. Injects stored WG private key and per-client AllowedIPs.
   */
  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);

    if (format === 'wireguard') {
      const persisted = this.clients.get(clientId);

      // Inject stored WG private key so exports produce valid, scannable configs
      if (persisted?.wgPrivateKey) {
        config = config.replace(
          '[Interface]\n',
          `[Interface]\nPrivateKey = ${persisted.wgPrivateKey}\n`,
        );
      }

      // Override AllowedIPs with per-client values based on tag-matched routes
      if (this.config.getClientAllowedIPs) {
        const clientTags = persisted?.serverDefinedClientTags || [];
        const allowedIPs = await this.config.getClientAllowedIPs(clientTags);
        config = config.replace(
          /AllowedIPs\s*=\s*.+/,
          `AllowedIPs = ${allowedIPs.join(', ')}`,
        );
      }
    }

    return config;
  }

  // ── Tag-based access control ───────────────────────────────────────────

  /**
   * Get assigned IPs for all enabled clients matching any of the given server-defined tags.
   */
  public getClientIpsForServerDefinedTags(tags: string[]): string[] {
    const ips: string[] = [];
    for (const client of this.clients.values()) {
      if (!client.enabled || !client.assignedIp) continue;
      if (client.serverDefinedClientTags?.some(t => tags.includes(t))) {
        ips.push(client.assignedIp);
      }
    }
    return ips;
  }

  // ── 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,
    };
  }

  // ── Per-client security ────────────────────────────────────────────────

  /**
   * Build per-client security settings for the smartvpn daemon.
   * Maps dcrouter-level fields (forceDestinationSmartproxy, allow/block lists)
   * to smartvpn's IClientSecurity with a destinationPolicy.
   */
  private buildClientSecurity(client: VpnClientDoc): plugins.smartvpn.IClientSecurity {
    const security: plugins.smartvpn.IClientSecurity = {};
    const forceSmartproxy = client.forceDestinationSmartproxy ?? true;

    if (!forceSmartproxy) {
      // Client traffic goes directly — not forced to SmartProxy
      security.destinationPolicy = {
        default: 'allow' as const,
        blockList: client.destinationBlockList,
      };
    } else if (client.destinationAllowList?.length || client.destinationBlockList?.length) {
      // Client is forced to SmartProxy, but with per-client allow/block overrides
      security.destinationPolicy = {
        default: 'forceTarget' as const,
        target: '127.0.0.1',
        allowList: client.destinationAllowList,
        blockList: client.destinationBlockList,
      };
    }
    // else: no per-client policy, server-wide applies

    return security;
  }

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

  private async loadOrGenerateServerKeys(): Promise<VpnServerKeysDoc> {
    const stored = await VpnServerKeysDoc.load();
    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 doc = stored || new VpnServerKeysDoc();
    doc.noisePrivateKey = noiseKeys.privateKey;
    doc.noisePublicKey = noiseKeys.publicKey;
    doc.wgPrivateKey = wgKeys.privateKey;
    doc.wgPublicKey = wgKeys.publicKey;
    await doc.save();

    logger.log('info', 'Generated and persisted new VPN server keys');
    return doc;
  }

  private async loadPersistedClients(): Promise<void> {
    const docs = await VpnClientDoc.findAll();
    for (const doc of docs) {
      // Migrate legacy `tags` → `serverDefinedClientTags`
      if (!doc.serverDefinedClientTags && (doc as any).tags) {
        doc.serverDefinedClientTags = (doc as any).tags;
        (doc as any).tags = undefined;
        await doc.save();
      }
      this.clients.set(doc.clientId, doc);
    }
    if (this.clients.size > 0) {
      logger.log('info', `Loaded ${this.clients.size} persisted VPN client(s)`);
    }
  }

  private async persistClient(client: VpnClientDoc): Promise<void> {
    await client.save();
  }
}