diff --git a/src/client/layout.ini b/src/client/layout.ini index bbc920d..4b7e7e7 100644 --- a/src/client/layout.ini +++ b/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 -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 diff --git a/src/client/main.nim b/src/client/main.nim index e276fc3..ff6f0c0 100644 --- a/src/client/main.nim +++ b/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.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 diff --git a/src/client/views/sessions.nim b/src/client/views/sessions.nim index cabf0ba..339b56b 100644 --- a/src/client/views/sessions.nim +++ b/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 - # 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() diff --git a/src/common/event.nim b/src/common/event.nim index 7b0a939..8629ca9 100644 --- a/src/common/event.nim +++ b/src/common/event.nim @@ -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 = diff --git a/src/common/types.nim b/src/common/types.nim index 11ba439..e55065c 100644 --- a/src/common/types.nim +++ b/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 - - diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 8213c7c..1282501 100644 --- a/src/server/api/handlers.nim +++ b/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.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 diff --git a/src/server/core/listener.nim b/src/server/core/listener.nim index a4d2751..06f7e50 100644 --- a/src/server/core/listener.nim +++ b/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() diff --git a/src/server/core/server.nim b/src/server/core/server.nim index 9cbee81..d94a2c8 100644 --- a/src/server/core/server.nim +++ b/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) - 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() diff --git a/src/server/event/send.nim b/src/server/event/send.nim index 311802d..d9558c5 100644 --- a/src/server/event/send.nim +++ b/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)