Started porting over functionality to the ImGui client via websocket communication.

This commit is contained in:
Jakob Friedl
2025-09-25 19:22:17 +02:00
parent f0dbcdfc58
commit 14771a4b50
22 changed files with 455 additions and 569 deletions

View File

@@ -4,6 +4,7 @@ import ../globals
import ../db/database
import ../protocol/packer
import ../core/logger
import ../event/send
import ../../common/[types, utils, serialize]
#[
@@ -36,6 +37,9 @@ proc register*(registrationData: seq[byte]): bool =
cq.agents[agent.agentId] = agent
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}.")
return true
@@ -69,6 +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)
# Return tasks
for task in cq.agents[agentId].tasks.mitems: # Iterate over agents as mutable items in order to modify GMAC tag

View File

@@ -141,7 +141,7 @@ proc agentBuild*(cq: Conquest, listener, sleepDelay: string, sleepTechnique: str
cq.error(fmt"Listener {listener.toUpperAscii} does not exist.")
return false
let listener = cq.listeners[listener.toUpperAscii].listener
let listener = cq.listeners[listener.toUpperAscii]
var config: seq[byte]
if sleepDelay.isEmptyOrWhitespace():

View File

@@ -8,7 +8,7 @@ import ../api/routes
import ../db/database
import ../core/logger
import ../../common/[types, utils, profile]
import ../websocket/send
import ../event/send
#[
Listener management
@@ -79,7 +79,9 @@ proc listenerStart*(cq: Conquest, name: string, host: string, port: int, protoco
createThread(thread, serve, listener)
server.waitUntilReady()
cq.listeners[name] = (listener, thread)
cq.listeners[name] = listener
cq.threads[name] = thread
if not cq.dbStoreListener(listener):
raise newException(CatchableError, "Failed to store listener in database.")
@@ -97,7 +99,6 @@ proc restartListeners*(cq: Conquest) =
for listener in listeners:
try:
# Create new listener
let name: string = generateUUID()
var router: Router
router.notFoundHandler = routes.error404
router.methodNotAllowedHandler = routes.error405
@@ -128,9 +129,11 @@ proc restartListeners*(cq: Conquest) =
createThread(thread, serve, listener)
server.waitUntilReady()
cq.listeners[listener.listenerId] = (listener, thread)
cq.listeners[listener.listenerId] = listener
cq.threads[listener.listenerId] = thread
cq.success("Restarted listener", fgGreen, fmt" {listener.listenerId} ", resetStyle, fmt"on {listener.address}:{$listener.port}.")
except CatchableError as err:
cq.error("Failed to restart listener: ", err.msg)

View File

@@ -5,8 +5,8 @@ import ./[agent, listener, builder]
import ../globals
import ../db/database
import ../core/logger
import ../../common/[types, crypto, profile]
import ../websocket/[receive, send]
import ../../common/[types, crypto, utils, profile]
import ../event/[recv, send]
import mummy, mummy/routers
#[
@@ -88,10 +88,10 @@ proc handleConsoleCommand(cq: Conquest, args: string) =
of "list":
cq.listenerList()
of "start":
#cq.listenerStart(opts.listener.get.start.get.ip, opts.listener.get.start.get.port)
cq.listenerStart(generateUUID(), opts.listener.get.start.get.ip, parseInt(opts.listener.get.start.get.port), HTTP)
discard
of "stop":
#cq.listenerStop(opts.listener.get.stop.get.name)
cq.listenerStop(opts.listener.get.stop.get.name)
discard
else:
cq.listenerUsage()
@@ -133,7 +133,8 @@ proc header() =
proc init*(T: type Conquest, profile: Profile): Conquest =
var cq = new Conquest
cq.prompt = Prompt.init()
cq.listeners = initTable[string, tuple[listener: Listener, thread: Thread[Listener]]]()
cq.listeners = initTable[string, Listener]()
cq.threads = initTable[string, Thread[Listener]]()
cq.agents = initTable[string, Agent]()
cq.interactAgent = nil
cq.profile = profile
@@ -148,27 +149,36 @@ proc upgradeHandler(request: Request) =
{.cast(gcsafe).}:
let ws = request.upgradeToWebSocket()
cq.ws = ws
# Send client connection message
ws.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.gcsafe.} =
{.cast(gcsafe).}:
case event:
of OpenEvent:
discard
# New client connected to team server
# Send profile, sessions and listeners to the UI client
ws.sendProfile(cq.profile)
for id, listener in cq.listeners:
ws.sendListener(listener)
for id, agent in cq.agents:
ws.sendAgent(agent)
ws.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
of MessageEvent:
# Continuously send heartbeat messages
ws.sendHeartbeat()
case message.getMessageType():
of CLIENT_AGENT_COMMAND:
discard
of CLIENT_LISTENER_START:
message.receiveStartListener()
of CLIENT_LISTENER_STOP:
message.receiveStopListener()
of CLIENT_AGENT_BUILD:
discard
else: discard
# case message.getMessageType():
# of CLIENT_AGENT_COMMAND:
# discard
# of CLIENT_LISTENER_START:
# message.receiveStartListener()
# of CLIENT_LISTENER_STOP:
# discard
# # message.receiveStopListener()
# of CLIENT_AGENT_BUILD:
# discard
# else: discard
of ErrorEvent:
discard
of CloseEvent:

40
src/server/event/recv.nim Normal file
View File

@@ -0,0 +1,40 @@
import mummy
import times, tables, json
import ./send
import ../globals
import ../core/[task, listener]
import ../../common/[types, utils, serialize, event]
#[
Client -> Server
]#
# proc getMessageType*(message: Message): EventType =
# var unpacker = Unpacker.init(message.data)
# return cast[EventType](unpacker.getUint8())
# proc receiveStartListener*(message: Message) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let
# listenerId = Uuid.toString(unpacker.getUint32())
# address = unpacker.getDataWithLengthPrefix()
# port = int(unpacker.getUint16())
# protocol = cast[Protocol](unpacker.getUint8())
# cq.listenerStart(listenerId, address, port, protocol)
# proc receiveStopListener*(message: Message) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let listenerId = Uuid.toString(unpacker.getUint32())
# cq.listenerStop(listenerId)
# proc receiveAgentCommand*(message: Message) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let
# agentId = Uuid.toString(unpacker.getUint32())
# command = unpacker.getDataWithLengthPrefix()

78
src/server/event/send.nim Normal file
View File

@@ -0,0 +1,78 @@
import mummy
import times, tables, json, base64, parsetoml
import ../utils
import ../../common/[types, utils, serialize, event]
export sendHeartbeat
#[
Server -> Client
]#
proc sendProfile*(ws: WebSocket, profile: Profile) =
let event = Event(
eventType: CLIENT_PROFILE,
timestamp: now().toTime().toUnix(),
data: %*{
"profile": profile.toTomlString()
}
)
ws.sendEvent(event)
proc sendEventlogItem*(ws: WebSocket, logType: LogType, message: string) =
let event = Event(
eventType: CLIENT_EVENTLOG_ITEM,
timestamp: now().toTime().toUnix(),
data: %*{
"logType": cast[uint8](logType),
"message": message
}
)
ws.sendEvent(event)
proc sendAgent*(ws: WebSocket, agent: Agent) =
let event = Event(
eventType: CLIENT_AGENT_ADD,
timestamp: now().toTime().toUnix(),
data: %agent
)
ws.sendEvent(event)
proc sendListener*(ws: WebSocket, listener: Listener) =
let event = Event(
eventType: CLIENT_LISTENER_ADD,
timestamp: now().toTime().toUnix(),
data: %listener
)
ws.sendEvent(event)
proc sendAgentCheckin*(ws: WebSocket, agentId: string) =
let event = Event(
eventType: CLIENT_AGENT_CHECKIN,
timestamp: now().toTime().toUnix(),
data: %*{
"agentId": agentId
}
)
ws.sendEvent(event)
proc sendAgentPayload*(ws: WebSocket, agentId: string, bytes: seq[byte]) =
let event = Event(
eventType: CLIENT_AGENT_PAYLOAD,
timestamp: now().toTime().toUnix(),
data: %*{
"agentId": agentId,
"payload": encode(bytes)
}
)
ws.sendEvent(event)
proc sendConsoleItem*(ws: WebSocket, agentId: string, logType: LogType, message: string) =
let event = Event(
eventType: CLIENT_CONSOLE_ITEM,
timestamp: now().toTime().toUnix(),
data: %*{
"agentId": agentId,
"logType": cast[uint8](logType),
"message": message
}
)
ws.sendEvent(event)

View File

@@ -1,15 +1,31 @@
import strutils, terminal, tables, sequtils, times, strformat, prompt
import strutils, terminal, tables, sequtils, times, strformat, prompt, json
import std/wordwrap
import ../common/types
import core/logger
proc validatePort*(portStr: string): bool =
try:
let port: int = portStr.parseInt
return port >= 1 and port <= 65535
except ValueError:
return false
proc `%`*(agent: Agent): JsonNode =
result = newJObject()
result["agentId"] = %agent.agentId
result["listenerId"] = %agent.listenerId
result["username"] = %agent.username
result["hostname"] = %agent.hostname
result["domain"] = %agent.domain
result["ip"] = %agent.ip
result["os"] = %agent.os
result["process"] = %agent.process
result["pid"] = %agent.pid
result["elevated"] = %agent.elevated
result["sleep"] = %agent.sleep
result["firstCheckin"] = %agent.firstCheckin.toTime().toUnix()
result["latestCheckin"] = %agent.latestCheckin.toTime().toUnix()
proc `%`*(listener: Listener): JsonNode =
result = newJObject()
result["listenerId"] = %listener.listenerId
result["address"] = %listener.address
result["port"] = %listener.port
result["protocol"] = %listener.protocol
# Table border characters
type

View File

@@ -1,42 +0,0 @@
import times, tables
import ../globals
import ../../common/[types, utils, serialize]
import mummy
import ./send
import ../core/[task, listener]
#[
[ Retrieval functions ]
Client -> Server
]#
proc getMessageType*(message: Message): WsPacketType =
var unpacker = Unpacker.init(message.data)
return cast[WsPacketType](unpacker.getUint8())
proc receiveStartListener*(message: Message) =
var unpacker = Unpacker.init(message.data)
discard unpacker.getUint8()
let
listenerId = Uuid.toString(unpacker.getUint32())
address = unpacker.getDataWithLengthPrefix()
port = int(unpacker.getUint16())
protocol = cast[Protocol](unpacker.getUint8())
cq.ws.sendEventlogItem(LOG_INFO_SHORT, "Attempting to start listener.")
cq.listenerStart(listenerId, address, port, protocol)
proc receiveStopListener*(message: Message) =
var unpacker = Unpacker.init(message.data)
discard unpacker.getUint8()
let listenerId = Uuid.toString(unpacker.getUint32())
cq.listenerStop(listenerId)
proc receiveAgentCommand*(message: Message) =
var unpacker = Unpacker.init(message.data)
discard unpacker.getUint8()
let
agentId = Uuid.toString(unpacker.getUint32())
command = unpacker.getDataWithLengthPrefix()

View File

@@ -1,79 +0,0 @@
import times, tables
import ../../common/[types, utils, serialize]
import mummy
#[
[ Sending functions ]
Server -> Client
]#
proc sendHeartbeat*(ws: WebSocket) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_HEARTBEAT))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc sendEventlogItem*(ws: WebSocket, logType: LogType, message: string, timestamp: int64 = now().toTime().toUnix()) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_EVENT_LOG))
packer.add(cast[uint8](logType))
packer.add(cast[uint32](timestamp))
packer.addDataWithLengthPrefix(string.toBytes(message))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc sendConsoleItem*(ws: WebSocket, agentId: string, logType: LogType, message: string, timestamp: int64 = now().toTime().toUnix()) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_CONSOLE_LOG))
packer.add(string.toUUid(agentId))
packer.add(cast[uint8](logType))
packer.add(cast[uint32](timestamp))
packer.addDataWithLengthPrefix(string.toBytes(message))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc sendAgentCheckin*(ws: WebSocket, agentId: string, timestamp: int64) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_AGENT_CHECKIN))
packer.add(string.toUUid(agentId))
packer.add(cast[uint32](timestamp))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc sendAgentPayload*(ws: WebSocket, payload: seq[byte]) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_AGENT_BINARY))
packer.addDataWithLengthPrefix(payload)
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
proc sendAgentConnection*(ws: WebSocket, agent: Agent) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_AGENT_CONNECTION))
packer.add(string.toUuid(agent.agentId))
packer.add(string.toUuid(agent.listenerId))
packer.addDataWithLengthPrefix(string.toBytes(agent.username))
packer.addDataWithLengthPrefix(string.toBytes(agent.hostname))
packer.addDataWithLengthPrefix(string.toBytes(agent.domain))
packer.addDataWithLengthPrefix(string.toBytes(agent.ip))
packer.addDataWithLengthPrefix(string.toBytes(agent.os))
packer.addDataWithLengthPrefix(string.toBytes(agent.process))
packer.add(uint32(agent.pid))
packer.add(uint8(agent.elevated))
packer.add(uint32(agent.sleep))
packer.add(cast[uint32](agent.firstCheckin))
packer.add(cast[uint32](agent.latestCheckin))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)