Implemented right-click menu to remove or download loot (files/screenshots).

src/client/core/websocket.nim

@@ -59,3 +59,23 @@ proc sendAgentTask*(connection: WsConnection, agentId: string, command: string,
         }
     )
     connection.ws.sendEvent(event, connection.sessionKey)
+
+proc sendRemoveLoot*(connection: WsConnection, lootId: string) = 
+    let event = Event(
+        eventType: CLIENT_LOOT_REMOVE, 
+        timestamp: now().toTime().toUnix(),
+        data: %*{
+            "lootId": lootId
+        }
+    )
+    connection.ws.sendEvent(event, connection.sessionKey)
+
+proc sendDownloadLoot*(connection: WsConnection, lootId: string) = 
+    let event = Event(
+        eventType: CLIENT_LOOT_SYNC, 
+        timestamp: now().toTime().toUnix(),
+        data: %*{
+            "lootId": lootId
+        }
+    )
+    connection.ws.sendEvent(event, connection.sessionKey)

src/client/main.nim

@@ -157,10 +157,18 @@ proc main(ip: string = "localhost", port: int = 37573) =
                 lootDownloads.items.add(lootItem)
             of SCREENSHOT:
                 lootScreenshots.addItem(lootItem)
-
             else: discard 
 
-        of CLIENT_SYNC_LOOT: 
+        of CLIENT_LOOT_SYNC:
+            let path = event.data["path"].getStr()
+            let file = decode(event.data["loot"].getStr())
+            try: 
+                # TODO: Using native file dialogs to have the client select the output file path (does not work in WSL) 
+                # let outFilePath = callDialogFileSave("Save Payload") 
+                writeFile(path & "_download", file)
+            except IOError:
+                discard 
+
             discard 
     
         else: discard 
@@ -169,8 +177,8 @@ proc main(ip: string = "localhost", port: int = 37573) =
         if showSessionsTable: sessionsTable.draw(addr showSessionsTable)   
         if showListeners: listenersTable.draw(addr showListeners, connection)
         if showEventlog: eventlog.draw(addr showEventlog)
-        if showDownloads: lootDownloads.draw(addr showDownloads)
+        if showDownloads: lootDownloads.draw(addr showDownloads, connection)
-        if showScreenshots: lootScreenshots.draw(addr showScreenshots)
+        if showScreenshots: lootScreenshots.draw(addr showScreenshots, connection)
 
         # Show console windows
         var newConsoleTable: Table[string, ConsoleComponent]

src/client/utils/globals.nim

@@ -3,7 +3,7 @@ import ../utils/fonticon/IconsFontAwesome6
 const CONQUEST_ROOT* {.strdefine.} = ""
 
 const WIDGET_SESSIONS* =  " " & ICON_FA_LIST & " " & "Sessions [Table View]"
-const WIDGET_LISTENERS* = " " & ICON_FA_HEADPHONES & " " & "Listeners"
+const WIDGET_LISTENERS* = " " & ICON_FA_SATELLITE_DISH & " " & "Listeners"
 const WIDGET_EVENTLOG* = "Eventlog"
 const WIDGET_DOWNLOADS* = " " & ICON_FA_DOWNLOAD & " " & "Downloads"
 const WIDGET_SCREENSHOTS* = " " & ICON_FA_IMAGE & " " & "Screenshots"

src/client/views/loot/downloads.nim

@@ -2,6 +2,7 @@ import strformat, strutils, times, os
 import imguin/[cimgui, glfw_opengl, simple]
 import ../../utils/[appImGui, colors]
 import ../../../common/[types, utils]
+import ../../core/websocket
 
 type
     DownloadsComponent* = ref object of RootObj
@@ -16,7 +17,7 @@ proc LootDownloads*(title: string): DownloadsComponent =
     result.items = @[]
     result.selectedIndex = -1
 
-proc draw*(component: DownloadsComponent, showComponent: ptr bool) =
+proc draw*(component: DownloadsComponent, showComponent: ptr bool, connection: WsConnection) =
     igBegin(component.title, showComponent, 0)
     defer: igEnd()
 
@@ -61,6 +62,10 @@ proc draw*(component: DownloadsComponent, showComponent: ptr bool) =
                     let isSelected = component.selectedIndex == i
                     if igSelectable_Bool(item.lootId.cstring, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32 or ImGuiSelectableFlags_AllowOverlap.int32, vec2(0, 0)):
                         component.selectedIndex = i
+                    
+                    if igIsItemHovered(ImGuiHoveredFlags_None.int32) and igIsMouseClicked_Bool(ImGuiMouseButton_Right.int32, false):
+                        component.selectedIndex = i
+                    
                     igPopID()
 
                 if igTableSetColumnIndex(1):
@@ -78,6 +83,23 @@ proc draw*(component: DownloadsComponent, showComponent: ptr bool) =
                 if igTableSetColumnIndex(5):
                     igText($item.size)
                                         
+            # Handle right-click context menu
+            if component.selectedIndex >= 0 and component.selectedIndex < component.items.len and igBeginPopupContextWindow("Downloads", ImGui_PopupFlags_MouseButtonRight.int32): 
+                let item = component.items[component.selectedIndex]
+
+                if igMenuItem("Download", nil, false, true):                     
+                    # Task team server to download file
+                    connection.sendDownloadLoot(item.lootId)
+                    igCloseCurrentPopup()
+
+                if igMenuItem("Remove", nil, false, true): 
+                    # Task team server to remove the loot item 
+                    connection.sendRemoveLoot(item.lootId)
+                    component.items.delete(component.selectedIndex)
+                    igCloseCurrentPopup()
+
+                igEndPopup()
+            
             igEndTable()
         
     igEndChild()
@@ -93,9 +115,11 @@ proc draw*(component: DownloadsComponent, showComponent: ptr bool) =
             igSameLine(0.0f, 0.0f)
             igText(item.path.extractFilename().replace("C_", "C:/").replace("_", "/"))
             
+            igDummy(vec2(0.0f, 5.0f))
             igSeparator()
+            igDummy(vec2(0.0f, 5.0f)) 
 
-            igText(item.data)
+            igTextUnformatted(item.data, nil)
             
         else:
             igText("Select item to preview contents")

src/client/views/loot/screenshots.nim

@@ -2,6 +2,7 @@ import strformat, strutils, times, os, tables
 import imguin/[cimgui, glfw_opengl, simple]
 import ../../utils/[appImGui, colors]
 import ../../../common/[types, utils]
+import ../../core/websocket
 
 type
     ScreenshotTexture* = ref object 
@@ -33,7 +34,7 @@ proc addItem*(component: ScreenshotsComponent, screenshot: LootItem) =
         height: height
     )
 
-proc draw*(component: ScreenshotsComponent, showComponent: ptr bool) =
+proc draw*(component: ScreenshotsComponent, showComponent: ptr bool, connection: WsConnection) =
     igBegin(component.title, showComponent, 0)
     defer: igEnd()
 
@@ -43,7 +44,7 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool) =
         
     # Left panel (file table) 
     let childFlags = ImGui_ChildFlags_ResizeX.int32 or ImGui_ChildFlags_NavFlattened.int32
-    if igBeginChild_Str("##Left", vec2(availableSize.x * 0.66f, 0.0f), childFlags, ImGui_WindowFlags_None.int32):
+    if igBeginChild_Str("##Left", vec2(availableSize.x * 0.5f, 0.0f), childFlags, ImGui_WindowFlags_None.int32):
 
         let tableFlags = (
             ImGui_TableFlags_Resizable.int32 or
@@ -77,6 +78,10 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool) =
                     let isSelected = component.selectedIndex == i
                     if igSelectable_Bool(item.lootId.cstring, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32 or ImGuiSelectableFlags_AllowOverlap.int32, vec2(0, 0)):
                         component.selectedIndex = i
+                
+                    if igIsItemHovered(ImGuiHoveredFlags_None.int32) and igIsMouseClicked_Bool(ImGuiMouseButton_Right.int32, false):
+                        component.selectedIndex = i
+                    
                     igPopID()                
 
                 if igTableSetColumnIndex(1):
@@ -91,6 +96,23 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool) =
                 if igTableSetColumnIndex(4):
                     igText($item.size)
 
+            # Handle right-click context menu
+            if component.selectedIndex >= 0 and component.selectedIndex < component.items.len and igBeginPopupContextWindow("Downloads", ImGui_PopupFlags_MouseButtonRight.int32): 
+                let item = component.items[component.selectedIndex]
+
+                if igMenuItem("Download", nil, false, true):                     
+                    # Task team server to download file
+                    connection.sendDownloadLoot(item.lootId)
+                    igCloseCurrentPopup()
+
+                if igMenuItem("Remove", nil, false, true): 
+                    # Task team server to remove the loot item 
+                    connection.sendRemoveLoot(item.lootId)
+                    component.items.delete(component.selectedIndex)
+                    igCloseCurrentPopup()
+
+                igEndPopup()
+
             igEndTable()
         
     igEndChild()
@@ -106,5 +128,5 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool) =
             igImage(ImTextureRef(internal_TexData: nil, internal_TexID: texture.textureId), vec2(texture.width, texture.height), vec2(0, 0), vec2(1, 1))
             
         else:
-            igText("Select item to preview contents")
+            igText("Select item for preview.")
     igEndChild()

src/common/types.nim

@@ -255,7 +255,8 @@ type
         CLIENT_AGENT_TASK = 2'u8            # Instruct TS to send queue a command for a specific agent
         CLIENT_LISTENER_START = 3'u8        # Start a listener on the TS
         CLIENT_LISTENER_STOP = 4'u8         # Stop a listener
-        CLIENT_REQUEST_SYNC = 5'u8          # Request to download a file/screenshot to the client
+        CLIENT_LOOT_REMOVE = 5'u8           # Remove loot on the team server
+        CLIENT_LOOT_SYNC = 6'u8         # Request to download a file/screenshot to the client
 
         # Sent by team server
         CLIENT_PROFILE = 100'u8             # Team server profile and configuration 
@@ -267,7 +268,6 @@ type
         CLIENT_EVENTLOG_ITEM = 106'u8       # Add entry to the eventlog   
         CLIENT_BUILDLOG_ITEM = 107'u8       # Add entry to the build log
         CLIENT_LOOT_ADD = 108'u8            # Add file or screenshot stored on the team server to preview on the client
-        CLIENT_SYNC_LOOT = 109'u8           # Download a file/screenshot to the operator desktop
 
     Event* = object 
         eventType*: EventType               

src/server/api/handlers.nim

@@ -111,11 +111,10 @@ proc handleResult*(resultData: seq[byte]) =
                 if int(taskResult.length) > 0:
                     cq.client.sendConsoleItem(agentId, LOG_INFO, "Output:") 
                     cq.info("Output:")
-                    cq.client.sendConsoleItem(agentId, LOG_OUTPUT, Bytes.toString(taskResult.data))
 
                     # Split result string on newline to keep formatting
                     for line in Bytes.toString(taskResult.data).split("\n"):
-                        cq.output(line)
+                        cq.client.sendConsoleItem(agentId, LOG_OUTPUT, line)
 
             of RESULT_BINARY:
                 # Write binary data to a file 
@@ -143,16 +142,11 @@ proc handleResult*(resultData: seq[byte]) =
                     host: cq.agents[agentId].hostname
                 )
 
-                if lootItem.itemType == SCREENSHOT: 
-                    lootItem.data = createThumbnail(readFile(downloadPath))    # Create a smaller thumbnail version of the screenshot for better transportability
-                elif lootItem.itemType == DOWNLOAD:
-                    lootItem.data = readFile(downloadPath)                     # Read downloaded file
-
                 # Store loot in database 
                 if not cq.dbStoreLoot(lootItem): 
                     raise newException(ValueError, fmt"Failed to store loot in database." & "\n")
 
-                # Send packet to client to display file/screenshot in the UI
+                # Send loot to client to display file/screenshot in the UI
                 cq.client.sendLoot(lootItem)
 
                 cq.output(fmt"File downloaded to {downloadPath} ({$fileData.len()} bytes).", "\n")

src/server/core/websocket.nim

@@ -1,4 +1,5 @@
-import times, json, base64, parsetoml, strformat
+import times, json, base64, parsetoml, strformat, pixie
+import stb_image/write as stbiw
 import ./logger
 import ../../common/[types, utils, event]
 export sendHeartbeat, recvEvent
@@ -144,7 +145,38 @@ proc sendBuildlogItem*(client: WsConnection, logType: LogType, message: string)
     if client != nil: 
         client.ws.sendEvent(event, client.sessionKey)
 
+proc createThumbnail(data: string, maxWidth: int = 1024, quality: int = 90): string =
+    let img: Image = decodeImage(data)
+    
+    let aspectRatio = img.height.float / img.width.float
+    let 
+        width = min(maxWidth, img.width)
+        height = int(width.float * aspectRatio)
+
+    # Resize image
+    let thumbnail = img.resize(width, height)
+
+    # Convert to JPEG image for smaller file size
+    var rgbaData = newSeq[byte](width * height * 4)
+    var i = 0
+    for y in 0..<height:
+        for x in 0..<width:
+            let color = thumbnail[x, y]
+            rgbaData[i] = color.r
+            rgbaData[i + 1] = color.g
+            rgbaData[i + 2] = color.b
+            rgbaData[i + 3] = color.a
+            i += 4
+    
+    return Bytes.toString(stbiw.writeJPG(width, height, 4, rgbaData, quality))
+
 proc sendLoot*(client: WsConnection, loot: LootItem) = 
+    var data: string
+    if loot.itemType == SCREENSHOT: 
+        loot.data = createThumbnail(readFile(loot.path))    # Create a smaller thumbnail version of the screenshot for better transportability
+    elif loot.itemType == DOWNLOAD:
+        loot.data = readFile(loot.path)                     # Read downloaded file
+    
     let event = Event(
         eventType: CLIENT_LOOT_ADD,
         timestamp: now().toTime().toUnix(),
@@ -152,3 +184,15 @@ proc sendLoot*(client: WsConnection, loot: LootItem) =
     )
     if client != nil: 
         client.ws.sendEvent(event, client.sessionKey)
+
+proc sendLootSync*(client: WsConnection, path: string, file: string) = 
+    let event = Event(
+        eventType: CLIENT_LOOT_SYNC,
+        timestamp: now().toTime().toUnix(),
+        data: %*{
+            "path": path,
+            "loot": encode(file)
+        }
+    )
+    if client != nil: 
+        client.ws.sendEvent(event, client.sessionKey)

src/server/db/dbLoot.nim

@@ -1,7 +1,6 @@
-import strutils, system, terminal, tiny_sqlite, pixie
+import system, terminal, tiny_sqlite
-import stb_image/write as stbiw
 import ../core/logger
-import ../../common/[types, utils]
+import ../../common/types
 
 proc dbStoreLoot*(cq: Conquest, loot: LootItem): bool = 
     try: 
@@ -19,31 +18,6 @@ proc dbStoreLoot*(cq: Conquest, loot: LootItem): bool =
     
     return true
 
-proc createThumbnail*(data: string, maxWidth: int = 1024, quality: int = 90): string =
-    let img: Image = decodeImage(data)
-    
-    let aspectRatio = img.height.float / img.width.float
-    let 
-        width = min(maxWidth, img.width)
-        height = int(width.float * aspectRatio)
-
-    # Resize image
-    let thumbnail = img.resize(width, height)
-
-    # Convert to JPEG image for smaller file size
-    var rgbaData = newSeq[byte](width * height * 4)
-    var i = 0
-    for y in 0..<height:
-        for x in 0..<width:
-            let color = thumbnail[x, y]
-            rgbaData[i] = color.r
-            rgbaData[i + 1] = color.g
-            rgbaData[i + 2] = color.b
-            rgbaData[i + 3] = color.a
-            i += 4
-    
-    return Bytes.toString(stbiw.writeJPG(width, height, 4, rgbaData, quality))
-
 proc dbGetLoot*(cq: Conquest): seq[LootItem] = 
     var loot: seq[LootItem] = @[]
 
@@ -53,7 +27,7 @@ proc dbGetLoot*(cq: Conquest): seq[LootItem] =
         for row in conquestDb.iterate("SELECT lootId, itemType, agentId, host, path, timestamp, size FROM loot;"):
             let (lootId, itemType, agentId, host, path, timestamp, size) = row.unpack((string, int, string, string, string, int64, int))
 
-            var l = LootItem(
+            let l = LootItem(
                 lootId: lootId, 
                 itemType: cast[LootItemType](itemType),
                 agentId: agentId,
@@ -63,11 +37,6 @@ proc dbGetLoot*(cq: Conquest): seq[LootItem] =
                 size: size
             )
 
-            if l.itemType == SCREENSHOT: 
-                l.data = createThumbnail(readFile(path))    # Create a smaller thumbnail version of the screenshot for better transportability
-            elif l.itemType == DOWNLOAD:
-                l.data = readFile(path)                     # Read downloaded file
- 
             loot.add(l)
 
         conquestDb.close()
@@ -75,3 +44,33 @@ proc dbGetLoot*(cq: Conquest): seq[LootItem] =
         cq.error(getCurrentExceptionMsg())
 
     return loot
+
+proc dbGetLootById*(cq: Conquest, lootId: string): LootItem = 
+    try: 
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+        for row in conquestDb.iterate("SELECT lootId, itemType, agentId, host, path, timestamp, size FROM loot WHERE lootId = ?;", lootId):
+            let (id, itemType, agentId, host, path, timestamp, size) = row.unpack((string, int, string, string, string, int64, int))
+            result = LootItem(
+                lootId: id,
+                itemType: cast[LootItemType](itemType),
+                agentId: agentId,
+                host: host,
+                path: path,
+                timestamp: timestamp,
+                size: size
+            )
+        conquestDb.close()
+    except: 
+        cq.error(getCurrentExceptionMsg())
+
+proc dbDeleteLootById*(cq: Conquest, lootId: string): bool =
+    try: 
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+
+        conquestDb.exec("DELETE FROM loot WHERE lootId = ?", lootId)
+
+        conquestDb.close()
+    except: 
+        return false
+    
+    return true

src/server/main.nim

@@ -102,6 +102,14 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
                 if payload.len() != 0: 
                     cq.client.sendAgentPayload(payload)
 
+            of CLIENT_LOOT_REMOVE: 
+                if not cq.dbDeleteLootById(event.data["lootId"].getStr()): 
+                    cq.client.sendEventlogItem(LOG_ERROR, "Failed to delete loot.")
+
+            of CLIENT_LOOT_SYNC: 
+                let path = cq.dbGetLootById(event.data["lootId"].getStr()).path
+                cq.client.sendLootSync(path, readFile(path))
+
             else: discard
 
         of ErrorEvent: