feat: Protocol v1.1.0 theta-agent C2 integration, agent install wizard, OpenBao 403 fix, nmap discovery fix, and restored documentation

2026-08-03 02:18:09 -04:00
parent 4945dec9c2
commit b7aac2d2ba
35 changed files with 3003 additions and 49 deletions
@@ -1,53 +1,105 @@
 'use strict';

+const express = require('express');
+const agentManager = require('../utils/agent_manager');
+
 module.exports = function initAgentWebSockets(app) {
-    if (!app.wss) {
-        console.warn("WebSocket server for agents is not initialized.");
-        return;
+  if (!app.wss) {
+    console.warn("WebSocket server for agents is not initialized.");
+    return;
+  }
+
+  app.wss.on('connection', (ws, req) => {
+    const url = new URL(req.url, `http://${req.headers.host || 'localhost'}`);
+    const token = url.searchParams.get('token') || req.headers['authorization'];
+
+    if (!token) {
+      ws.close(4001, 'Unauthorized: Missing token');
+      return;
    }

-    app.wss.on('connection', (ws, req) => {
-        // Parse the token from query param or header (e.g. ?token=XYZ)
-        // For the beta, we will just accept it if a token is present.
-        const url = new URL(req.url, `http://${req.headers.host}`);
-        const token = url.searchParams.get('token') || req.headers['authorization'];
+    const remoteAddr = req.socket.remoteAddress;
+    console.log(`[Theta Agent] Agent connected from ${remoteAddr} with token ${token.substring(0, 8)}...`);

-        if (!token) {
-            ws.close(4001, 'Unauthorized: Missing token');
-            return;
+    agentManager.registerAgent(token, ws, remoteAddr);
+
+    ws.on('message', (message) => {
+      try {
+        const data = JSON.parse(message);
+        if (!data || typeof data.type !== 'string') return;
+
+        const payload = data.payload || {};
+
+        switch (data.type) {
+          case 'discovery':
+            agentManager.handleDiscovery(token, payload);
+            if (app.io) app.io.emit('agent.discovery', { token, payload });
+            break;
+          case 'telemetry':
+            agentManager.handleTelemetry(token, payload);
+            if (app.io) app.io.emit('agent.telemetry', { token, payload });
+            break;
+          case 'heartbeat':
+            agentManager.handleHeartbeat(token, payload, ws);
+            break;
+          case 'response':
+            agentManager.handleResponse(token, payload);
+            if (app.io) app.io.emit('agent.response', { token, payload });
+            break;
+          default:
+            console.log(`[Theta Agent] Received message type '${data.type}' from ${token}`);
        }
-
-        console.log(`[Theta Agent] Agent connected from ${req.socket.remoteAddress}`);
-
-        ws.on('message', (message) => {
-            try {
-                const data = JSON.parse(message);
-                
-                // Example handling incoming telemetry
-                if (data.type === 'telemetry') {
-                    // Send to discovery service or log
-                    // console.log(`[Theta Agent] Received telemetry from ${data.host}`);
-                    
-                    // We can publish it to the event bus for the UI
-                    if(app.contoller && app.contoller.ps) {
-                        app.contoller.ps.publish('agent.telemetry', data);
-                    }
-                }
-            } catch (err) {
-                console.error("[Theta Agent] Error parsing message:", err);
-            }
-        });
-
-        ws.on('close', () => {
-            console.log(`[Theta Agent] Agent disconnected`);
-        });
-
-        // Example: Send a welcome config payload to the agent
-        ws.send(JSON.stringify({
-            type: 'config',
-            payload: {
-                message: 'Welcome to SSO Manager C2'
-            }
-        }));
+      } catch (err) {
+        console.error("[Theta Agent] Error parsing message:", err);
+      }
    });
+
+    ws.on('close', () => {
+      console.log(`[Theta Agent] Agent disconnected (${token})`);
+      agentManager.unregisterAgent(token, ws);
+    });
+
+    // Send initial welcome/config payload
+    try {
+      ws.send(JSON.stringify({
+        type: 'config',
+        payload: {
+          message: 'Connected to SSO Manager C2',
+          protocol_version: '1.1.0'
+        }
+      }));
+    } catch (e) {}
+  });
+
+  // REST API routes for Agent Management (mounted under /api/agent)
+  const router = express.Router();
+
+  router.get('/nodes', (req, res) => {
+    res.json({
+      status: 'ok',
+      agents: agentManager.getConnectedAgents(),
+      publicKey: agentManager.publicKeyPem
+    });
+  });
+
+  router.post('/nodes/:token/command', (req, res) => {
+    const { token } = req.params;
+    const { command, payload, isHighRisk } = req.body;
+
+    if (!command) {
+      return res.status(400).json({ status: 'error', message: 'Command type is required' });
+    }
+
+    try {
+      const HIGH_RISK_COMMANDS = ['reboot', 'service_restart', 'configure_ldap', 'arbitrary_bash', 'update_binary'];
+      const requiresSigning = isHighRisk || HIGH_RISK_COMMANDS.includes(command);
+
+      const msg = agentManager.sendCommand(token, command, payload || {}, requiresSigning);
+      res.json({ status: 'ok', sentMessage: msg });
+    } catch (err) {
+      res.status(400).json({ status: 'error', message: err.message });
+    }
+  });
+
+  app.use('/api/agent', router);
 };