Screenshots and downloads are now only retrieved once from the team server when the user selects them in the client for preview, which leads to faster start-up times and less blocking UI.

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

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:
+        of CLIENT_LOOT_DATA:
-            let path = event.data["path"].getStr()
+            let
-            let file = decode(event.data["loot"].getStr())
+                lootItem = event.data["loot"].to(LootItem)
-            try: 
+                data = decode(event.data["data"].getStr())
-                # TODO: Using native file dialogs to have the client select the output file path (does not work in WSL) 
+            
-                # let outFilePath = callDialogFileSave("Save Payload") 
+            case lootItem.itemType: 
-                writeFile(path & "_download", file)
+            of DOWNLOAD: 
-            except IOError:
+                lootDownloads.contents[lootItem.lootId] = data
-                discard 
+            of SCREENSHOT: 
+                lootScreenshots.addTexture(lootItem.lootId, data)
+            else: discard 
     
         else: discard 
     

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
+            - https://github.com/ocornut/imgui/issues/950        
-        
-        Huge thanks to @dinau for implementing ImGuiTextSelect into imguin very rapidly after I requested it.
     ]#
     let consolePadding: float = 10.0f 
     let footerHeight = (consolePadding * 2) + (igGetStyle().ItemSpacing.y + igGetFrameHeightWithSpacing()) * 0.75f

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
+                    # Download file
-                    connection.sendDownloadLoot(item.lootId)
+                    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}] "))
+            if not component.contents.hasKey(item.lootId): 
-            igSameLine(0.0f, 0.0f)
+                connection.sendGetLoot(item.lootId)     
-            igText(item.path.extractFilename().replace("C_", "C:/").replace("_", "/"))
+                component.contents[item.lootId] = ""       # Ensure that the sendGetLoot() function is sent only once by setting a value for the table key
-            
-            igDummy(vec2(0.0f, 5.0f))
-            igSeparator()
-            igDummy(vec2(0.0f, 5.0f)) 
 
-            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")

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) = 
+proc addTexture*(component: ScreenshotsComponent, lootId: string, data: string) = 
-    component.items.add(screenshot)
-
     var textureId: GLuint
-    let (width, height) = loadTextureFromBytes(string.toBytes(screenshot.data), textureId)
+    let (width, height) = loadTextureFromBytes(string.toBytes(data), textureId)
-    component.textures[screenshot.path] = ScreenshotTexture(
+    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
+                    # Download screenshot 
-                    connection.sendDownloadLoot(item.lootId)
+                    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.")

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)

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": %loot,
-            "loot": encode(file)
+            "data": encode(data)
         }
     )
     if client != nil: 

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: 
+            of CLIENT_LOOT_GET: 
-                let path = cq.dbGetLootById(event.data["lootId"].getStr()).path
+                let loot = cq.dbGetLootById(event.data["lootId"].getStr())
-                cq.client.sendLootSync(path, readFile(path))
+                cq.client.sendLootData(loot, readFile(loot.path))
 
             else: discard