diff --git a/src/client/main.nim b/src/client/main.nim index 76cd4f8..f116e67 100644 --- a/src/client/main.nim +++ b/src/client/main.nim @@ -73,119 +73,118 @@ proc main(ip: string = "localhost", port: int = 37573) = #[ WebSocket communication with the team server ]# - # Continuously send heartbeat messages - connection.ws.sendHeartbeat() - - # Receive and parse websocket response message try: - let event = recvEvent(connection.ws.receiveMessage().get(), connection.sessionKey) - case event.eventType: - of CLIENT_KEY_EXCHANGE: - connection.sessionKey = deriveSessionKey(clientKeyPair, decode(event.data["publicKey"].getStr()).toKey()) - connection.sendPublicKey(clientKeyPair.publicKey) - wipeKey(clientKeyPair.privateKey) + # Receive and parse websocket response message + let message = connection.ws.receiveMessage(timeout = 16) # Use a 16ms timeout to reduce CPU load = ~60FPS + if message.isSome(): + let event = recvEvent(message.get(), connection.sessionKey) + case event.eventType: + of CLIENT_KEY_EXCHANGE: + connection.sessionKey = deriveSessionKey(clientKeyPair, decode(event.data["publicKey"].getStr()).toKey()) + connection.sendPublicKey(clientKeyPair.publicKey) + wipeKey(clientKeyPair.privateKey) - of CLIENT_PROFILE: - profile = parsetoml.parseString(event.data["profile"].getStr()) - - of CLIENT_LISTENER_ADD: - let listener = event.data.to(UIListener) - listenersTable.listeners.add(listener) - - of CLIENT_AGENT_ADD: - let agent = event.data.to(UIAgent) - - # The ImGui Multi Select only works well with seq's, so we maintain a - # separate table of the latest agent heartbeats to have the benefit of quick and direct O(1) access - sessionsTable.agents.add(agent) - sessionsTable.agentActivity[agent.agentId] = agent.latestCheckin - - if not agent.impersonationToken.isEmptyOrWhitespace(): - sessionsTable.agentImpersonation[agent.agentId] = agent.impersonationToken - - # 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), - # so if you place your listeners somewhere else, the console windows show up somewhere else too - # The only case that is not covered is when the listeners table is hidden and the bottom panel was split - var agentConsole = Console(agent) - consoles[agent.agentId] = agentConsole - let listenersWindow = igFindWindowByName(WIDGET_LISTENERS) - if listenersWindow != nil and listenersWindow.DockNode != nil: - igSetNextWindowDockID(listenersWindow.DockNode.ID, ImGuiCond_FirstUseEver.int32) - else: - igSetNextWindowDockID(dockBottom, ImGuiCond_FirstUseEver.int32) - consoles[agent.agentId].draw(connection) - consoles[agent.agentId].showConsole = false - - of CLIENT_AGENT_CHECKIN: - sessionsTable.agentActivity[event.data["agentId"].getStr()] = event.timestamp - - of CLIENT_AGENT_PAYLOAD: - let payload = decode(event.data["payload"].getStr()) - try: - let path = callDialogFileSave("Save Payload") - writeFile(path, payload) - except IOError: - discard - - # Close and reset the payload generation modal window when the payload was received - listenersTable.generatePayloadModal.resetModalValues() - igClosePopupToLevel(0, false) - - of CLIENT_CONSOLE_ITEM: - let agentId = event.data["agentId"].getStr() - consoles[agentId].console.addItem( - cast[LogType](event.data["logType"].getInt()), - event.data["message"].getStr(), - event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") - ) - - of CLIENT_EVENTLOG_ITEM: - eventlog.textarea.addItem( - cast[LogType](event.data["logType"].getInt()), - event.data["message"].getStr(), - event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") - ) - - of CLIENT_BUILDLOG_ITEM: - listenersTable.generatePayloadModal.buildLog.addItem( - cast[LogType](event.data["logType"].getInt()), - event.data["message"].getStr(), - event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") - ) - - of CLIENT_LOOT_ADD: - let lootItem = event.data.to(LootItem) - case lootItem.itemType: - of DOWNLOAD: - lootDownloads.items.add(lootItem) - of SCREENSHOT: - lootScreenshots.items.add(lootItem) - else: discard - - of CLIENT_LOOT_DATA: - let - lootItem = event.data["loot"].to(LootItem) - data = decode(event.data["data"].getStr()) + of CLIENT_PROFILE: + profile = parsetoml.parseString(event.data["profile"].getStr()) - case lootItem.itemType: - of DOWNLOAD: - lootDownloads.contents[lootItem.lootId] = data - of SCREENSHOT: - lootScreenshots.addTexture(lootItem.lootId, data) + of CLIENT_LISTENER_ADD: + let listener = event.data.to(UIListener) + listenersTable.listeners.add(listener) + + of CLIENT_AGENT_ADD: + let agent = event.data.to(UIAgent) + + # The ImGui Multi Select only works well with seq's, so we maintain a + # separate table of the latest agent heartbeats to have the benefit of quick and direct O(1) access + sessionsTable.agents.add(agent) + sessionsTable.agentActivity[agent.agentId] = agent.latestCheckin + + if not agent.impersonationToken.isEmptyOrWhitespace(): + sessionsTable.agentImpersonation[agent.agentId] = agent.impersonationToken + + # 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), + # so if you place your listeners somewhere else, the console windows show up somewhere else too + # The only case that is not covered is when the listeners table is hidden and the bottom panel was split + var agentConsole = Console(agent) + consoles[agent.agentId] = agentConsole + let listenersWindow = igFindWindowByName(WIDGET_LISTENERS) + if listenersWindow != nil and listenersWindow.DockNode != nil: + igSetNextWindowDockID(listenersWindow.DockNode.ID, ImGuiCond_FirstUseEver.int32) + else: + igSetNextWindowDockID(dockBottom, ImGuiCond_FirstUseEver.int32) + consoles[agent.agentId].draw(connection) + consoles[agent.agentId].showConsole = false + + of CLIENT_AGENT_CHECKIN: + sessionsTable.agentActivity[event.data["agentId"].getStr()] = event.timestamp + + of CLIENT_AGENT_PAYLOAD: + let payload = decode(event.data["payload"].getStr()) + try: + let path = callDialogFileSave("Save Payload") + writeFile(path, payload) + except IOError: + discard + + # Close and reset the payload generation modal window when the payload was received + listenersTable.generatePayloadModal.resetModalValues() + igClosePopupToLevel(0, false) + + of CLIENT_CONSOLE_ITEM: + let agentId = event.data["agentId"].getStr() + consoles[agentId].console.addItem( + cast[LogType](event.data["logType"].getInt()), + event.data["message"].getStr(), + event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") + ) + + of CLIENT_EVENTLOG_ITEM: + eventlog.textarea.addItem( + cast[LogType](event.data["logType"].getInt()), + event.data["message"].getStr(), + event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") + ) + + of CLIENT_BUILDLOG_ITEM: + listenersTable.generatePayloadModal.buildLog.addItem( + cast[LogType](event.data["logType"].getInt()), + event.data["message"].getStr(), + event.timestamp.fromUnix().local().format("dd-MM-yyyy HH:mm:ss") + ) + + of CLIENT_LOOT_ADD: + let lootItem = event.data.to(LootItem) + case lootItem.itemType: + of DOWNLOAD: + lootDownloads.items.add(lootItem) + of SCREENSHOT: + lootScreenshots.items.add(lootItem) + else: discard + + of CLIENT_LOOT_DATA: + let + lootItem = event.data["loot"].to(LootItem) + data = decode(event.data["data"].getStr()) + + case lootItem.itemType: + of DOWNLOAD: + lootDownloads.contents[lootItem.lootId] = data + of SCREENSHOT: + lootScreenshots.addTexture(lootItem.lootId, data) + else: discard + + of CLIENT_IMPERSONATE_TOKEN: + let + agentId = event.data["agentId"].getStr() + impersonationToken = event.data["username"].getStr() + sessionsTable.agentImpersonation[agentId] = impersonationToken + + of CLIENT_REVERT_TOKEN: + sessionsTable.agentImpersonation.del(event.data["agentId"].getStr()) + else: discard - - of CLIENT_IMPERSONATE_TOKEN: - let - agentId = event.data["agentId"].getStr() - impersonationToken = event.data["username"].getStr() - sessionsTable.agentImpersonation[agentId] = impersonationToken - - of CLIENT_REVERT_TOKEN: - sessionsTable.agentImpersonation.del(event.data["agentId"].getStr()) - else: discard - # Draw/update UI components/views if showSessionsTable: sessionsTable.draw(addr showSessionsTable, connection) if showListeners: listenersTable.draw(addr showListeners, connection) diff --git a/src/server/main.nim b/src/server/main.nim index 65cab5f..cbae33f 100644 --- a/src/server/main.nim +++ b/src/server/main.nim @@ -44,10 +44,7 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {. # Send the public key for the key exchange, all other information with be transmitted when the key exchange is completed cq.client.sendPublicKey(cq.keyPair.publicKey) - of MessageEvent: - # Continuously send heartbeat messages - ws.sendHeartbeat() - + of MessageEvent: let event = message.recvEvent(cq.client.sessionKey) case event.eventType: