Fixed team server crashing when no UI client is connected.

This commit is contained in:
Jakob Friedl
2025-09-26 15:30:14 +02:00
parent 19c073c261
commit a8154bb7d6
9 changed files with 207 additions and 126 deletions

View File

@@ -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
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002,0
[Window][Eventlog]
Pos=1489,43
Size=409,305
Pos=1398,43
Size=500,340
Collapsed=0
DockId=0x00000004,0
@@ -26,68 +26,68 @@ Size=1908,999
Collapsed=0
[Window][[7E248CBA] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002,0
[Window][[7BE69219] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
DockId=0x00000002,2
[Window][[CD44669B] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
[Window][[F300DB27] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
DockId=0x00000002,1
[Window][[AB3464CE] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
[Window][[3FC9903D] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
DockId=0x00000002,1
[Window][[6CD04F2C] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
[Window][[4E93619C] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
DockId=0x00000002,1
[Window][[0C56BE7D] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
DockId=0x00000002,1
[Window][[37C08F6C] jakob@AURA]
Pos=10,151
Size=1004,639
Pos=10,385
Size=1888,604
Collapsed=0
DockId=0x00000002
[Window][[BCC8B616] jakob@AURA]
Pos=10,151
Size=1004,639
Pos=10,186
Size=1004,604
Collapsed=0
DockId=0x00000002
@@ -97,21 +97,81 @@ Size=400,400
Collapsed=0
[Window][[9592878D] jakob@AURA]
Pos=10,350
Size=1888,639
Pos=10,385
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 1 Weight=1.0000
Column 2 Weight=1.0000
Column 3 Weight=1.0000
Column 4 Weight=1.0000
Column 5 Weight=1.0000
Column 6 Weight=1.0000
Column 7 Weight=1.0000
Column 8 Weight=1.0000
Column 0 Weight=0.9944 Visible=1
Column 1 Weight=1.2029 Visible=1
Column 2 Weight=0.7801 Visible=1
Column 3 Weight=0.6312 Visible=1
Column 4 Weight=1.3658 Visible=0
Column 5 Weight=1.0064 Visible=1
Column 6 Weight=1.1671 Visible=1
Column 7 Weight=0.8456 Visible=1
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=0x00000003 Parent=0x00000001 SizeRef=613,159 CentralNode=1 Selected=0x61E02D75
DockNode ID=0x00000004 Parent=0x00000001 SizeRef=409,159 Selected=0x0FA43D88
DockNode ID=0x00000002 Parent=0x85940918 SizeRef=1024,639 Selected=0x326CCD2B
DockNode ID=0x00000001 Parent=0x85940918 SizeRef=1024,340 Split=X
DockNode ID=0x00000003 Parent=0x00000001 SizeRef=1386,159 CentralNode=1 Selected=0x61E02D75
DockNode ID=0x00000004 Parent=0x00000001 SizeRef=500,159 Selected=0x0FA43D88
DockNode ID=0x00000002 Parent=0x85940918 SizeRef=1024,604 Selected=0x6BE22050

View File

@@ -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.add(agent)
sessionsTable.agents[agent.agentId] = 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

View File

@@ -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
# Create a new console window
if not component.consoles[].hasKey(agent.agentId):
component.consoles[][agent.agentId] = Console(agent)
if selectedAgent != nil:
# Create a new console window
if not component.consoles[].hasKey(selectedAgent.agentId):
component.consoles[][selectedAgent.agentId] = Console(selectedAgent)
# Focus the existing console window
else:
igSetWindowFocus_Str(fmt"[{agent.agentId}] {agent.username}@{agent.hostname}")
# Focus the existing console window
else:
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] = @[]
for i, agent in component.agents:
if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
newAgents.add(agent)
var agentsToDelete: seq[string] = @[]
for agentId, agent in component.agents.mpairs():
let i = igGetID_Str(agentId)
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()

View File

@@ -13,7 +13,7 @@ proc sendEvent*(ws: WebSocket, event: Event) =
packer.add(cast[uint32](event.timestamp))
packer.addDataWithLengthPrefix(string.toBytes($event.data))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc recvEvent*(message: Message): Event =

View File

@@ -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

View File

@@ -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.ws.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.")
cq.client.sendAgent(agent)
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

View File

@@ -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()

View File

@@ -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)
ws.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
cq.client.sendAgent(agent)
cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
of MessageEvent:
# Continuously send heartbeat messages
@@ -180,16 +183,17 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
# else: discard
of ErrorEvent:
discard
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:
server.serve(Port(12345))
except Exception:
discard
proc startServer*(profilePath: string) =
# Ensure that the conquest root directory was passed as a compile-time define
@@ -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()

View File

@@ -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)