54 lines
1.8 KiB
JavaScript
54 lines
1.8 KiB
JavaScript
'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'
|
|
}
|
|
}));
|
|
});
|
|
};
|