From ceba377939a92ca69ef53aa3c0d53d958bb6e600 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Sat, 27 Sep 2025 13:54:12 +0200 Subject: [PATCH] Implemented starting and stopping listeners from the ImGui client. --- src/client/event/recv.nim | 77 -------------------- src/client/main.nim | 2 +- src/client/views/listeners.nim | 15 ++-- src/client/views/modals/startListener.nim | 1 - src/client/{event/send.nim => websocket.nim} | 41 +++++------ src/server/api/handlers.nim | 19 ++--- src/server/core/listener.nim | 4 +- src/server/core/server.nim | 36 +++++---- src/server/db/database.nim | 3 +- src/server/event/recv.nim | 40 ---------- src/server/{event/send.nim => websocket.nim} | 6 +- 11 files changed, 61 insertions(+), 183 deletions(-) delete mode 100644 src/client/event/recv.nim rename src/client/{event/send.nim => websocket.nim} (58%) delete mode 100644 src/server/event/recv.nim rename src/server/{event/send.nim => websocket.nim} (95%) diff --git a/src/client/event/recv.nim b/src/client/event/recv.nim deleted file mode 100644 index 53a2eff..0000000 --- a/src/client/event/recv.nim +++ /dev/null @@ -1,77 +0,0 @@ -import whisky -import times, tables -import ../views/[sessions, listeners, console, eventlog] -import ../../common/[types, utils, event] -export recvEvent - -#[ - Server -> Client -]# - - - - -# proc getMessageType*(message: Message): EventType = -# var unpacker = Unpacker.init(message.data) -# return cast[EventType](unpacker.getUint8()) - -# proc receiveAgentPayload*(message: Message): seq[byte] = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# return string.toBytes(unpacker.getDataWithLengthPrefix()) - -# proc receiveAgentConnection*(message: Message, sessions: ptr SessionsTableComponent) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let agent = Agent( -# agentId: Uuid.toString(unpacker.getUint32()), -# listenerId: Uuid.toString(unpacker.getUint32()), -# username: unpacker.getDataWithLengthPrefix(), -# hostname: unpacker.getDataWithLengthPrefix(), -# domain: unpacker.getDataWithLengthPrefix(), -# ip: unpacker.getDataWithLengthPrefix(), -# os: unpacker.getDataWithLengthPrefix(), -# process: unpacker.getDataWithLengthPrefix(), -# pid: int(unpacker.getUint32()), -# elevated: unpacker.getUint8() != 0, -# sleep: int(unpacker.getUint32()), -# tasks: @[], -# firstCheckin: cast[int64](unpacker.getUint32()).fromUnix().utc(), -# latestCheckin: cast[int64](unpacker.getUint32()).fromUnix().utc(), -# ) - -# sessions.agents.add(agent) - -# proc receiveAgentCheckin*(message: Message, sessions: ptr SessionsTableComponent)= -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let agentId = Uuid.toString(unpacker.getUint32()) -# let timestamp = cast[int64](unpacker.getUint32()) - -# # TODO: Update checkin - -# proc receiveConsoleItem*(message: Message, consoles: ptr Table[string, ConsoleComponent]) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let -# agentId = Uuid.toString(unpacker.getUint32()) -# logType = cast[LogType](unpacker.getUint8()) -# timestamp = cast[int64](unpacker.getUint32()) -# message = unpacker.getDataWithLengthPrefix() - -# consoles[][agentId].addItem(logType, message, timestamp) - -# proc receiveEventlogItem*(message: Message, eventlog: ptr EventlogComponent) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let -# logType = cast[LogType](unpacker.getUint8()) -# timestamp = cast[int64](unpacker.getUint32()) -# message = unpacker.getDataWithLengthPrefix() - -# eventlog[].addItem(logType, message, timestamp) \ No newline at end of file diff --git a/src/client/main.nim b/src/client/main.nim index 765b000..794c777 100644 --- a/src/client/main.nim +++ b/src/client/main.nim @@ -3,7 +3,7 @@ import tables, strutils, json, parsetoml import ./utils/appImGui import ./views/[dockspace, sessions, listeners, eventlog, console] import ../common/[types, utils] -import ./event/[send, recv] +import ./websocket import sugar diff --git a/src/client/views/listeners.nim b/src/client/views/listeners.nim index 4efc091..aec80a2 100644 --- a/src/client/views/listeners.nim +++ b/src/client/views/listeners.nim @@ -3,7 +3,7 @@ import imguin/[cimgui, glfw_opengl, simple] import ../utils/appImGui import ../../common/[types, utils] import ./modals/[startListener, generatePayload] -import ../event/send +import ../websocket import whisky type @@ -76,15 +76,14 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS var multiSelectIO = igBeginMultiSelect(ImGuiMultiSelectFlags_ClearOnEscape.int32 or ImGuiMultiSelectFlags_BoxSelect1d.int32, component.selection[].Size, int32(component.listeners.len())) ImGuiSelectionBasicStorage_ApplyRequests(component.selection, multiSelectIO) - for row in 0 ..< component.listeners.len(): + for i, listener in component.listeners: igTableNextRow(ImGuiTableRowFlags_None.int32, 0.0f) - let listener = component.listeners[row] if igTableSetColumnIndex(0): # Enable multi-select functionality - igSetNextItemSelectionUserData(row) - var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](row)) + igSetNextItemSelectionUserData(i) + var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)) discard igSelectable_Bool(listener.listenerId, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32, vec2(0.0f, 0.0f)) if igTableSetColumnIndex(1): @@ -101,11 +100,11 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS if igMenuItem("Stop", nil, false, true): # Update agents table with only non-selected ones var newListeners: seq[UIListener] = @[] - for i in 0 ..< component.listeners.len(): + for i, listener in component.listeners: if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)): - newListeners.add(component.listeners[i]) + newListeners.add(listener) else: - ws.sendStopListener(component.listeners[i].listenerId) + ws.sendStopListener(listener.listenerId) component.listeners = newListeners ImGuiSelectionBasicStorage_Clear(component.selection) diff --git a/src/client/views/modals/startListener.nim b/src/client/views/modals/startListener.nim index 877b49e..7aeca28 100644 --- a/src/client/views/modals/startListener.nim +++ b/src/client/views/modals/startListener.nim @@ -75,7 +75,6 @@ proc draw*(component: ListenerModalComponent): UIListener = igBeginDisabled(($(addr component.address[0]) == "") or (component.port <= 0)) if igButton("Start", vec2(availableSize.x * 0.5 - textSpacing * 0.5, 0.0f)): - result = UIListener( listenerId: generateUUID(), address: $(addr component.address[0]), diff --git a/src/client/event/send.nim b/src/client/websocket.nim similarity index 58% rename from src/client/event/send.nim rename to src/client/websocket.nim index 3d9ab34..8eb98a3 100644 --- a/src/client/event/send.nim +++ b/src/client/websocket.nim @@ -1,34 +1,29 @@ import whisky -import times, tables -import ../views/[sessions, listeners, console, eventlog] -import ../../common/[types, utils, serialize, event] -export sendHeartbeat +import times, tables, json +import ./views/[sessions, listeners, console, eventlog] +import ../common/[types, utils, serialize, event] +export sendHeartbeat, recvEvent #[ Client -> Server ]# proc sendStartListener*(ws: WebSocket, listener: UIListener) = - var packer = Packer.init() - - packer.add(cast[uint8](CLIENT_LISTENER_START)) - packer.add(string.toUUid(listener.listenerId)) - packer.addDataWithLengthPrefix(string.toBytes(listener.address)) - packer.add(cast[uint16](listener.port)) - packer.add(cast[uint8](listener.protocol)) - - let data = packer.pack() - - ws.send(Bytes.toString(data), BinaryMessage) + let event = Event( + eventType: CLIENT_LISTENER_START, + timestamp: now().toTime().toUnix(), + data: %listener + ) + ws.sendEvent(event) proc sendStopListener*(ws: WebSocket, listenerId: string) = - discard - # var packer = Packer.init() - - # packer.add(cast[uint8](CLIENT_LISTENER_STOP)) - # packer.add(string.toUuid(listenerId)) - # let data = packer.pack() - - # ws.send(Bytes.toString(data), BinaryMessage) + let event = Event( + eventType: CLIENT_LISTENER_STOP, + timestamp: now().toTime().toUnix(), + data: %*{ + "listenerId": listenerId + } + ) + ws.sendEvent(event) # proc sendAgentCommand*(ws: WebSocket, agentId: string, command: string) = # var packer = Packer.init() diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 1282501..6763b2a 100644 --- a/src/server/api/handlers.nim +++ b/src/server/api/handlers.nim @@ -4,7 +4,7 @@ import ../globals import ../db/database import ../protocol/packer import ../core/logger -import ../event/send +import ../websocket import ../../common/[types, utils, serialize] #[ @@ -21,26 +21,23 @@ proc register*(registrationData: seq[byte]): bool = # Validate that listener exists if not cq.dbListenerExists(agent.listenerId.toUpperAscii): - cq.error(fmt"{agent.ip} attempted to register to non-existent listener: {agent.listenerId}.", "\n") - return false + raise newException(CatchableError, fmt"{agent.ip} attempted to register to non-existent listener: {agent.listenerId}." & "\n") # Store agent in database if not cq.dbStoreAgent(agent): - cq.error(fmt"Failed to insert agent {agent.agentId} into database.", "\n") - return false + raise newException(CatchableError, fmt"Failed to insert agent {agent.agentId} into database." & "\n") # Create log directory if not cq.makeAgentLogDirectory(agent.agentId): - cq.error("Failed to create log directory.", "\n") - return false + raise newException(CatchableError, "Failed to create log directory.\n") cq.agents[agent.agentId] = agent cq.info("Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") + # Send new agent to client cq.client.sendAgent(agent) cq.client.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.") - return true except CatchableError as err: @@ -63,13 +60,11 @@ proc getTasks*(heartbeat: seq[byte]): seq[seq[byte]] = # Check if listener exists if not cq.dbListenerExists(listenerId): - cq.error(fmt"Task-retrieval request made to non-existent listener: {listenerId}.", "\n") - raise newException(ValueError, "Invalid listener.") + raise newException(ValueError, fmt"Task-retrieval request made to non-existent listener: {listenerId}." & "\n") # Check if agent exists if not cq.dbAgentExists(agentId): - cq.error(fmt"Task-retrieval request made to non-existent agent: {agentId}.", "\n") - raise newException(ValueError, "Invalid agent.") + raise newException(ValueError, fmt"Task-retrieval request made to non-existent agent: {agentId}." & "\n") # Update the last check-in date for the accessed agent cq.agents[agentId].latestCheckin = cast[int64](timestamp).fromUnix().local() diff --git a/src/server/core/listener.nim b/src/server/core/listener.nim index 06f7e50..200e0a5 100644 --- a/src/server/core/listener.nim +++ b/src/server/core/listener.nim @@ -8,7 +8,7 @@ import ../api/routes import ../db/database import ../core/logger import ../../common/[types, utils, profile] -import ../event/send +import ../websocket #[ Listener management @@ -86,6 +86,7 @@ proc listenerStart*(cq: Conquest, name: string, host: string, port: int, protoco raise newException(CatchableError, "Failed to store listener in database.") cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{$port}.") + cq.client.sendListener(listener) cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Started listener {name} on {host}:{$port}.") except CatchableError as err: @@ -132,6 +133,7 @@ proc restartListeners*(cq: Conquest) = cq.listeners[listener.listenerId] = listener cq.threads[listener.listenerId] = thread + cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Restarted listener {listener.listenerId} on {listener.address}:{$listener.port}.") cq.success("Restarted listener", fgGreen, fmt" {listener.listenerId} ", resetStyle, fmt"on {listener.address}:{$listener.port}.") except CatchableError as err: diff --git a/src/server/core/server.nim b/src/server/core/server.nim index d94a2c8..b58e2db 100644 --- a/src/server/core/server.nim +++ b/src/server/core/server.nim @@ -1,12 +1,12 @@ -import prompt, terminal, argparse, parsetoml, times +import prompt, terminal, argparse, parsetoml, times, json import strutils, strformat, system, tables import ./[agent, listener, builder] import ../globals import ../db/database import ../core/logger -import ../../common/[types, crypto, utils, profile] -import ../event/[recv, send] +import ../../common/[types, crypto, utils, profile, event] +import ../websocket import mummy, mummy/routers #[ @@ -170,17 +170,23 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {. # Continuously send heartbeat messages ws.sendHeartbeat() - # case message.getMessageType(): - # of CLIENT_AGENT_COMMAND: - # discard - # of CLIENT_LISTENER_START: - # message.receiveStartListener() - # of CLIENT_LISTENER_STOP: - # discard - # # message.receiveStopListener() - # of CLIENT_AGENT_BUILD: - # discard - # else: discard + let event = message.recvEvent() + + case event.eventType: + of CLIENT_AGENT_COMMAND: + discard + + of CLIENT_LISTENER_START: + let listener = event.data.to(UIListener) + cq.listenerStart(listener.listenerId, listener.address, listener.port, listener.protocol) + + of CLIENT_LISTENER_STOP: + let listenerId = event.data["listenerId"].getStr() + cq.listenerStop(listenerId) + + of CLIENT_AGENT_BUILD: + discard + else: discard of ErrorEvent: discard @@ -210,7 +216,7 @@ proc startServer*(profilePath: string) = try: # Initialize framework context # Load and parse profile - let profile = parseFile(profilePath) + let profile = parsetoml.parseFile(profilePath) cq = Conquest.init(profile) cq.info("Using profile \"", profile.getString("name"), "\" (", profilePath ,").") diff --git a/src/server/db/database.nim b/src/server/db/database.nim index da8b4db..2084ae4 100644 --- a/src/server/db/database.nim +++ b/src/server/db/database.nim @@ -35,8 +35,7 @@ proc dbInit*(cq: Conquest) = sleep INTEGER DEFAULT 10, firstCheckin DATETIME NOT NULL, latestCheckin DATETIME NOT NULL, - sessionKey BLOB NOT NULL, - FOREIGN KEY (listener) REFERENCES listeners(name) + sessionKey BLOB NOT NULL ); """) diff --git a/src/server/event/recv.nim b/src/server/event/recv.nim deleted file mode 100644 index 73c3b4f..0000000 --- a/src/server/event/recv.nim +++ /dev/null @@ -1,40 +0,0 @@ -import mummy -import times, tables, json -import ./send -import ../globals -import ../core/[task, listener] -import ../../common/[types, utils, serialize, event] - -#[ - Client -> Server -]# -# proc getMessageType*(message: Message): EventType = -# var unpacker = Unpacker.init(message.data) -# return cast[EventType](unpacker.getUint8()) - -# proc receiveStartListener*(message: Message) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let -# listenerId = Uuid.toString(unpacker.getUint32()) -# address = unpacker.getDataWithLengthPrefix() -# port = int(unpacker.getUint16()) -# protocol = cast[Protocol](unpacker.getUint8()) -# cq.listenerStart(listenerId, address, port, protocol) - -# proc receiveStopListener*(message: Message) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let listenerId = Uuid.toString(unpacker.getUint32()) -# cq.listenerStop(listenerId) - -# proc receiveAgentCommand*(message: Message) = -# var unpacker = Unpacker.init(message.data) - -# discard unpacker.getUint8() -# let -# agentId = Uuid.toString(unpacker.getUint32()) -# command = unpacker.getDataWithLengthPrefix() - diff --git a/src/server/event/send.nim b/src/server/websocket.nim similarity index 95% rename from src/server/event/send.nim rename to src/server/websocket.nim index d9558c5..d149beb 100644 --- a/src/server/event/send.nim +++ b/src/server/websocket.nim @@ -1,8 +1,8 @@ import mummy import times, tables, json, base64, parsetoml -import ../utils -import ../../common/[types, utils, serialize, event] -export sendHeartbeat +import ./utils +import ../common/[types, utils, serialize, event] +export sendHeartbeat, recvEvent #[ Server -> Client