Improved agent generation modal.

src/client/layout.ini

@@ -110,7 +110,7 @@ Size=76,76
 Collapsed=0
 
 [Window][Start Listener]
-Pos=704,411
+Pos=704,387
 Size=500,225
 Collapsed=0
 
@@ -119,13 +119,13 @@ IsChild=1
 Size=1363,540
 
 [Window][Generate Payload]
-Pos=1075,509
+Pos=704,161
-Size=717,677
+Size=500,677
 Collapsed=0
 
 [Window][Generate Payload/0_B6B17D5F]
 IsChild=1
-Size=326,310
+Size=217,310
 
 [Table][0x32886A44,8]
 Column 0  Weight=0.6432
@@ -153,5 +153,5 @@ DockSpace     ID=0x85940918 Window=0x260A4489 Pos=10,43 Size=1888,946 Split=Y
   DockNode    ID=0x00000005 Parent=0x85940918 SizeRef=1888,279 Split=X
     DockNode  ID=0x00000003 Parent=0x00000005 SizeRef=1310,159 CentralNode=1 Selected=0x61E02D75
     DockNode  ID=0x00000004 Parent=0x00000005 SizeRef=576,159 Selected=0x5E5F7166
-  DockNode    ID=0x00000006 Parent=0x85940918 SizeRef=1888,665 Selected=0x65D642C0
+  DockNode    ID=0x00000006 Parent=0x85940918 SizeRef=1888,665 Selected=0x6BE22050
 

src/client/views/listeners.nim

@@ -35,7 +35,7 @@ proc ListenersTable*(title: string): ListenersTableComponent =
     result.listeners = exampleListeners
     result.selection = ImGuiSelectionBasicStorage_ImGuiSelectionBasicStorage()
     result.startListenerModal = ListenerModal()
-    result.generatePayloadModal = AgentModal(result.listeners)
+    result.generatePayloadModal = AgentModal()
 
 
 proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebSocket) = 
@@ -57,7 +57,7 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS
         ws.sendStartListener(listener)
         component.listeners.add(listener)    
 
-    component.generatePayloadModal.draw()
+    component.generatePayloadModal.draw(component.listeners)
 
     #[
         Listener table

src/client/views/modals/generatePayload.nim

@@ -10,20 +10,16 @@ type
         sleepDelay: uint32 
         sleepMask: int32 
         spoofStack: bool 
-        listeners: seq[string]
         sleepMaskTechniques: seq[string]
         moduleSelection: DualListSelectionComponent[ModuleType]
 
-proc AgentModal*(listeners: seq[Listener]): AgentModalComponent =
+proc AgentModal*(): AgentModalComponent =
     result = new AgentModalComponent
     result.listener = 0
     result.sleepDelay = 5
     result.sleepMask = 0
     result.spoofStack = false
 
-    for l in listeners: 
-        result.listeners.add(l.listenerId)
-
     for technique in SleepObfuscationTechnique.low .. SleepObfuscationTechnique.high:
         result.sleepMaskTechniques.add($technique)
 
@@ -37,9 +33,14 @@ proc AgentModal*(listeners: seq[Listener]): AgentModalComponent =
     result.moduleSelection = DualListSelection(modules, moduleName)
 
 proc resetModalValues(component: AgentModalComponent) = 
-    discard 
+    component.listener = 0
+    component.sleepDelay = 5
+    component.sleepMask = 0
+    component.spoofStack = false 
+    component.moduleSelection.reset()
+
+proc draw*(component: AgentModalComponent, listeners: seq[Listener]) =
 
-proc draw*(component: AgentModalComponent) =
     let textSpacing = igGetStyle().ItemSpacing.x    
     
     # Center modal
@@ -64,7 +65,7 @@ proc draw*(component: AgentModalComponent) =
         igSameLine(0.0f, textSpacing)
         igGetContentRegionAvail(addr availableSize)
         igSetNextItemWidth(availableSize.x)
-        igCombo_Str("##InputListener", addr component.listener, (component.listeners.join("\0") & "\0").cstring , component.listeners.len().int32)
+        igCombo_Str("##InputListener", addr component.listener, (listeners.mapIt(it.listenerId).join("\0") & "\0").cstring , listeners.len().int32)
 
         # Sleep delay
         let step: uint32 = 1
@@ -110,7 +111,7 @@ proc draw*(component: AgentModalComponent) =
         if igButton("Build", vec2(availableSize.x * 0.5 - textSpacing * 0.5, 0.0f)):
 
             # Get values 
-            echo component.listeners[component.listener]
+            echo listeners[component.listener].listenerId
             echo $component.sleepDelay
             echo component.sleepMaskTechniques[component.sleepMask]
             echo $component.spoofStack

src/client/views/widgets/dualListSelection.nim

@@ -43,6 +43,9 @@ proc moveSelection[T](component: DualListSelectionComponent[T], src, dst: int) =
     ImGuiSelectionBasicStorage_Swap(component.selection[src], component.selection[dst])
     ImGuiSelectionBasicStorage_Clear(component.selection[src])
 
+proc reset*[T](component: DualListSelectionComponent[T]) = 
+    component.moveAll(1, 0)
+
 proc draw*[T](component: DualListSelectionComponent[T]) = 
 
     if igBeginTable("split", 3, ImGuiTableFlags_None.int32, vec2(0.0f, 0.0f), 0.0f): 

src/client/websocket.nim

@@ -69,9 +69,9 @@ proc sendAgentBuild*(ws: WebSocket, listenerId: string, sleepDelay: int, sleepMa
     [ Retrieval Functions ]
     Server -> Client 
 ]#
-proc getMessageType*(message: Message): WsMessageAction = 
+proc getMessageType*(message: Message): WsPacketType = 
     var unpacker = Unpacker.init(message.data)
-    return cast[WsMessageAction](unpacker.getUint8()) 
+    return cast[WsPacketType](unpacker.getUint8()) 
 
 proc receiveAgentPayload*(message: Message): seq[byte] = 
     var unpacker = Unpacker.init(message.data)

src/common/types.nim

@@ -280,13 +280,13 @@ type
     Client <-> Server WebSocket communication
 ]#
 type 
-    WsMessageAction* = enum
+    WsPacketType* = enum
         # Sent by client 
         CLIENT_HEARTBEAT = 0'u8             # Basic checkin 
-        CLIENT_AGENT_COMMAND = 1'u8               # Instruct TS to send queue a command for a specific agent
+        CLIENT_AGENT_BUILD = 1'u8           # Generate an agent binary for a specific listener
-        CLIENT_LISTENER_START = 2'u8        # Start a listener on the TS
+        CLIENT_AGENT_COMMAND = 2'u8               # Instruct TS to send queue a command for a specific agent
-        CLIENT_LISTENER_STOP = 3'u8         # Stop a listener
+        CLIENT_LISTENER_START = 3'u8        # Start a listener on the TS
-        CLIENT_AGENT_BUILD = 4'u8           # Generate an agent binary for a specific listener
+        CLIENT_LISTENER_STOP = 4'u8         # Stop a listener
 
         # Sent by team server
         CLIENT_AGENT_BINARY = 100'u8        # Return the agent binary to write to the operator's client machine
@@ -296,63 +296,3 @@ type
         CLIENT_EVENT_LOG = 104'u8           # Add entry to the eventlog        
         
         CLIENT_CONNECTION = 200'u8          # Return team server profile 
-
-    # Client -> Server 
-    WsHeartbeat* = object 
-        msgType* = CLIENT_HEARTBEAT         
-
-    WsCommand* = object 
-        msgType* = CLIENT_AGENT_COMMAND           
-        agentId*: uint32                    
-        command*: seq[byte]                 # Command input field in the console window, prefixed with length
-
-    WsListenerStart* = object  
-        msgType* = CLIENT_LISTENER_START    
-        listener*: Listener
-
-    WsListenerStop* = object 
-        msgType* = CLIENT_LISTENER_STOP     
-        listenerId*: uint32                 
-
-    WsAgentBuild* = object 
-        msgType* = CLIENT_AGENT_BUILD   
-        listenerId*: uint32                 
-        sleepDelay*: uint32                  
-        sleepMask*: SleepObfuscationTechnique
-        spoofStack*: uint8  
-        modules*: uint64
-
-    # Server -> Client 
-    WsAgentBinary* = object
-        msgType* = CLIENT_AGENT_BINARY 
-        agentPayload*: seq[byte]            # Agent binary in byte-form, opens file browser to select location on the client
-
-    WsAgentConnection* = object
-        msgType* = CLIENT_AGENT_CONNECTION
-        agent*: Agent                      
-    
-    WsAgentCheckin* = object 
-        msgType* = CLIENT_AGENT_CHECKIN
-        agentId*: uint32                   
-        timestamp*: uint32
-
-    WsConsoleLog* = object 
-        msgType* = CLIENT_CONSOLE_LOG
-        agentId*: uint32 
-        logType*: LogType
-        timestamp*: uint32 
-        data*: seq[byte]
-
-    WsEventLog* = object 
-        msgType* = CLIENT_EVENT_LOG
-        logType*: LogType
-        timestamp*: uint32 
-        data*: seq[byte]
-
-    WsClientConnection* = object 
-        msgType* = CLIENT_CONNECTION
-        version: uint8
-        profile*: seq[byte]
-        agents*: seq[Agent]
-        listeners*: seq[Listener] 
-    

src/server/websocket/receive.nim

@@ -9,9 +9,9 @@ import ../core/[task, listener]
     [ Retrieval functions ]
     Client -> Server
 ]#
-proc getMessageType*(message: Message): WsMessageAction = 
+proc getMessageType*(message: Message): WsPacketType = 
     var unpacker = Unpacker.init(message.data)
-    return cast[WsMessageAction](unpacker.getUint8()) 
+    return cast[WsPacketType](unpacker.getUint8()) 
 
 proc receiveStartListener*(message: Message) = 
     var unpacker = Unpacker.init(message.data)