Implemented websocket (client <-> server) traffic encryption & compression.
This commit is contained in:
@@ -2,7 +2,7 @@ import whisky
|
||||
import tables, strutils, strformat, json, parsetoml, base64, os # native_dialogs
|
||||
import ./utils/[appImGui, globals]
|
||||
import ./views/[dockspace, sessions, listeners, eventlog, console]
|
||||
import ../common/[types, utils]
|
||||
import ../common/[types, utils, crypto]
|
||||
import ./websocket
|
||||
|
||||
import sugar
|
||||
@@ -39,9 +39,15 @@ proc main(ip: string = "localhost", port: int = 37573) =
|
||||
|
||||
let io = igGetIO()
|
||||
|
||||
# Create key pair
|
||||
let clientKeyPair = generateKeyPair()
|
||||
|
||||
# Initiate WebSocket connection
|
||||
let ws = newWebSocket(fmt"ws://{ip}:{$port}")
|
||||
defer: ws.close()
|
||||
var connection = WsConnection(
|
||||
ws: newWebSocket(fmt"ws://{ip}:{$port}"),
|
||||
sessionKey: default(Key)
|
||||
)
|
||||
defer: connection.ws.close()
|
||||
|
||||
# main loop
|
||||
while not app.handle.windowShouldClose:
|
||||
@@ -59,13 +65,17 @@ proc main(ip: string = "localhost", port: int = 37573) =
|
||||
WebSocket communication with the team server
|
||||
]#
|
||||
# Continuously send heartbeat messages
|
||||
ws.sendHeartbeat()
|
||||
connection.ws.sendHeartbeat()
|
||||
|
||||
# Receive and parse websocket response message
|
||||
let event = recvEvent(ws.receiveMessage().get())
|
||||
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)
|
||||
|
||||
of CLIENT_PROFILE:
|
||||
profile = parsetoml.parseString(event.data["profile"].getStr())
|
||||
profile = parsetoml.parseString(event.data["profile"].getStr())
|
||||
|
||||
of CLIENT_LISTENER_ADD:
|
||||
let listener = event.data.to(UIListener)
|
||||
@@ -90,7 +100,7 @@ proc main(ip: string = "localhost", port: int = 37573) =
|
||||
igSetNextWindowDockID(listenersWindow.DockNode.ID, ImGuiCond_FirstUseEver.int32)
|
||||
else:
|
||||
igSetNextWindowDockID(dockBottom, ImGuiCond_FirstUseEver.int32)
|
||||
consoles[agent.agentId].draw(ws)
|
||||
consoles[agent.agentId].draw(connection)
|
||||
consoles[agent.agentId].showConsole = false
|
||||
|
||||
of CLIENT_AGENT_CHECKIN:
|
||||
@@ -127,7 +137,7 @@ proc main(ip: string = "localhost", port: int = 37573) =
|
||||
|
||||
# Draw/update UI components/views
|
||||
if showSessionsTable: sessionsTable.draw(addr showSessionsTable)
|
||||
if showListeners: listenersTable.draw(addr showListeners, ws)
|
||||
if showListeners: listenersTable.draw(addr showListeners, connection)
|
||||
if showEventlog: eventlog.draw(addr showEventlog)
|
||||
|
||||
# Show console windows
|
||||
@@ -136,7 +146,7 @@ proc main(ip: string = "localhost", port: int = 37573) =
|
||||
if console.showConsole:
|
||||
# Ensure that new console windows are docked to the bottom panel by default
|
||||
igSetNextWindowDockID(dockBottom, ImGuiCond_FirstUseEver.int32)
|
||||
console.draw(ws)
|
||||
console.draw(connection)
|
||||
newConsoleTable[agentId] = console
|
||||
|
||||
# Update the consoles table with only those sessions that have not been closed yet
|
||||
|
||||
@@ -170,7 +170,7 @@ proc handleHelp(component: ConsoleComponent, parsed: seq[string]) =
|
||||
# Command was not found
|
||||
component.addItem(LOG_ERROR, fmt"The command '{parsed[1]}' does not exist.")
|
||||
|
||||
proc handleAgentCommand*(component: ConsoleComponent, ws: WebSocket, input: string) =
|
||||
proc handleAgentCommand*(component: ConsoleComponent, connection: WsConnection, input: string) =
|
||||
|
||||
# Convert user input into sequence of string arguments
|
||||
let parsedArgs = parseInput(input)
|
||||
@@ -186,7 +186,7 @@ proc handleAgentCommand*(component: ConsoleComponent, ws: WebSocket, input: stri
|
||||
command = getCommandByName(parsedArgs[0])
|
||||
task = createTask(component.agent.agentId, component.agent.listenerId, command, parsedArgs[1..^1])
|
||||
|
||||
ws.sendAgentTask(component.agent.agentId, task)
|
||||
connection.sendAgentTask(component.agent.agentId, task)
|
||||
component.addItem(LOG_INFO, fmt"Tasked agent to {command.description.toLowerAscii()} ({Uuid.toString(task.taskId)})")
|
||||
|
||||
except CatchableError:
|
||||
@@ -219,7 +219,7 @@ proc print(item: ConsoleItem) =
|
||||
igSameLine(0.0f, 0.0f)
|
||||
igTextUnformatted(item.text.cstring, nil)
|
||||
|
||||
proc draw*(component: ConsoleComponent, ws: WebSocket) =
|
||||
proc draw*(component: ConsoleComponent, connection: WsConnection) =
|
||||
igBegin(fmt"[{component.agent.agentId}] {component.agent.username}@{component.agent.hostname}".cstring, addr component.showConsole, 0)
|
||||
defer: igEnd()
|
||||
|
||||
@@ -340,7 +340,7 @@ proc draw*(component: ConsoleComponent, ws: WebSocket) =
|
||||
component.addItem(LOG_COMMAND, command)
|
||||
|
||||
# Send command to team server
|
||||
component.handleAgentCommand(ws, command)
|
||||
component.handleAgentCommand(connection, command)
|
||||
|
||||
# Add command to console history
|
||||
component.history.add(command)
|
||||
|
||||
@@ -22,7 +22,7 @@ proc ListenersTable*(title: string): ListenersTableComponent =
|
||||
result.startListenerModal = ListenerModal()
|
||||
result.generatePayloadModal = AgentModal()
|
||||
|
||||
proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebSocket) =
|
||||
proc draw*(component: ListenersTableComponent, showComponent: ptr bool, connection: WsConnection) =
|
||||
igBegin(component.title, showComponent, 0)
|
||||
defer: igEnd()
|
||||
|
||||
@@ -41,11 +41,11 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS
|
||||
|
||||
let listener = component.startListenerModal.draw()
|
||||
if listener != nil:
|
||||
ws.sendStartListener(listener)
|
||||
connection.sendStartListener(listener)
|
||||
|
||||
let buildInformation = component.generatePayloadModal.draw(component.listeners)
|
||||
if buildInformation != nil:
|
||||
ws.sendAgentBuild(buildInformation)
|
||||
connection.sendAgentBuild(buildInformation)
|
||||
|
||||
#[
|
||||
Listener table
|
||||
@@ -106,7 +106,7 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS
|
||||
if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
|
||||
newListeners.add(listener)
|
||||
else:
|
||||
ws.sendStopListener(listener.listenerId)
|
||||
connection.sendStopListener(listener.listenerId)
|
||||
|
||||
component.listeners = newListeners
|
||||
ImGuiSelectionBasicStorage_Clear(component.selection)
|
||||
|
||||
@@ -1,20 +1,30 @@
|
||||
import whisky
|
||||
import times, tables, json
|
||||
import times, tables, json, base64
|
||||
import ../common/[types, utils, serialize, event]
|
||||
export sendHeartbeat, recvEvent
|
||||
|
||||
#[
|
||||
Client -> Server
|
||||
]#
|
||||
proc sendStartListener*(ws: WebSocket, listener: UIListener) =
|
||||
proc sendPublicKey*(connection: WsConnection, publicKey: Key) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_KEY_EXCHANGE,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
data: %*{
|
||||
"publicKey": encode(Bytes.toString(publicKey))
|
||||
}
|
||||
)
|
||||
connection.ws.sendEvent(event, connection.sessionKey)
|
||||
|
||||
proc sendStartListener*(connection: WsConnection, listener: UIListener) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_LISTENER_START,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
data: %listener
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
connection.ws.sendEvent(event, connection.sessionKey)
|
||||
|
||||
proc sendStopListener*(ws: WebSocket, listenerId: string) =
|
||||
proc sendStopListener*(connection: WsConnection, listenerId: string) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_LISTENER_STOP,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
@@ -22,9 +32,9 @@ proc sendStopListener*(ws: WebSocket, listenerId: string) =
|
||||
"listenerId": listenerId
|
||||
}
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
connection.ws.sendEvent(event, connection.sessionKey)
|
||||
|
||||
proc sendAgentBuild*(ws: WebSocket, buildInformation: AgentBuildInformation) =
|
||||
proc sendAgentBuild*(connection: WsConnection, buildInformation: AgentBuildInformation) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_AGENT_BUILD,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
@@ -36,9 +46,9 @@ proc sendAgentBuild*(ws: WebSocket, buildInformation: AgentBuildInformation) =
|
||||
"modules": buildInformation.modules
|
||||
}
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
connection.ws.sendEvent(event, connection.sessionKey)
|
||||
|
||||
proc sendAgentTask*(ws: WebSocket, agentId: string, task: Task) =
|
||||
proc sendAgentTask*(connection: WsConnection, agentId: string, task: Task) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_AGENT_TASK,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
@@ -47,4 +57,4 @@ proc sendAgentTask*(ws: WebSocket, agentId: string, task: Task) =
|
||||
"task": task
|
||||
}
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
connection.ws.sendEvent(event, connection.sessionKey)
|
||||
|
||||
Reference in New Issue
Block a user