feat: add websocket server endpoint for theta-agent

nodejs/app.js

@@ -43,6 +43,9 @@ app.onListen.push(function(){
       // socket.broadcast.emit('P2PSub', msg);
     });
   });
+  
+  // Initialize Theta Agent WebSockets
+  require('./routes/api_agent')(app);
 }); 
 
 // Gzip text responses (HTML/JS/CSS/JSON). The admin UI loads ~13 separate,

nodejs/bin/www

@@ -25,6 +25,19 @@ var server = http.createServer(app);
 var io = require('socket.io')(server);
 app.io = io;
 
+const WebSocket = require('ws');
+const wss = new WebSocket.Server({ noServer: true });
+server.on('upgrade', (request, socket, head) => {
+  // We only handle upgrade for /api/agent/ws.
+  // Socket.IO handles its own upgrades natively because it attaches directly to `server`.
+  if (request.url.startsWith('/api/agent/ws')) {
+    wss.handleUpgrade(request, socket, head, (ws) => {
+      wss.emit('connection', ws, request);
+    });
+  }
+});
+app.wss = wss;
+
 const models = require('../models');
 
 /**

nodejs/package-lock.json

@@ -42,6 +42,7 @@
         "nodemailer": "^9.0.0",
         "p2psub": "^0.2.0",
         "socket.io": "^4.8.3",
+        "ws": "^8.21.1",
         "xss": "^1.0.15"
       },
       "devDependencies": {

nodejs/package.json

@@ -54,6 +54,7 @@
     "nodemailer": "^9.0.0",
     "p2psub": "^0.2.0",
     "socket.io": "^4.8.3",
+    "ws": "^8.21.1",
     "xss": "^1.0.15"
   },
   "license": "MIT",

nodejs/routes/api_agent.js

@@ -0,0 +1,53 @@
+'use strict';
+
+module.exports = function initAgentWebSockets(app) {
+    if (!app.wss) {
+        console.warn("WebSocket server for agents is not initialized.");
+        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'];
+
+        if (!token) {
+            ws.close(4001, 'Unauthorized: Missing token');
+            return;
+        }
+
+        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'
+            }
+        }));
+    });
+};