whale-code 6.5.5 → 6.5.6

package/dist/local-agent/connection.js

@@ -1,5 +1,5 @@
 /**
- * Local Agent Connection — maintains WebSocket connection to SwagManager cloud.
+ * Local Agent Connection — maintains WebSocket connection to WhaleTools cloud.
  *
  * Features:
  * - Outbound WSS connection (no firewall issues)
@@ -7,326 +7,385 @@
  * - Heartbeat/pong handling
  * - Command routing to executor
  */
+
 import WebSocket from "ws";
 import os from "node:os";
 import { discoverTools } from "./discovery.js";
 import { exec, getSystemInfo } from "./executor.js";
+
 // ============================================================================
 // CONFIG
 // ============================================================================
+
 const DEFAULT_SERVER_URL = "wss://whale-agent.fly.dev/agent/ws";
 const MAX_RECONNECT_DELAY_MS = 60_000;
 const INITIAL_RECONNECT_DELAY_MS = 1_000;
 // ============================================================================
 // AGENT CONNECTION
 // ============================================================================
+
 export class LocalAgentConnection {
-    ws = null;
-    config;
-    reconnectAttempts = 0;
-    reconnectTimer = null;
-    intentionalClose = false;
-    discoveredTools = [];
-    _agentId = null;
-    constructor(config) {
-        this.config = {
-            serverUrl: DEFAULT_SERVER_URL,
-            autoReconnect: true,
-            onConnected: () => { },
-            onDisconnected: () => { },
-            onError: () => { },
-            onLog: () => { },
-            ...config,
-        };
-    }
-    /**
-     * Connect to the SwagManager cloud gateway.
-     */
-    async connect() {
-        this.intentionalClose = false;
-        // Discover tools before connecting
-        this.log("Scanning for installed security tools...");
-        this.discoveredTools = discoverTools();
-        this.log(`Found ${this.discoveredTools.length} tools: ${this.discoveredTools.map(t => t.name).join(", ") || "none"}`);
-        this.log(`Connecting to ${this.config.serverUrl}...`);
-        return new Promise((resolve, reject) => {
-            try {
-                this.ws = new WebSocket(this.config.serverUrl, {
-                    headers: {
-                        "User-Agent": "swag-agent/1.0",
-                    },
-                });
-            }
-            catch (err) {
-                reject(err);
-                return;
-            }
-            const connectTimeout = setTimeout(() => {
-                if (this.ws?.readyState !== WebSocket.OPEN) {
-                    this.ws?.terminate();
-                    reject(new Error("Connection timeout"));
-                }
-            }, 15_000);
-            this.ws.on("open", () => {
-                clearTimeout(connectTimeout);
-                this.reconnectAttempts = 0;
-                // Send auth message
-                this.send({
-                    type: "auth",
-                    api_key: this.config.apiKey,
-                    capabilities: this.discoveredTools.map(t => t.name),
-                    platform: process.platform,
-                    hostname: os.hostname(),
-                    version: "1.0.0",
-                });
-            });
-            this.ws.on("message", (raw) => {
-                let msg;
-                try {
-                    msg = JSON.parse(raw.toString());
-                }
-                catch {
-                    this.log("Received invalid JSON");
-                    return;
-                }
-                this.handleMessage(msg, resolve);
-            });
-            this.ws.on("close", (code, reason) => {
-                clearTimeout(connectTimeout);
-                const reasonStr = reason?.toString() || `code ${code}`;
-                this.log(`Disconnected: ${reasonStr}`);
-                this.config.onDisconnected(reasonStr);
-                if (!this.intentionalClose && this.config.autoReconnect) {
-                    this.scheduleReconnect();
-                }
-            });
-            this.ws.on("error", (err) => {
-                clearTimeout(connectTimeout);
-                this.config.onError(err);
-                if (this.ws?.readyState !== WebSocket.OPEN) {
-                    reject(err);
-                }
-            });
+  ws = null;
+  reconnectAttempts = 0;
+  reconnectTimer = null;
+  intentionalClose = false;
+  discoveredTools = [];
+  _agentId = null;
+  constructor(config) {
+    this.config = {
+      serverUrl: DEFAULT_SERVER_URL,
+      autoReconnect: true,
+      onConnected: () => {},
+      onDisconnected: () => {},
+      onError: () => {},
+      onLog: () => {},
+      ...config
+    };
+  }
+
+  /**
+   * Connect to the WhaleTools cloud gateway.
+   */
+  async connect() {
+    this.intentionalClose = false;
+
+    // Discover tools before connecting
+    this.log("Scanning for installed security tools...");
+    this.discoveredTools = discoverTools();
+    this.log(`Found ${this.discoveredTools.length} tools: ${this.discoveredTools.map(t => t.name).join(", ") || "none"}`);
+    this.log(`Connecting to ${this.config.serverUrl}...`);
+    return new Promise((resolve, reject) => {
+      try {
+        this.ws = new WebSocket(this.config.serverUrl, {
+          headers: {
+            "User-Agent": "whale-agent/1.0"
+          }
         });
-    }
-    /**
-     * Disconnect gracefully.
-     */
-    disconnect() {
-        this.intentionalClose = true;
-        if (this.reconnectTimer) {
-            clearTimeout(this.reconnectTimer);
-            this.reconnectTimer = null;
+      } catch (err) {
+        reject(err);
+        return;
+      }
+      const connectTimeout = setTimeout(() => {
+        if (this.ws?.readyState !== WebSocket.OPEN) {
+          this.ws?.terminate();
+          reject(new Error("Connection timeout"));
         }
-        if (this.ws) {
-            this.ws.close(1000, "Agent shutting down");
-            this.ws = null;
-        }
-    }
-    /**
-     * Check if connected.
-     */
-    get connected() {
-        return this.ws?.readyState === WebSocket.OPEN;
-    }
-    // ============================================================================
-    // MESSAGE HANDLING
-    // ============================================================================
-    handleMessage(msg, connectResolve) {
-        switch (msg.type) {
-            case "authenticated":
-                this._agentId = msg.agent_id;
-                this.log(msg.message || "Authenticated");
-                this.config.onConnected();
-                connectResolve?.();
-                break;
-            case "error":
-                this.log(`Server error: ${msg.error}`);
-                if (msg.error?.includes("Authentication")) {
-                    this.intentionalClose = true; // Don't reconnect on auth failure
-                }
-                break;
-            case "ping":
-                this.send({ type: "pong" });
-                break;
-            case "exec":
-                this.handleExec(msg);
-                break;
-            case "tool_exec":
-                this.handleToolExec(msg);
-                break;
-            case "discover":
-                this.handleDiscover(msg);
-                break;
-            case "shutdown":
-                this.log("Server requested shutdown — will reconnect");
-                // Do NOT set intentionalClose — server restarts (deploys) should auto-reconnect
-                this.ws?.close(1000, "Server shutdown");
-                break;
-            default:
-                this.log(`Unknown message type: ${msg.type}`);
-        }
-    }
-    async handleExec(msg) {
-        const { request_id, command, session_id, timeout } = msg;
-        if (!command) {
-            this.sendResult(request_id, { success: false, error: "No command provided" });
-            return;
-        }
-        this.log(`Executing: ${command.substring(0, 80)}${command.length > 80 ? "..." : ""}`);
+      }, 15_000);
+      this.ws.on("open", () => {
+        clearTimeout(connectTimeout);
+        this.reconnectAttempts = 0;
+
+        // Send auth message
+        this.send({
+          type: "auth",
+          api_key: this.config.apiKey,
+          capabilities: this.discoveredTools.map(t => t.name),
+          platform: process.platform,
+          hostname: os.hostname(),
+          version: "1.0.0"
+        });
+      });
+      this.ws.on("message", raw => {
+        let msg;
         try {
-            const result = await exec(command, { session_id, timeout });
-            this.sendResult(request_id, result);
-        }
-        catch (err) {
-            this.sendResult(request_id, {
-                success: false,
-                error: err.message || "Execution failed",
-            });
+          msg = JSON.parse(raw.toString());
+        } catch {
+          this.log("Received invalid JSON");
+          return;
         }
-    }
-    async handleToolExec(msg) {
-        const { request_id, tool, args, timeout } = msg;
-        this.log(`Tool exec: ${tool}`);
-        // Route to tool-specific handlers or fallback to generic exec
-        try {
-            let result;
-            switch (tool) {
-                case "tshark": {
-                    const iface = args.interface || (process.platform === "darwin" ? "en0" : "eth0");
-                    const filter = args.filter ? `-f "${args.filter}"` : "";
-                    const count = args.count || 100;
-                    const cmd = `tshark -i ${iface} ${filter} -c ${count} 2>&1`;
-                    result = await exec(cmd, { timeout: timeout || 60000 });
-                    break;
-                }
-                case "burpsuite": {
-                    const port = args.port || 1337;
-                    const action = args.action || "status";
-                    if (action === "status") {
-                        result = await exec(`curl -s http://localhost:${port}/v0.1/ 2>&1 || echo "Burp Suite not running on port ${port}"`, { timeout: 5000 });
-                    }
-                    else if (action === "scan") {
-                        result = await exec(`curl -s -X POST http://localhost:${port}/v0.1/scan -H "Content-Type: application/json" -d '${JSON.stringify({ urls: [args.target_url] })}' 2>&1`, { timeout: 10000 });
-                    }
-                    else {
-                        result = { success: false, error: `Unknown burpsuite action: ${action}` };
-                    }
-                    break;
-                }
-                case "zap": {
-                    const port = args.port || 8090;
-                    const action = args.action || "status";
-                    const apiKey = args.api_key || "";
-                    const base = `http://localhost:${port}/JSON`;
-                    if (action === "status") {
-                        result = await exec(`curl -s "${base}/core/view/version/?apikey=${apiKey}" 2>&1`, { timeout: 5000 });
-                    }
-                    else if (action === "spider") {
-                        result = await exec(`curl -s "${base}/spider/action/scan/?apikey=${apiKey}&url=${encodeURIComponent(args.target_url || "")}" 2>&1`, { timeout: 10000 });
-                    }
-                    else if (action === "active_scan") {
-                        result = await exec(`curl -s "${base}/ascan/action/scan/?apikey=${apiKey}&url=${encodeURIComponent(args.target_url || "")}" 2>&1`, { timeout: 10000 });
-                    }
-                    else {
-                        result = { success: false, error: `Unknown zap action: ${action}` };
-                    }
-                    break;
-                }
-                case "aircrack-ng": {
-                    const action = args.action || "help";
-                    if (action === "scan") {
-                        const iface = args.interface || "wlan0";
-                        result = await exec(`airmon-ng start ${iface} && airodump-ng ${iface}mon --write /tmp/airodump --output-format csv -w 10 2>&1`, { timeout: timeout || 30000 });
-                    }
-                    else if (action === "crack") {
-                        const capFile = args.cap_file;
-                        const wordlist = args.wordlist || "/usr/share/wordlists/rockyou.txt";
-                        if (!capFile) {
-                            result = { success: false, error: "cap_file is required" };
-                        }
-                        else {
-                            result = await exec(`aircrack-ng -w ${wordlist} ${capFile} 2>&1`, { timeout: timeout || 300000 });
-                        }
-                    }
-                    else {
-                        result = await exec(`aircrack-ng --help 2>&1`, { timeout: 5000 });
-                    }
-                    break;
-                }
-                case "netdiscover": {
-                    const iface = args.interface || (process.platform === "darwin" ? "en0" : "eth0");
-                    const range = args.range || "";
-                    const rangeFlag = range ? `-r ${range}` : "";
-                    result = await exec(`netdiscover -i ${iface} ${rangeFlag} -P -N 2>&1 | head -50`, { timeout: timeout || 30000 });
-                    break;
-                }
-                default:
-                    // Generic: try to run as command
-                    result = await exec(`${tool} ${args.args || ""} 2>&1`, { timeout: timeout || 30000 });
-            }
-            this.sendResult(request_id, result);
+        this.handleMessage(msg, resolve);
+      });
+      this.ws.on("close", (code, reason) => {
+        clearTimeout(connectTimeout);
+        const reasonStr = reason?.toString() || `code ${code}`;
+        this.log(`Disconnected: ${reasonStr}`);
+        this.config.onDisconnected(reasonStr);
+        if (!this.intentionalClose && this.config.autoReconnect) {
+          this.scheduleReconnect();
         }
-        catch (err) {
-            this.sendResult(request_id, {
-                success: false,
-                error: err.message || `Tool ${tool} failed`,
-            });
+      });
+      this.ws.on("error", err => {
+        clearTimeout(connectTimeout);
+        this.config.onError(err);
+        if (this.ws?.readyState !== WebSocket.OPEN) {
+          reject(err);
         }
+      });
+    });
+  }
+
+  /**
+   * Disconnect gracefully.
+   */
+  disconnect() {
+    this.intentionalClose = true;
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
     }
-    handleDiscover(msg) {
-        const tools = discoverTools();
-        this.discoveredTools = tools;
-        // Update capabilities on the server
-        this.send({
-            type: "tools",
-            tools: tools.map(t => t.name),
-        });
-        // Also send as result
-        this.sendResult(msg.request_id, {
-            success: true,
-            data: {
-                tools: tools.map(t => ({
-                    name: t.name,
-                    path: t.path,
-                    version: t.version,
-                    category: t.category,
-                    running: t.running,
-                })),
-                system: getSystemInfo(),
-            },
-        });
+    if (this.ws) {
+      this.ws.close(1000, "Agent shutting down");
+      this.ws = null;
     }
-    // ============================================================================
-    // HELPERS
-    // ============================================================================
-    send(data) {
-        if (this.ws?.readyState === WebSocket.OPEN) {
-            this.ws.send(JSON.stringify(data));
+  }
+
+  /**
+   * Check if connected.
+   */
+  get connected() {
+    return this.ws?.readyState === WebSocket.OPEN;
+  }
+
+  // ============================================================================
+  // MESSAGE HANDLING
+  // ============================================================================
+
+  handleMessage(msg, connectResolve) {
+    switch (msg.type) {
+      case "authenticated":
+        this._agentId = msg.agent_id;
+        this.log(msg.message || "Authenticated");
+        this.config.onConnected();
+        connectResolve?.();
+        break;
+      case "error":
+        this.log(`Server error: ${msg.error}`);
+        if (msg.error?.includes("Authentication")) {
+          this.intentionalClose = true; // Don't reconnect on auth failure
         }
-    }
-    sendResult(requestId, result) {
+        break;
+      case "ping":
         this.send({
-            type: "result",
-            request_id: requestId,
-            ...result,
+          type: "pong"
         });
+        break;
+      case "exec":
+        this.handleExec(msg);
+        break;
+      case "tool_exec":
+        this.handleToolExec(msg);
+        break;
+      case "discover":
+        this.handleDiscover(msg);
+        break;
+      case "shutdown":
+        this.log("Server requested shutdown — will reconnect");
+        // Do NOT set intentionalClose — server restarts (deploys) should auto-reconnect
+        this.ws?.close(1000, "Server shutdown");
+        break;
+      default:
+        this.log(`Unknown message type: ${msg.type}`);
+    }
+  }
+  async handleExec(msg) {
+    const {
+      request_id,
+      command,
+      session_id,
+      timeout
+    } = msg;
+    if (!command) {
+      this.sendResult(request_id, {
+        success: false,
+        error: "No command provided"
+      });
+      return;
+    }
+    this.log(`Executing: ${command.substring(0, 80)}${command.length > 80 ? "..." : ""}`);
+    try {
+      const result = await exec(command, {
+        session_id,
+        timeout
+      });
+      this.sendResult(request_id, result);
+    } catch (err) {
+      this.sendResult(request_id, {
+        success: false,
+        error: err.message || "Execution failed"
+      });
     }
-    scheduleReconnect() {
-        this.reconnectAttempts++;
-        const delay = Math.min(INITIAL_RECONNECT_DELAY_MS * Math.pow(2, this.reconnectAttempts - 1), MAX_RECONNECT_DELAY_MS);
-        this.log(`Reconnecting in ${Math.round(delay / 1000)}s (attempt ${this.reconnectAttempts})...`);
-        this.reconnectTimer = setTimeout(async () => {
-            try {
-                await this.connect();
+  }
+  async handleToolExec(msg) {
+    const {
+      request_id,
+      tool,
+      args,
+      timeout
+    } = msg;
+    this.log(`Tool exec: ${tool}`);
+
+    // Route to tool-specific handlers or fallback to generic exec
+    try {
+      let result;
+      switch (tool) {
+        case "tshark":
+          {
+            const iface = args.interface || (process.platform === "darwin" ? "en0" : "eth0");
+            const filter = args.filter ? `-f "${args.filter}"` : "";
+            const count = args.count || 100;
+            const cmd = `tshark -i ${iface} ${filter} -c ${count} 2>&1`;
+            result = await exec(cmd, {
+              timeout: timeout || 60000
+            });
+            break;
+          }
+        case "burpsuite":
+          {
+            const port = args.port || 1337;
+            const action = args.action || "status";
+            if (action === "status") {
+              result = await exec(`curl -s http://localhost:${port}/v0.1/ 2>&1 || echo "Burp Suite not running on port ${port}"`, {
+                timeout: 5000
+              });
+            } else if (action === "scan") {
+              result = await exec(`curl -s -X POST http://localhost:${port}/v0.1/scan -H "Content-Type: application/json" -d '${JSON.stringify({
+                urls: [args.target_url]
+              })}' 2>&1`, {
+                timeout: 10000
+              });
+            } else {
+              result = {
+                success: false,
+                error: `Unknown burpsuite action: ${action}`
+              };
             }
-            catch (err) {
-                this.log(`Reconnect failed: ${err.message}`);
-                // connect() will trigger close event which will schedule another reconnect
+            break;
+          }
+        case "zap":
+          {
+            const port = args.port || 8090;
+            const action = args.action || "status";
+            const apiKey = args.api_key || "";
+            const base = `http://localhost:${port}/JSON`;
+            if (action === "status") {
+              result = await exec(`curl -s "${base}/core/view/version/?apikey=${apiKey}" 2>&1`, {
+                timeout: 5000
+              });
+            } else if (action === "spider") {
+              result = await exec(`curl -s "${base}/spider/action/scan/?apikey=${apiKey}&url=${encodeURIComponent(args.target_url || "")}" 2>&1`, {
+                timeout: 10000
+              });
+            } else if (action === "active_scan") {
+              result = await exec(`curl -s "${base}/ascan/action/scan/?apikey=${apiKey}&url=${encodeURIComponent(args.target_url || "")}" 2>&1`, {
+                timeout: 10000
+              });
+            } else {
+              result = {
+                success: false,
+                error: `Unknown zap action: ${action}`
+              };
+            }
+            break;
+          }
+        case "aircrack-ng":
+          {
+            const action = args.action || "help";
+            if (action === "scan") {
+              const iface = args.interface || "wlan0";
+              result = await exec(`airmon-ng start ${iface} && airodump-ng ${iface}mon --write /tmp/airodump --output-format csv -w 10 2>&1`, {
+                timeout: timeout || 30000
+              });
+            } else if (action === "crack") {
+              const capFile = args.cap_file;
+              const wordlist = args.wordlist || "/usr/share/wordlists/rockyou.txt";
+              if (!capFile) {
+                result = {
+                  success: false,
+                  error: "cap_file is required"
+                };
+              } else {
+                result = await exec(`aircrack-ng -w ${wordlist} ${capFile} 2>&1`, {
+                  timeout: timeout || 300000
+                });
+              }
+            } else {
+              result = await exec(`aircrack-ng --help 2>&1`, {
+                timeout: 5000
+              });
             }
-        }, delay);
+            break;
+          }
+        case "netdiscover":
+          {
+            const iface = args.interface || (process.platform === "darwin" ? "en0" : "eth0");
+            const range = args.range || "";
+            const rangeFlag = range ? `-r ${range}` : "";
+            result = await exec(`netdiscover -i ${iface} ${rangeFlag} -P -N 2>&1 | head -50`, {
+              timeout: timeout || 30000
+            });
+            break;
+          }
+        default:
+          // Generic: try to run as command
+          result = await exec(`${tool} ${args.args || ""} 2>&1`, {
+            timeout: timeout || 30000
+          });
+      }
+      this.sendResult(request_id, result);
+    } catch (err) {
+      this.sendResult(request_id, {
+        success: false,
+        error: err.message || `Tool ${tool} failed`
+      });
     }
-    log(msg) {
-        this.config.onLog(`[swag-agent] ${msg}`);
+  }
+  handleDiscover(msg) {
+    const tools = discoverTools();
+    this.discoveredTools = tools;
+
+    // Update capabilities on the server
+    this.send({
+      type: "tools",
+      tools: tools.map(t => t.name)
+    });
+
+    // Also send as result
+    this.sendResult(msg.request_id, {
+      success: true,
+      data: {
+        tools: tools.map(t => ({
+          name: t.name,
+          path: t.path,
+          version: t.version,
+          category: t.category,
+          running: t.running
+        })),
+        system: getSystemInfo()
+      }
+    });
+  }
+
+  // ============================================================================
+  // HELPERS
+  // ============================================================================
+
+  send(data) {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      this.ws.send(JSON.stringify(data));
     }
+  }
+  sendResult(requestId, result) {
+    this.send({
+      type: "result",
+      request_id: requestId,
+      ...result
+    });
+  }
+  scheduleReconnect() {
+    this.reconnectAttempts++;
+    const delay = Math.min(INITIAL_RECONNECT_DELAY_MS * Math.pow(2, this.reconnectAttempts - 1), MAX_RECONNECT_DELAY_MS);
+    this.log(`Reconnecting in ${Math.round(delay / 1000)}s (attempt ${this.reconnectAttempts})...`);
+    this.reconnectTimer = setTimeout(async () => {
+      try {
+        await this.connect();
+      } catch (err) {
+        this.log(`Reconnect failed: ${err.message}`);
+        // connect() will trigger close event which will schedule another reconnect
+      }
+    }, delay);
+  }
+  log(msg) {
+    this.config.onLog(`[whale-agent] ${msg}`);
+  }
 }
+//# sourceMappingURL=connection.js.map