Added more websocket commands and started agent generation modal window.

This commit is contained in:
Jakob Friedl
2025-09-23 15:51:57 +02:00
parent d3b37aa4a1
commit d4bdf56224
13 changed files with 326 additions and 112 deletions

View File

@@ -8,6 +8,7 @@ import ../api/routes
import ../db/database
import ../core/logger
import ../../common/[types, utils, profile]
import ../websocket/send
#[
Listener management
@@ -32,22 +33,16 @@ proc listenerList*(cq: Conquest) =
cq.drawTable(listeners)
proc serve(listener: Listener) {.thread.} =
try:
try:
listener.server.serve(Port(listener.port), listener.address)
except Exception:
except Exception as err:
discard
proc listenerStart*(cq: Conquest, host: string, portStr: string) =
proc listenerStart*(cq: Conquest, name: string, host: string, port: int, protocol: Protocol) =
# Validate arguments
try:
if not validatePort(portStr):
raise newException(CatchableError, fmt"[ - ] Invalid port number: {portStr}")
let port = portStr.parseInt
# Create new listener
let name: string = generateUUID()
var router: Router
router.notFoundHandler = routes.error404
router.methodNotAllowedHandler = routes.error405
@@ -78,7 +73,7 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
listenerId: name,
address: host,
port: port,
protocol: HTTP
protocol: protocol
)
# Start serving
@@ -90,10 +85,12 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
if not cq.dbStoreListener(listener):
raise newException(CatchableError, "Failed to store listener in database.")
cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{portStr}.")
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}.")
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}.")
proc restartListeners*(cq: Conquest) =
var listeners: seq[Listener] = cq.dbGetAllListeners()

View File

@@ -6,7 +6,7 @@ import ../globals
import ../db/database
import ../core/logger
import ../../common/[types, crypto, profile]
import ../protocol/websocket
import ../websocket/[receive, send]
import mummy, mummy/routers
#[
@@ -88,9 +88,11 @@ 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(opts.listener.get.start.get.ip, opts.listener.get.start.get.port)
discard
of "stop":
cq.listenerStop(opts.listener.get.stop.get.name)
#cq.listenerStop(opts.listener.get.stop.get.name)
discard
else:
cq.listenerUsage()
@@ -143,22 +145,34 @@ proc init*(T: type Conquest, profile: Profile): Conquest =
WebSocket
]#
proc upgradeHandler(request: Request) =
let ws = request.upgradeToWebSocket()
{.cast(gcsafe).}:
let ws = request.upgradeToWebSocket()
cq.ws = ws
# Send client connection message
ws.sendEventlogItem(LOG_SUCCESS_SHORT, "CQ-V1")
# Send client connection message
ws.sendEventlogItem(LOG_SUCCESS_SHORT, now().toTime().toUnix(), "CQ-V1")
proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.gcsafe.} =
{.cast(gcsafe).}:
case event:
of OpenEvent:
discard
of MessageEvent:
ws.sendHeartbeat()
proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) =
case event:
of OpenEvent:
discard
of MessageEvent:
ws.sendHeartbeat()
of ErrorEvent:
discard
of CloseEvent:
discard
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
of ErrorEvent:
discard
of CloseEvent:
discard
proc serve(server: Server) {.thread.} =
try:
@@ -196,6 +210,7 @@ proc startServer*(profilePath: string) =
cq.restartListeners()
cq.addMultiple(cq.dbGetAllAgents())
# Start websocket server
var router: Router
router.get("/*", upgradeHandler)
let server = newServer(router, websocketHandler)

View File

@@ -1,47 +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, timestamp: int64, message: string) =
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)
#[
[ Retrieval functions ]
Client -> Server
]#
proc getMessageType*(message: Message): WsMessageAction =
var unpacker = Unpacker.init(message.data)
return cast[WsMessageAction](unpacker.getUint8())
proc receiveStartListener*(message: Message): Listener =
var unpacker = Unpacker.init(message.data)
discard unpacker.getUint8()
return Listener(
server: nil,
listenerId: Uuid.toString(unpacker.getUint32()),
address: unpacker.getDataWithLengthPrefix(),
port: int(unpacker.getUint16()),
protocol: cast[Protocol](unpacker.getUint8())
)

View File

@@ -0,0 +1,42 @@
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): WsMessageAction =
var unpacker = Unpacker.init(message.data)
return cast[WsMessageAction](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

@@ -0,0 +1,79 @@
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)