diff --git a/nodejs/app.js b/nodejs/app.js index 4eafe4b..10a094e 100755 --- a/nodejs/app.js +++ b/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, diff --git a/nodejs/bin/www b/nodejs/bin/www index 6550c5a..ecd75dd 100755 --- a/nodejs/bin/www +++ b/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'); /** diff --git a/nodejs/package-lock.json b/nodejs/package-lock.json index afbf31b..1a3119f 100644 --- a/nodejs/package-lock.json +++ b/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": { diff --git a/nodejs/package.json b/nodejs/package.json index 4c7ca0c..651d45c 100755 --- a/nodejs/package.json +++ b/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", diff --git a/nodejs/routes/api_agent.js b/nodejs/routes/api_agent.js new file mode 100644 index 0000000..09abafe --- /dev/null +++ b/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' + } + })); + }); +};