diff --git a/src/client/core/websocket.nim b/src/client/core/websocket.nim index a6ca52c..2b60fda 100644 --- a/src/client/core/websocket.nim +++ b/src/client/core/websocket.nim @@ -70,9 +70,9 @@ proc sendRemoveLoot*(connection: WsConnection, lootId: string) = ) connection.ws.sendEvent(event, connection.sessionKey) -proc sendDownloadLoot*(connection: WsConnection, lootId: string) = +proc sendGetLoot*(connection: WsConnection, lootId: string) = let event = Event( - eventType: CLIENT_LOOT_SYNC, + eventType: CLIENT_LOOT_GET, timestamp: now().toTime().toUnix(), data: %*{ "lootId": lootId diff --git a/src/client/main.nim b/src/client/main.nim index c995544..b8d6f6e 100644 --- a/src/client/main.nim +++ b/src/client/main.nim @@ -157,18 +157,20 @@ proc main(ip: string = "localhost", port: int = 37573) = of DOWNLOAD: lootDownloads.items.add(lootItem) of SCREENSHOT: - lootScreenshots.addItem(lootItem) + lootScreenshots.items.add(lootItem) else: discard - 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 + of CLIENT_LOOT_DATA: + let + lootItem = event.data["loot"].to(LootItem) + data = decode(event.data["data"].getStr()) + + case lootItem.itemType: + of DOWNLOAD: + lootDownloads.contents[lootItem.lootId] = data + of SCREENSHOT: + lootScreenshots.addTexture(lootItem.lootId, data) + else: discard else: discard diff --git a/src/client/views/console.nim b/src/client/views/console.nim index 4fc4b0b..38f328e 100644 --- a/src/client/views/console.nim +++ b/src/client/views/console.nim @@ -224,9 +224,7 @@ proc draw*(component: ConsoleComponent, connection: WsConnection) = Problems I encountered with other approaches (Multi-line Text Input, TextEditor, ...): - https://github.com/ocornut/imgui/issues/383#issuecomment-2080346129 - - https://github.com/ocornut/imgui/issues/950 - - Huge thanks to @dinau for implementing ImGuiTextSelect into imguin very rapidly after I requested it. + - https://github.com/ocornut/imgui/issues/950 ]# let consolePadding: float = 10.0f let footerHeight = (consolePadding * 2) + (igGetStyle().ItemSpacing.y + igGetFrameHeightWithSpacing()) * 0.75f diff --git a/src/client/views/loot/downloads.nim b/src/client/views/loot/downloads.nim index d608e90..d53c0f9 100644 --- a/src/client/views/loot/downloads.nim +++ b/src/client/views/loot/downloads.nim @@ -1,4 +1,4 @@ -import strformat, strutils, times, os +import strformat, strutils, times, os, tables import imguin/[cimgui, glfw_opengl, simple] import ../../utils/[appImGui, colors] import ../../../common/[types, utils] @@ -8,6 +8,7 @@ type DownloadsComponent* = ref object of RootObj title: string items*: seq[LootItem] + contents*: Table[string, string] selectedIndex: int @@ -15,6 +16,7 @@ proc LootDownloads*(title: string): DownloadsComponent = result = new DownloadsComponent result.title = title result.items = @[] + result.contents = initTable[string, string]() result.selectedIndex = -1 proc draw*(component: DownloadsComponent, showComponent: ptr bool, connection: WsConnection) = @@ -88,8 +90,14 @@ proc draw*(component: DownloadsComponent, showComponent: ptr bool, connection: W let item = component.items[component.selectedIndex] if igMenuItem("Download", nil, false, true): - # Task team server to download file - connection.sendDownloadLoot(item.lootId) + # Download file + try: + # TODO: Use native dialogs to select the download location + let path = item.path & ".download" + let data = component.contents[item.lootId] + writeFile(path, data) + except IOError: + discard igCloseCurrentPopup() if igMenuItem("Remove", nil, false, true): @@ -111,15 +119,20 @@ proc draw*(component: DownloadsComponent, showComponent: ptr bool, connection: W if component.selectedIndex >= 0 and component.selectedIndex < component.items.len: let item = component.items[component.selectedIndex] - igText(fmt("[{item.host}] ")) - 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)) + if not component.contents.hasKey(item.lootId): + connection.sendGetLoot(item.lootId) + component.contents[item.lootId] = "" # Ensure that the sendGetLoot() function is sent only once by setting a value for the table key - igTextUnformatted(item.data, nil) + else: + igText(fmt("[{item.host}] ")) + 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)) + + igTextUnformatted(component.contents[item.lootId], 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 d6dfdb5..1b72288 100644 --- a/src/client/views/loot/screenshots.nim +++ b/src/client/views/loot/screenshots.nim @@ -7,12 +7,13 @@ import ../../core/websocket type ScreenshotTexture* = ref object textureId*: GLuint + data*: string width: int height: int ScreenshotsComponent* = ref object of RootObj title: string - items: seq[LootItem] + items*: seq[LootItem] selectedIndex: int textures: Table[string, ScreenshotTexture] @@ -23,13 +24,12 @@ proc LootScreenshots*(title: string): ScreenshotsComponent = result.selectedIndex = -1 result.textures = initTable[string, ScreenshotTexture]() -proc addItem*(component: ScreenshotsComponent, screenshot: LootItem) = - component.items.add(screenshot) - +proc addTexture*(component: ScreenshotsComponent, lootId: string, data: string) = var textureId: GLuint - let (width, height) = loadTextureFromBytes(string.toBytes(screenshot.data), textureId) - component.textures[screenshot.path] = ScreenshotTexture( + let (width, height) = loadTextureFromBytes(string.toBytes(data), textureId) + component.textures[lootId] = ScreenshotTexture( textureId: textureId, + data: data, width: width, height: height ) @@ -98,11 +98,18 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool, connection: # 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) + # Download screenshot + try: + # TODO: Use native dialogs to select the download location + let path = item.path & "_download.jpeg" + let data = component.textures[item.lootId].data + writeFile(path, data) + except IOError: + discard igCloseCurrentPopup() if igMenuItem("Remove", nil, false, true): @@ -122,10 +129,20 @@ proc draw*(component: ScreenshotsComponent, showComponent: ptr bool, connection: if igBeginChild_Str("##Preview", vec2(0.0f, 0.0f), ImGui_ChildFlags_Borders.int32, ImGui_WindowFlags_None.int32): if component.selectedIndex >= 0 and component.selectedIndex < component.items.len: - let item = component.items[component.selectedIndex] - let texture = component.textures[item.path] - igImage(ImTextureRef(internal_TexData: nil, internal_TexID: texture.textureId), vec2(texture.width, texture.height), vec2(0, 0), vec2(1, 1)) + let item = component.items[component.selectedIndex] + + # Check if the texture for the loot item has already been loaded from the team server + # If the texture doesn't exist yet, send a request to the team server to retrieve and render it + if not component.textures.hasKey(item.lootId): + connection.sendGetLoot(item.lootId) + component.textures[item.lootId] = nil # Ensure that the sendGetLoot() function is sent only once by setting a value for the table key + + # Display the image preview + else: + let texture = component.textures[item.lootId] + if not texture.isNil(): + igImage(ImTextureRef(internal_TexData: nil, internal_TexID: texture.textureId), vec2(texture.width, texture.height), vec2(0, 0), vec2(1, 1)) else: igText("Select item for preview.") diff --git a/src/common/types.nim b/src/common/types.nim index 83f28d9..01adf26 100644 --- a/src/common/types.nim +++ b/src/common/types.nim @@ -258,8 +258,8 @@ type CLIENT_LISTENER_START = 3'u8 # Start a listener on the TS CLIENT_LISTENER_STOP = 4'u8 # Stop a listener CLIENT_LOOT_REMOVE = 5'u8 # Remove loot on the team server - CLIENT_LOOT_SYNC = 6'u8 # Download a file/screenshot to the client - + CLIENT_LOOT_GET = 6'u8 # Request file/screenshot from the team server for preview or download + # Sent by team server CLIENT_PROFILE = 100'u8 # Team server profile and configuration CLIENT_LISTENER_ADD = 101'u8 # Add listener to listeners table @@ -269,7 +269,8 @@ type CLIENT_CONSOLE_ITEM = 105'u8 # Add entry to a agent's console 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_LOOT_ADD = 108'u8 # Add file or screenshot stored on the team server to preview on the client, only sends metadata and not the actual file content + CLIENT_LOOT_DATA = 109'u8 # Send file/screenshot bytes to the client to display as preview or to download to the client desktop Event* = object eventType*: EventType @@ -363,4 +364,3 @@ type path*: string timestamp*: int64 size*: int - data*: string # Image bytes or file content (binary data prefixed with length) diff --git a/src/server/core/websocket.nim b/src/server/core/websocket.nim index 4279220..3b9cfbf 100644 --- a/src/server/core/websocket.nim +++ b/src/server/core/websocket.nim @@ -171,12 +171,6 @@ proc createThumbnail(data: string, maxHeight: int = 1024, quality: int = 80): st 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(), @@ -185,13 +179,13 @@ proc sendLoot*(client: WsConnection, loot: LootItem) = if client != nil: client.ws.sendEvent(event, client.sessionKey) -proc sendLootSync*(client: WsConnection, path: string, file: string) = +proc sendLootData*(client: WsConnection, loot: LootItem, data: string) = let event = Event( - eventType: CLIENT_LOOT_SYNC, + eventType: CLIENT_LOOT_DATA, timestamp: now().toTime().toUnix(), data: %*{ - "path": path, - "loot": encode(file) + "loot": %loot, + "data": encode(data) } ) if client != nil: diff --git a/src/server/main.nim b/src/server/main.nim index d340346..c3bd4a9 100644 --- a/src/server/main.nim +++ b/src/server/main.nim @@ -67,7 +67,7 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {. for id, agent in cq.agents: cq.client.sendAgent(agent) - # Downloads & Screenshots + # Downloads & Screenshots metadata for lootItem in cq.dbGetLoot(): cq.client.sendLoot(lootItem) @@ -106,9 +106,9 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {. 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)) + of CLIENT_LOOT_GET: + let loot = cq.dbGetLootById(event.data["lootId"].getStr()) + cq.client.sendLootData(loot, readFile(loot.path)) else: discard