feat(mitm): parameterize addDNSEntry/removeDNSEntry for dynamic hosts (fix1)

2026-07-26 09:52:11 +03:00 · 2026-05-28 11:23:45 -03:00
parent 53754ae93c
commit 785f8e4117
2 changed files with 182 additions and 85 deletions
--- a/src/mitm/dns/dnsConfig.ts
+++ b/src/mitm/dns/dnsConfig.ts
@@ -7,16 +7,96 @@ import {
  runElevatedPowerShell,
 } from "../systemCommands.ts";

-const TARGET_HOST = "daily-cloudcode-pa.googleapis.com";
+// Legacy Antigravity defaults preserved for backward compat.
+const ANTIGRAVITY_HOSTS = [
+  "daily-cloudcode-pa.googleapis.com",
+  "cloudcode-pa.googleapis.com",
+  "daily-cloudcode-pa.sandbox.googleapis.com",
+  "autopush-cloudcode-pa.sandbox.googleapis.com",
+];
+
 const IS_WIN = process.platform === "win32";
 const HOSTS_FILE = IS_WIN
  ? path.join(process.env.SystemRoot || "C:\\Windows", "System32", "drivers", "etc", "hosts")
  : "/etc/hosts";

-// Both IPv4 and IPv6 entries are needed — modern Windows apps often resolve
-// to IPv6 first, bypassing an IPv4-only MITM redirect.
-const DNS_ENTRIES = [`127.0.0.1 ${TARGET_HOST}`, `::1 ${TARGET_HOST}`];
+/**
+ * Build the set of /etc/hosts lines for a given hostname.
+ * Both IPv4 and IPv6 are needed — modern systems often resolve IPv6 first.
+ */
+function dnsLines(hostname: string): string[] {
+  return [`127.0.0.1 ${hostname}`, `::1 ${hostname}`];
+}

+/**
+ * Read the current hosts file content. Returns empty string on error.
+ */
+function readHostsFile(): string {
+  try {
+    return fs.readFileSync(HOSTS_FILE, "utf8");
+  } catch {
+    return "";
+  }
+}
+
+/**
+ * Check whether all IPv4+IPv6 lines for `hostname` are present in the hosts file.
+ */
+function hasHostEntry(hostsContent: string, hostname: string): boolean {
+  const lines = hostsContent.split(/\r?\n/);
+  return dnsLines(hostname).every((entry) => {
+    const [ip, host] = entry.split(/\s+/);
+    return lines.some((line) => {
+      const parts = line.trim().split(/\s+/).filter(Boolean);
+      return parts.length >= 2 && parts[0] === ip && parts.includes(host);
+    });
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Public API — parametrized (new)
+// ---------------------------------------------------------------------------
+
+/**
+ * Add /etc/hosts entries for every hostname in `hosts`.
+ * Idempotent — existing entries are not duplicated.
+ * Complies with Hard Rule #13: no string interpolation in shell commands.
+ */
+export async function addDNSEntries(hosts: string[], sudoPassword: string): Promise<void> {
+  const hostsContent = readHostsFile();
+
+  for (const hostname of hosts) {
+    const lines = dnsLines(hostname);
+    const missing = lines.filter((entry) => {
+      const [ip, host] = entry.split(/\s+/);
+      const existing = hostsContent.split(/\r?\n/);
+      return !existing.some((line) => {
+        const parts = line.trim().split(/\s+/).filter(Boolean);
+        return parts.length >= 2 && parts[0] === ip && parts.includes(host);
+      });
+    });
+
+    for (const entry of missing) {
+      if (IS_WIN) {
+        await runElevatedPowerShell(
+          `Add-Content -LiteralPath ${quotePowerShell(HOSTS_FILE)} -Value ${quotePowerShell(entry)}`
+        );
+      } else {
+        // Hard Rule #13: entry is passed as stdin data, not interpolated into the command.
+        await execFileWithPassword(
+          "sudo",
+          ["-S", "tee", "-a", HOSTS_FILE],
+          sudoPassword,
+          `${entry}\n`
+        );
+      }
+      console.log(`[DNS] Added entry: ${entry}`);
+    }
+  }
+}
+
+// Node.js inline script for removing hosts entries — uses process.argv so no
+// values are interpolated into the script body (Hard Rule #13).
 const REMOVE_HOSTS_ENTRY_SCRIPT = `
 const fs = require("fs");
 const filePath = process.argv[1];
@@ -29,91 +109,74 @@ const filtered = content.split(/\\r?\\n/).filter((line) => {
 fs.writeFileSync(filePath, filtered.join("\\n").replace(/\\n*$/, "\\n"));
 `;

-export function checkDNSEntry(): boolean {
-  try {
-    const hostsContent = fs.readFileSync(HOSTS_FILE, "utf8");
-    const lines = hostsContent.split(/\r?\n/);
-    return DNS_ENTRIES.every((entry) => {
-      const entryIp = entry.split(/\s+/)[0];
-      return lines.some((line) => {
-        const parts = line.trim().split(/\s+/);
-        return parts.length >= 2 && parts[0] === entryIp && parts.some((p) => p === TARGET_HOST);
-      });
-    });
-  } catch {
-    return false;
+/**
+ * Remove /etc/hosts entries for every hostname in `hosts`.
+ * Idempotent — silently skips hosts that are not present.
+ * Complies with Hard Rule #13: HOSTS_FILE and hostname are passed as argv, not interpolated.
+ */
+export async function removeDNSEntries(hosts: string[], sudoPassword: string): Promise<void> {
+  const hostsContent = readHostsFile();
+
+  for (const hostname of hosts) {
+    if (!hasHostEntry(hostsContent, hostname)) {
+      console.log(`[DNS] Entry for ${hostname} not present — skipping`);
+      continue;
+    }
+
+    try {
+      if (IS_WIN) {
+        const psHostsFile = quotePowerShell(HOSTS_FILE);
+        const psTargetHost = quotePowerShell(hostname);
+        await runElevatedPowerShell(`
+          $hostsFile = ${psHostsFile};
+          $targetHost = ${psTargetHost};
+          $lines = Get-Content -LiteralPath $hostsFile;
+          $filtered = $lines | Where-Object {
+            $parts = ($_ -split '\\s+') | Where-Object { $_ };
+            -not (($parts.Length -ge 2) -and ($parts -contains $targetHost))
+          };
+          Set-Content -LiteralPath $hostsFile -Value $filtered;
+        `);
+      } else {
+        // Hard Rule #13: HOSTS_FILE and hostname are argv arguments, not interpolated.
+        await execFileWithPassword(
+          "sudo",
+          ["-S", process.execPath, "-e", REMOVE_HOSTS_ENTRY_SCRIPT, HOSTS_FILE, hostname],
+          sudoPassword
+        );
+      }
+      console.log(`[DNS] Removed entries for ${hostname}`);
+    } catch (error) {
+      throw new Error(`Failed to remove DNS entry for ${hostname}: ${getErrorMessage(error)}`);
+    }
  }
 }

-export async function addDNSEntry(sudoPassword: string): Promise<void> {
-  if (checkDNSEntry()) {
-    console.log(`DNS entries for ${TARGET_HOST} already exist (IPv4 + IPv6)`);
-    return;
-  }
+// ---------------------------------------------------------------------------
+// Legacy API — backward compat wrappers for manager.ts callers
+// ---------------------------------------------------------------------------

-  const entriesToAdd = DNS_ENTRIES.filter((entry) => {
-    const entryIp = entry.split(/\s+/)[0];
-    try {
-      const hostsContent = fs.readFileSync(HOSTS_FILE, "utf8");
-      const lines = hostsContent.split(/\r?\n/);
-      return !lines.some((line) => {
-        const parts = line.trim().split(/\s+/);
-        return parts.length >= 2 && parts[0] === entryIp && parts.some((p) => p === TARGET_HOST);
-      });
-    } catch {
-      return true;
-    }
-  });
-
-  if (entriesToAdd.length === 0) return;
-
-  for (const entry of entriesToAdd) {
-    if (IS_WIN) {
-      await runElevatedPowerShell(
-        `Add-Content -LiteralPath ${quotePowerShell(HOSTS_FILE)} -Value ${quotePowerShell(entry)}`
-      );
-    } else {
-      await execFileWithPassword(
-        "sudo",
-        ["-S", "tee", "-a", HOSTS_FILE],
-        sudoPassword,
-        `${entry}\n`
-      );
-    }
-    console.log(`Added DNS entry: ${entry}`);
-  }
+/**
+ * Check whether the Antigravity default DNS entries are present.
+ * Preserved for backward compat (called by getMitmStatus and other callers).
+ */
+export function checkDNSEntry(): boolean {
+  const hostsContent = readHostsFile();
+  return ANTIGRAVITY_HOSTS.every((h) => hasHostEntry(hostsContent, h));
 }

 /**
- * Remove DNS entry from hosts file
+ * Add DNS entries for the Antigravity default hosts.
+ * Delegates to `addDNSEntries` — backward compat wrapper.
+ */
+export async function addDNSEntry(sudoPassword: string): Promise<void> {
+  await addDNSEntries(ANTIGRAVITY_HOSTS, sudoPassword);
+}
+
+/**
+ * Remove DNS entries for the Antigravity default hosts.
+ * Delegates to `removeDNSEntries` — backward compat wrapper.
 */
 export async function removeDNSEntry(sudoPassword: string): Promise<void> {
-  if (!checkDNSEntry()) {
-    console.log(`DNS entry for ${TARGET_HOST} does not exist`);
-    return;
-  }
-
-  try {
-    if (IS_WIN) {
-      await runElevatedPowerShell(`
-        $hostsFile = ${quotePowerShell(HOSTS_FILE)};
-        $targetHost = ${quotePowerShell(TARGET_HOST)};
-        $lines = Get-Content -LiteralPath $hostsFile;
-        $filtered = $lines | Where-Object {
-          $parts = ($_ -split '\\s+') | Where-Object { $_ };
-          -not (($parts.Length -ge 2) -and ($parts -contains $targetHost))
-        };
-        Set-Content -LiteralPath $hostsFile -Value $filtered;
-      `);
-    } else {
-      await execFileWithPassword(
-        "sudo",
-        ["-S", process.execPath, "-e", REMOVE_HOSTS_ENTRY_SCRIPT, HOSTS_FILE, TARGET_HOST],
-        sudoPassword
-      );
-    }
-    console.log(`✅ Removed DNS entry for ${TARGET_HOST}`);
-  } catch (error) {
-    throw new Error(`Failed to remove DNS entry: ${getErrorMessage(error)}`);
-  }
+  await removeDNSEntries(ANTIGRAVITY_HOSTS, sudoPassword);
 }
--- a/src/mitm/manager.ts
+++ b/src/mitm/manager.ts
@@ -2,12 +2,14 @@ import { spawn, type ChildProcess } from "child_process";
 import path from "path";
 import fs from "fs";
 import { resolveMitmDataDir } from "./dataDir.ts";
-import { addDNSEntry, removeDNSEntry } from "./dns/dnsConfig.ts";
+import { addDNSEntry, addDNSEntries, removeDNSEntry } from "./dns/dnsConfig.ts";
 import { generateCert } from "./cert/generate.ts";
 import { installCert } from "./cert/install.ts";
 import { ALL_TARGETS } from "./targets/index.ts";
 import { detectAgent } from "./detection/index.ts";
 import type { AgentId, DetectionResult, MitmTarget } from "./types.ts";
+import { getAllAgentBridgeStates } from "@/lib/db/agentBridgeState.ts";
+import { listCustomHosts } from "@/lib/db/inspectorCustomHosts.ts";

 // Store server process
 let serverProcess: ChildProcess | null = null;
@@ -178,10 +180,42 @@ export async function startMitm(
  // 2. Install certificate to system keychain
  await installCert(sudoPassword, certPath);

-  // 3. Add DNS entry
-  console.log("Adding DNS entry...");
+  // 3. Add DNS entries: Antigravity defaults + all agents with dns_enabled=true +
+  //    all custom hosts with enabled=true.
+  console.log("Adding DNS entries...");
  await addDNSEntry(sudoPassword);

+  // Collect hosts from agents that have dns_enabled=true in the DB.
+  try {
+    const agentStates = getAllAgentBridgeStates();
+    const agentHostsToAdd: string[] = [];
+    for (const state of agentStates) {
+      if (!state.dns_enabled) continue;
+      const target = ALL_TARGETS.find((t) => t.id === state.agent_id);
+      if (target) {
+        agentHostsToAdd.push(...target.hosts);
+      }
+    }
+    if (agentHostsToAdd.length > 0) {
+      console.log(`[MITM] Adding DNS for ${agentHostsToAdd.length} agent host(s)...`);
+      await addDNSEntries(agentHostsToAdd, sudoPassword);
+    }
+  } catch (err) {
+    console.error(`[MITM] Failed to add agent DNS entries (continuing): ${(err as Error).message ?? err}`);
+  }
+
+  // Collect enabled custom hosts.
+  try {
+    const customHosts = listCustomHosts({ enabledOnly: true });
+    const customHostNames = customHosts.map((h) => h.host);
+    if (customHostNames.length > 0) {
+      console.log(`[MITM] Adding DNS for ${customHostNames.length} custom host(s)...`);
+      await addDNSEntries(customHostNames, sudoPassword);
+    }
+  } catch (err) {
+    console.error(`[MITM] Failed to add custom host DNS entries (continuing): ${(err as Error).message ?? err}`);
+  }
+
  // 4. Start MITM server
  console.log("Starting MITM server...");
  const port =