Fixed team server crashing when no UI client is connected.

src/client/layout.ini

@@ -1,18 +1,18 @@
 [Window][Sessions [Table View]]
 Pos=10,43
-Size=1477,305
+Size=1386,340
 Collapsed=0
 DockId=0x00000003,0
 
 [Window][Listeners]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002,0
 
 [Window][Eventlog]
-Pos=1489,43
+Pos=1398,43
-Size=409,305
+Size=500,340
 Collapsed=0
 DockId=0x00000004,0
 
@@ -26,68 +26,68 @@ Size=1908,999
 Collapsed=0
 
 [Window][[7E248CBA] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002,0
 
 [Window][[7BE69219] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
-DockId=0x00000002
+DockId=0x00000002,2
 
 [Window][[CD44669B] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002
 
 [Window][[F300DB27] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
-DockId=0x00000002
+DockId=0x00000002,1
 
 [Window][[AB3464CE] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002
 
 [Window][[3FC9903D] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
-DockId=0x00000002
+DockId=0x00000002,1
 
 [Window][[6CD04F2C] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002
 
 [Window][[4E93619C] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
-DockId=0x00000002
+DockId=0x00000002,1
 
 [Window][[0C56BE7D] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
 Collapsed=0
-DockId=0x00000002
+DockId=0x00000002,1
 
 [Window][[37C08F6C] jakob@AURA]
-Pos=10,151
+Pos=10,385
-Size=1004,639
+Size=1888,604
 Collapsed=0
 DockId=0x00000002
 
 [Window][[BCC8B616] jakob@AURA]
-Pos=10,151
+Pos=10,186
-Size=1004,639
+Size=1004,604
 Collapsed=0
 DockId=0x00000002
 
@@ -97,21 +97,81 @@ Size=400,400
 Collapsed=0
 
 [Window][[9592878D] jakob@AURA]
-Pos=10,350
+Pos=10,385
-Size=1888,639
+Size=1888,604
+Collapsed=0
+DockId=0x00000002,1
+
+[Window][[8F8DC95F] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[E05185F6] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002,1
+
+[Window][[022E62E0] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[F5CE46E3] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[37DAA990] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[824B6EC7] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[386BFA92] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[209155FC] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002,1
+
+[Window][[2498CB04] jakob@AURA]
+Pos=10,385
+Size=1888,604
+Collapsed=0
+DockId=0x00000002
+
+[Window][[0EBC9C82] jakob@AURA]
+Pos=10,186
+Size=1004,604
 Collapsed=0
 DockId=0x00000002,1
 
 [Table][0x32886A44,9]
-Column 0  Weight=1.0000
+Column 0  Weight=0.9944 Visible=1
-Column 1  Weight=1.0000
+Column 1  Weight=1.2029 Visible=1
-Column 2  Weight=1.0000
+Column 2  Weight=0.7801 Visible=1
-Column 3  Weight=1.0000
+Column 3  Weight=0.6312 Visible=1
-Column 4  Weight=1.0000
+Column 4  Weight=1.3658 Visible=0
-Column 5  Weight=1.0000
+Column 5  Weight=1.0064 Visible=1
-Column 6  Weight=1.0000
+Column 6  Weight=1.1671 Visible=1
-Column 7  Weight=1.0000
+Column 7  Weight=0.8456 Visible=1
-Column 8  Weight=1.0000
+Column 8  Weight=1.0064 Visible=1
 
 [Table][0x064A67CC,4]
 Column 0  Weight=1.0000
@@ -121,8 +181,8 @@ Column 3  Weight=1.0000
 
 [Docking][Data]
 DockSpace     ID=0x85940918 Window=0x260A4489 Pos=10,43 Size=1888,946 Split=Y
-  DockNode    ID=0x00000001 Parent=0x85940918 SizeRef=1024,159 Split=X
+  DockNode    ID=0x00000001 Parent=0x85940918 SizeRef=1024,340 Split=X
-    DockNode  ID=0x00000003 Parent=0x00000001 SizeRef=613,159 CentralNode=1 Selected=0x61E02D75
+    DockNode  ID=0x00000003 Parent=0x00000001 SizeRef=1386,159 CentralNode=1 Selected=0x61E02D75
-    DockNode  ID=0x00000004 Parent=0x00000001 SizeRef=409,159 Selected=0x0FA43D88
+    DockNode  ID=0x00000004 Parent=0x00000001 SizeRef=500,159 Selected=0x0FA43D88
-  DockNode    ID=0x00000002 Parent=0x85940918 SizeRef=1024,639 Selected=0x326CCD2B
+  DockNode    ID=0x00000002 Parent=0x85940918 SizeRef=1024,604 Selected=0x6BE22050
 

src/client/main.nim

@@ -69,13 +69,11 @@ proc main() =
         
         of CLIENT_LISTENER_ADD: 
             let listener = event.data.to(UIListener)
-            dump listener.listenerId
             listenersTable.listeners.add(listener)
 
         of CLIENT_AGENT_ADD: 
             let agent = event.data.to(UIAgent)
-            dump agent.agentId
+            sessionsTable.agents[agent.agentId] = agent
-            sessionsTable.agents.add(agent)
 
             # Initialize position of console windows to bottom by drawing them once when they are added
             # By default, the consoles are attached to the same DockNode as the Listeners table (Default: bottom), 
@@ -92,7 +90,7 @@ proc main() =
             consoles[agent.agentId].showConsole = false
 
         of CLIENT_AGENT_CHECKIN: 
-            discard 
+            sessionsTable.agents[event.data["agentId"].getStr()].latestCheckin = event.timestamp
 
         of CLIENT_AGENT_PAYLOAD: 
             discard

src/client/views/sessions.nim

@@ -8,14 +8,14 @@ import ../../common/[types, utils]
 type 
     SessionsTableComponent* = ref object of RootObj
         title: string 
-        agents*: seq[UIAgent]
+        agents*: Table[string, UIAgent]
         selection: ptr ImGuiSelectionBasicStorage
         consoles: ptr Table[string, ConsoleComponent]
 
 proc SessionsTable*(title: string, consoles: ptr Table[string, ConsoleComponent]): SessionsTableComponent = 
     result = new SessionsTableComponent
     result.title = title
-    result.agents = @[]
+    result.agents = initTable[string, UIAgent]() 
     result.selection = ImGuiSelectionBasicStorage_ImGuiSelectionBasicStorage()
     result.consoles = consoles
 
@@ -24,15 +24,19 @@ proc interact(component: SessionsTableComponent) =
     var it: pointer = nil
     var row: ImGuiID
     while ImGuiSelectionBasicStorage_GetNextSelectedItem(component.selection, addr it, addr row):
-        let agent = component.agents[cast[int](row)]
+        var selectedAgent: UIAgent = nil
+        for agentId, agent in component.agents.mpairs(): 
+            if igGetID_Str(agentId) == row:
+                selectedAgent = agent
 
+        if selectedAgent != nil: 
             # Create a new console window
-        if not component.consoles[].hasKey(agent.agentId):
+            if not component.consoles[].hasKey(selectedAgent.agentId):
-            component.consoles[][agent.agentId] = Console(agent)
+                component.consoles[][selectedAgent.agentId] = Console(selectedAgent)
 
             # Focus the existing console window
             else:
-            igSetWindowFocus_Str(fmt"[{agent.agentId}] {agent.username}@{agent.hostname}")
+                igSetWindowFocus_Str(fmt"[{selectedAgent.agentId}] {selectedAgent.username}@{selectedAgent.hostname}")
     
 proc draw*(component: SessionsTableComponent, showComponent: ptr bool) = 
     igBegin(component.title, showComponent, 0)
@@ -71,13 +75,15 @@ proc draw*(component: SessionsTableComponent, showComponent: ptr bool) =
         var multiSelectIO = igBeginMultiSelect(ImGuiMultiSelectFlags_ClearOnEscape.int32 or ImGuiMultiSelectFlags_BoxSelect1d.int32, component.selection[].Size, int32(component.agents.len())) 
         ImGuiSelectionBasicStorage_ApplyRequests(component.selection, multiSelectIO)
 
-        for row, agent in component.agents: 
+        for agentId, agent in component.agents.mpairs(): 
             
             igTableNextRow(ImGuiTableRowFlags_None.int32, 0.0f)
 
+            let row = igGetID_Str(agentId)
+
             if igTableSetColumnIndex(0):          
                 # Enable multi-select functionality       
-                igSetNextItemSelectionUserData(row)
+                igSetNextItemSelectionUserData(cast[ImGuiSelectionUserData](row))
                 var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](row))
                 discard igSelectable_Bool(agent.agentId, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32, vec2(0.0f, 0.0f))
                 
@@ -122,12 +128,15 @@ proc draw*(component: SessionsTableComponent, showComponent: ptr bool) =
         
             if igMenuItem("Remove", nil, false, true): 
                 # Update agents table with only non-selected ones
-                var newAgents: seq[UIAgent] = @[]
+                var agentsToDelete: seq[string] = @[]
-                for i, agent in component.agents:
+                for agentId, agent in component.agents.mpairs():
-                    if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
+                    let i = igGetID_Str(agentId)
-                        newAgents.add(agent)
+                    if ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
+                        agentsToDelete.add(agentId)
+
+                for agentId in agentsToDelete:
+                    component.agents.del(agentId)
 
-                component.agents = newAgents
                 ImGuiSelectionBasicStorage_Clear(component.selection)
                 igCloseCurrentPopup()
 

src/common/types.nim

@@ -3,6 +3,7 @@ import tables
 import times
 import parsetoml, json
 import mummy
+import system
 
 # Custom Binary Task structure
 const   
@@ -236,6 +237,33 @@ type
         port*: int
         protocol*: Protocol
 
+#[
+    Client <-> Server WebSocket communication
+]#
+type 
+    EventType* = enum
+        CLIENT_HEARTBEAT = 0'u8             # Basic checkin 
+       
+        # Sent by client 
+        CLIENT_AGENT_BUILD = 1'u8           # Generate an agent binary for a specific listener
+        CLIENT_AGENT_COMMAND = 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
+
+        # Sent by team server
+        CLIENT_PROFILE = 100'u8             # Team server profile and configuration 
+        CLIENT_LISTENER_ADD = 101'u8        # Add listener to listeners table
+        CLIENT_AGENT_ADD = 102'u8           # Add agent to sessions table
+        CLIENT_AGENT_CHECKIN = 103'u8       # Update agent checkin
+        CLIENT_AGENT_PAYLOAD = 104'u8       # Return agent payload binary 
+        CLIENT_CONSOLE_ITEM = 105'u8         # Add entry to a agent's console 
+        CLIENT_EVENTLOG_ITEM = 106'u8           # Add entry to the eventlog        
+
+    Event* = object 
+        eventType*: EventType               
+        timestamp*: int64 
+        data*: JsonNode   
+
 # Context structures
 type 
     KeyPair* = object 
@@ -244,6 +272,9 @@ type
 
     Profile* = TomlValueRef
 
+    UIClient* = ref object
+        ws*: WebSocket
+
     Conquest* = ref object
         prompt*: Prompt
         dbPath*: string
@@ -253,7 +284,7 @@ type
         interactAgent*: Agent
         keyPair*: KeyPair
         profile*: Profile
-        ws*: WebSocket
+        client*: UIClient
 
     AgentCtx* = ref object
         agentId*: string
@@ -298,32 +329,3 @@ type
 
     ConsoleItems* = ref object
         items*: seq[ConsoleItem]
-
-#[
-    Client <-> Server WebSocket communication
-]#
-type 
-    EventType* = enum
-        CLIENT_HEARTBEAT = 0'u8             # Basic checkin 
-       
-        # Sent by client 
-        CLIENT_AGENT_BUILD = 1'u8           # Generate an agent binary for a specific listener
-        CLIENT_AGENT_COMMAND = 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
-
-        # Sent by team server
-        CLIENT_PROFILE = 100'u8             # Team server profile and configuration 
-        CLIENT_LISTENER_ADD = 101'u8        # Add listener to listeners table
-        CLIENT_AGENT_ADD = 102'u8           # Add agent to sessions table
-        CLIENT_AGENT_CHECKIN = 103'u8       # Update agent checkin
-        CLIENT_AGENT_PAYLOAD = 104'u8       # Return agent payload binary 
-        CLIENT_CONSOLE_ITEM = 105'u8         # Add entry to a agent's console 
-        CLIENT_EVENTLOG_ITEM = 106'u8           # Add entry to the eventlog        
-
-    Event* = object 
-        eventType*: EventType               
-        timestamp*: int64 
-        data*: JsonNode   
-
-

src/server/api/handlers.nim

@@ -38,8 +38,8 @@ proc register*(registrationData: seq[byte]): bool =
 
             cq.info("Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") 
             
-            cq.ws.sendAgent(agent)
+            cq.client.sendAgent(agent)
-            cq.ws.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.")
+            cq.client.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.")
 
             return true
         
@@ -73,7 +73,7 @@ proc getTasks*(heartbeat: seq[byte]): seq[seq[byte]] =
 
             # Update the last check-in date for the accessed agent
             cq.agents[agentId].latestCheckin = cast[int64](timestamp).fromUnix().local()
-            # cq.ws.sendAgentCheckin(agentId)
+            cq.client.sendAgentCheckin(agentId)
 
             # Return tasks
             for task in cq.agents[agentId].tasks.mitems: # Iterate over agents as mutable items in order to modify GMAC tag

src/server/core/listener.nim

@@ -86,11 +86,11 @@ proc listenerStart*(cq: Conquest, name: string, host: string, port: int, protoco
             raise newException(CatchableError, "Failed to store listener in database.")
 
         cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{$port}.")
-        cq.ws.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Started listener {name} on {host}:{$port}.")
+        cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Started listener {name} on {host}:{$port}.")
 
     except CatchableError as err: 
         cq.error("Failed to start listener: ", err.msg)
-        cq.ws.sendEventlogItem(LOG_ERROR_SHORT, fmt"Failed to start listener: {err.msg}.")
+        cq.client.sendEventlogItem(LOG_ERROR_SHORT, fmt"Failed to start listener: {err.msg}.")
 
 proc restartListeners*(cq: Conquest) =
     var listeners: seq[Listener] = cq.dbGetAllListeners()

src/server/core/server.nim

@@ -140,6 +140,7 @@ proc init*(T: type Conquest, profile: Profile): Conquest =
     cq.profile = profile
     cq.keyPair = loadKeyPair(CONQUEST_ROOT & "/" & profile.getString("private-key-file"))
     cq.dbPath = CONQUEST_ROOT & "/" & profile.getString("database-file")
+    cq.client = nil
     return cq
 
 #[
@@ -148,7 +149,9 @@ proc init*(T: type Conquest, profile: Profile): Conquest =
 proc upgradeHandler(request: Request) = 
     {.cast(gcsafe).}:
         let ws = request.upgradeToWebSocket()
-        cq.ws = ws
+        cq.client = UIClient(
+            ws: ws
+        )
 
 proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.gcsafe.} = 
     {.cast(gcsafe).}:
@@ -156,12 +159,12 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
         of OpenEvent:
             # New client connected to team server
             # Send profile, sessions and listeners to the UI client
-            ws.sendProfile(cq.profile)
+            cq.client.sendProfile(cq.profile)
             for id, listener in cq.listeners: 
-                ws.sendListener(listener)
+                cq.client.sendListener(listener)
             for id, agent in cq.agents: 
-                ws.sendAgent(agent)
+                cq.client.sendAgent(agent)
-            ws.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
+            cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
     
         of MessageEvent:
             # Continuously send heartbeat messages
@@ -182,7 +185,8 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
         of ErrorEvent:
             discard 
         of CloseEvent:
-            discard
+            # Set the client instance to nil again to prevent debug error messages
+            cq.client = nil
     
 proc serve(server: Server) {.thread.} = 
     try:
@@ -230,6 +234,7 @@ proc startServer*(profilePath: string) =
 
     # Main loop
     while true: 
+    
         cq.prompt.setIndicator("[conquest]> ")
         cq.prompt.setStatusBar(@[("[mode]", "manage"), ("[listeners]", $len(cq.listeners)), ("[agents]", $len(cq.agents))])    
         cq.prompt.showPrompt() 

src/server/event/send.nim

@@ -7,7 +7,7 @@ export sendHeartbeat
 #[
     Server -> Client
 ]#
-proc sendProfile*(ws: WebSocket, profile: Profile) = 
+proc sendProfile*(client: UIClient, profile: Profile) = 
     let event = Event(
         eventType: CLIENT_PROFILE,
         timestamp: now().toTime().toUnix(),
@@ -15,9 +15,10 @@ proc sendProfile*(ws: WebSocket, profile: Profile) =
             "profile": profile.toTomlString()
         }
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendEventlogItem*(ws: WebSocket, logType: LogType, message: string) = 
+proc sendEventlogItem*(client: UIClient, logType: LogType, message: string) = 
     let event = Event(
         eventType: CLIENT_EVENTLOG_ITEM,
         timestamp: now().toTime().toUnix(),
@@ -26,25 +27,28 @@ proc sendEventlogItem*(ws: WebSocket, logType: LogType, message: string) =
             "message": message
         }
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendAgent*(ws: WebSocket, agent: Agent) = 
+proc sendAgent*(client: UIClient, agent: Agent) = 
     let event = Event(
         eventType: CLIENT_AGENT_ADD, 
         timestamp: now().toTime().toUnix(),
         data: %agent
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendListener*(ws: WebSocket, listener: Listener) =
+proc sendListener*(client: UIClient, listener: Listener) =
     let event = Event(
         eventType: CLIENT_LISTENER_ADD,
         timestamp: now().toTime().toUnix(),
         data: %listener
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendAgentCheckin*(ws: WebSocket, agentId: string) = 
+proc sendAgentCheckin*(client: UIClient, agentId: string) = 
     let event = Event(
         eventType: CLIENT_AGENT_CHECKIN,
         timestamp: now().toTime().toUnix(),
@@ -52,9 +56,10 @@ proc sendAgentCheckin*(ws: WebSocket, agentId: string) =
             "agentId": agentId
         }
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendAgentPayload*(ws: WebSocket, agentId: string, bytes: seq[byte]) =
+proc sendAgentPayload*(client: UIClient, agentId: string, bytes: seq[byte]) =
     let event = Event(
         eventType: CLIENT_AGENT_PAYLOAD, 
         timestamp: now().toTime().toUnix(),
@@ -63,9 +68,10 @@ proc sendAgentPayload*(ws: WebSocket, agentId: string, bytes: seq[byte]) =
             "payload": encode(bytes)
         }
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)
 
-proc sendConsoleItem*(ws: WebSocket, agentId: string, logType: LogType, message: string) = 
+proc sendConsoleItem*(client: UIClient, agentId: string, logType: LogType, message: string) = 
     let event = Event(
         eventType: CLIENT_CONSOLE_ITEM,
         timestamp: now().toTime().toUnix(),
@@ -75,4 +81,5 @@ proc sendConsoleItem*(ws: WebSocket, agentId: string, logType: LogType, message:
             "message": message
         }
     )
-    ws.sendEvent(event)
+    if client != nil: 
+        client.ws.sendEvent(event)