diff --git a/src/client/core/websocket.nim b/src/client/core/websocket.nim index 04d9fae..a6ca52c 100644 --- a/src/client/core/websocket.nim +++ b/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) \ No newline at end of file diff --git a/src/client/main.nim b/src/client/main.nim index 78a7c25..c86820b 100644 --- a/src/client/main.nim +++ b/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 showScreenshots: lootScreenshots.draw(addr showScreenshots) + if showDownloads: lootDownloads.draw(addr showDownloads, connection) + if showScreenshots: lootScreenshots.draw(addr showScreenshots, connection) # Show console windows var newConsoleTable: Table[string, ConsoleComponent] diff --git a/src/client/utils/globals.nim b/src/client/utils/globals.nim index 9ca1dfd..d231dcb 100644 --- a/src/client/utils/globals.nim +++ b/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" diff --git a/src/client/views/loot/downloads.nim b/src/client/views/loot/downloads.nim index 08c3014..d608e90 100644 --- a/src/client/views/loot/downloads.nim +++ b/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,7 +62,11 @@ 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 - igPopID() + + if igIsItemHovered(ImGuiHoveredFlags_None.int32) and igIsMouseClicked_Bool(ImGuiMouseButton_Right.int32, false): + component.selectedIndex = i + + igPopID() if igTableSetColumnIndex(1): igText(item.agentId) @@ -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() - - igText(item.data) + igDummy(vec2(0.0f, 5.0f)) + + igTextUnformatted(item.data, nil) else: igText("Select item to preview contents") diff --git a/src/client/views/loot/screenshots.nim b/src/client/views/loot/screenshots.nim index 2337081..07962d9 100644 --- a/src/client/views/loot/screenshots.nim +++ b/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() \ No newline at end of file diff --git a/src/common/types.nim b/src/common/types.nim index cb92564..60c3267 100644 --- a/src/common/types.nim +++ b/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 diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 574c7c7..9f0babf 100644 --- a/src/server/api/handlers.nim +++ b/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") diff --git a/src/server/core/websocket.nim b/src/server/core/websocket.nim index 605d1a5..57f5fcc 100644 --- a/src/server/core/websocket.nim +++ b/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..