smartproxy/ts/classes.iptablesproxy.ts

import { exec, execSync } from 'child_process';
import { promisify } from 'util';

const execAsync = promisify(exec);

/**
 * Settings for IPTablesProxy.
 */
export interface IIpTableProxySettings {
  fromPort: number;
  toPort: number;
  toHost?: string; // Target host for proxying; defaults to 'localhost'
  preserveSourceIP?: boolean; // If true, the original source IP is preserved.
  deleteOnExit?: boolean;   // If true, clean up marked iptables rules before process exit.
}

/**
 * IPTablesProxy sets up iptables NAT rules to forward TCP traffic.
 * It only supports basic port forwarding and uses iptables comments to tag rules.
 */
export class IPTablesProxy {
  public settings: IIpTableProxySettings;
  private rulesInstalled: boolean = false;
  private ruleTag: string;

  constructor(settings: IIpTableProxySettings) {
    this.settings = {
      ...settings,
      toHost: settings.toHost || 'localhost',
    };
    // Generate a unique identifier for the rules added by this instance.
    this.ruleTag = `IPTablesProxy:${Date.now()}:${Math.random().toString(36).substr(2, 5)}`;

    // If deleteOnExit is true, register cleanup handlers.
    if (this.settings.deleteOnExit) {
      const cleanup = () => {
        try {
          IPTablesProxy.cleanSlateSync();
        } catch (err) {
          console.error('Error cleaning iptables rules on exit:', err);
        }
      };
      process.on('exit', cleanup);
      process.on('SIGINT', () => {
        cleanup();
        process.exit();
      });
      process.on('SIGTERM', () => {
        cleanup();
        process.exit();
      });
    }
  }

  /**
   * Sets up iptables rules for port forwarding.
   * The rules are tagged with a unique comment so that they can be identified later.
   */
  public async start(): Promise<void> {
    const dnatCmd = `iptables -t nat -A PREROUTING -p tcp --dport ${this.settings.fromPort} ` +
      `-j DNAT --to-destination ${this.settings.toHost}:${this.settings.toPort} ` +
      `-m comment --comment "${this.ruleTag}:DNAT"`;
    try {
      await execAsync(dnatCmd);
      console.log(`Added iptables rule: ${dnatCmd}`);
      this.rulesInstalled = true;
    } catch (err) {
      console.error(`Failed to add iptables DNAT rule: ${err}`);
      throw err;
    }

    // If preserveSourceIP is false, add a MASQUERADE rule.
    if (!this.settings.preserveSourceIP) {
      const masqueradeCmd = `iptables -t nat -A POSTROUTING -p tcp -d ${this.settings.toHost} ` +
        `--dport ${this.settings.toPort} -j MASQUERADE ` +
        `-m comment --comment "${this.ruleTag}:MASQ"`;
      try {
        await execAsync(masqueradeCmd);
        console.log(`Added iptables rule: ${masqueradeCmd}`);
      } catch (err) {
        console.error(`Failed to add iptables MASQUERADE rule: ${err}`);
        // Roll back the DNAT rule if MASQUERADE fails.
        try {
          const rollbackCmd = `iptables -t nat -D PREROUTING -p tcp --dport ${this.settings.fromPort} ` +
            `-j DNAT --to-destination ${this.settings.toHost}:${this.settings.toPort} ` +
            `-m comment --comment "${this.ruleTag}:DNAT"`;
          await execAsync(rollbackCmd);
          this.rulesInstalled = false;
        } catch (rollbackErr) {
          console.error(`Rollback failed: ${rollbackErr}`);
        }
        throw err;
      }
    }
  }

  /**
   * Removes the iptables rules that were added in start(), by matching the unique comment.
   */
  public async stop(): Promise<void> {
    if (!this.rulesInstalled) return;

    const dnatDelCmd = `iptables -t nat -D PREROUTING -p tcp --dport ${this.settings.fromPort} ` +
      `-j DNAT --to-destination ${this.settings.toHost}:${this.settings.toPort} ` +
      `-m comment --comment "${this.ruleTag}:DNAT"`;
    try {
      await execAsync(dnatDelCmd);
      console.log(`Removed iptables rule: ${dnatDelCmd}`);
    } catch (err) {
      console.error(`Failed to remove iptables DNAT rule: ${err}`);
    }

    if (!this.settings.preserveSourceIP) {
      const masqueradeDelCmd = `iptables -t nat -D POSTROUTING -p tcp -d ${this.settings.toHost} ` +
        `--dport ${this.settings.toPort} -j MASQUERADE ` +
        `-m comment --comment "${this.ruleTag}:MASQ"`;
      try {
        await execAsync(masqueradeDelCmd);
        console.log(`Removed iptables rule: ${masqueradeDelCmd}`);
      } catch (err) {
        console.error(`Failed to remove iptables MASQUERADE rule: ${err}`);
      }
    }

    this.rulesInstalled = false;
  }

  /**
   * Asynchronously cleans up any iptables rules in the nat table that were added by this module.
   * It looks for rules with comments containing "IPTablesProxy:".
   */
  public static async cleanSlate(): Promise<void> {
    try {
      const { stdout } = await execAsync('iptables-save -t nat');
      const lines = stdout.split('\n');
      const proxyLines = lines.filter(line => line.includes('IPTablesProxy:'));
      for (const line of proxyLines) {
        const trimmedLine = line.trim();
        if (trimmedLine.startsWith('-A')) {
          // Replace the "-A" with "-D" to form a deletion command.
          const deleteRule = trimmedLine.replace('-A', '-D');
          const cmd = `iptables -t nat ${deleteRule}`;
          try {
            await execAsync(cmd);
            console.log(`Cleaned up iptables rule: ${cmd}`);
          } catch (err) {
            console.error(`Failed to remove iptables rule: ${cmd}`, err);
          }
        }
      }
    } catch (err) {
      console.error(`Failed to run iptables-save: ${err}`);
    }
  }

  /**
   * Synchronously cleans up any iptables rules in the nat table that were added by this module.
   * It looks for rules with comments containing "IPTablesProxy:".
   * This method is intended for use in process exit handlers.
   */
  public static cleanSlateSync(): void {
    try {
      const stdout = execSync('iptables-save -t nat').toString();
      const lines = stdout.split('\n');
      const proxyLines = lines.filter(line => line.includes('IPTablesProxy:'));
      for (const line of proxyLines) {
        const trimmedLine = line.trim();
        if (trimmedLine.startsWith('-A')) {
          const deleteRule = trimmedLine.replace('-A', '-D');
          const cmd = `iptables -t nat ${deleteRule}`;
          try {
            execSync(cmd);
            console.log(`Cleaned up iptables rule: ${cmd}`);
          } catch (err) {
            console.error(`Failed to remove iptables rule: ${cmd}`, err);
          }
        }
      }
    } catch (err) {
      console.error(`Failed to run iptables-save: ${err}`);
    }
  }
}