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

const execAsync = promisify(exec);

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.
}

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

  constructor(settings: IIpTableProxySettings) {
    this.settings = {
      ...settings,
      toHost: settings.toHost || 'localhost',
    };
  }

  /**
   * Sets up iptables rules for port forwarding.
   */
  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}`;
    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 (!this.settings.preserveSourceIP) {
      const masqueradeCmd = `iptables -t nat -A POSTROUTING -p tcp -d ${this.settings.toHost} --dport ${this.settings.toPort} -j MASQUERADE`;
      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}`;
          await execAsync(rollbackCmd);
          this.rulesInstalled = false;
        } catch (rollbackErr) {
          console.error(`Rollback failed: ${rollbackErr}`);
        }
        throw err;
      }
    }
  }

  /**
   * Removes the iptables rules that were added in start().
   */
  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}`;
    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`;
      try {
        await execAsync(masqueradeDelCmd);
        console.log(`Removed iptables rule: ${masqueradeDelCmd}`);
      } catch (err) {
        console.error(`Failed to remove iptables MASQUERADE rule: ${err}`);
      }
    }

    this.rulesInstalled = false;
  }
}