Reworked websocket communication to avoid high CPU usage by client application.
This commit is contained in:
Jakob Friedl
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user