Implemented starting and stopping listeners from the ImGui client.

This commit is contained in:
Jakob Friedl
2025-09-27 13:54:12 +02:00
parent 933a72f920
commit ceba377939
11 changed files with 61 additions and 183 deletions

View File

@@ -1,77 +0,0 @@
import whisky
import times, tables
import ../views/[sessions, listeners, console, eventlog]
import ../../common/[types, utils, event]
export recvEvent
#[
Server -> Client
]#
# proc getMessageType*(message: Message): EventType =
# var unpacker = Unpacker.init(message.data)
# return cast[EventType](unpacker.getUint8())
# proc receiveAgentPayload*(message: Message): seq[byte] =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# return string.toBytes(unpacker.getDataWithLengthPrefix())
# proc receiveAgentConnection*(message: Message, sessions: ptr SessionsTableComponent) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let agent = Agent(
# agentId: Uuid.toString(unpacker.getUint32()),
# listenerId: Uuid.toString(unpacker.getUint32()),
# username: unpacker.getDataWithLengthPrefix(),
# hostname: unpacker.getDataWithLengthPrefix(),
# domain: unpacker.getDataWithLengthPrefix(),
# ip: unpacker.getDataWithLengthPrefix(),
# os: unpacker.getDataWithLengthPrefix(),
# process: unpacker.getDataWithLengthPrefix(),
# pid: int(unpacker.getUint32()),
# elevated: unpacker.getUint8() != 0,
# sleep: int(unpacker.getUint32()),
# tasks: @[],
# firstCheckin: cast[int64](unpacker.getUint32()).fromUnix().utc(),
# latestCheckin: cast[int64](unpacker.getUint32()).fromUnix().utc(),
# )
# sessions.agents.add(agent)
# proc receiveAgentCheckin*(message: Message, sessions: ptr SessionsTableComponent)=
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let agentId = Uuid.toString(unpacker.getUint32())
# let timestamp = cast[int64](unpacker.getUint32())
# # TODO: Update checkin
# proc receiveConsoleItem*(message: Message, consoles: ptr Table[string, ConsoleComponent]) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let
# agentId = Uuid.toString(unpacker.getUint32())
# logType = cast[LogType](unpacker.getUint8())
# timestamp = cast[int64](unpacker.getUint32())
# message = unpacker.getDataWithLengthPrefix()
# consoles[][agentId].addItem(logType, message, timestamp)
# proc receiveEventlogItem*(message: Message, eventlog: ptr EventlogComponent) =
# var unpacker = Unpacker.init(message.data)
# discard unpacker.getUint8()
# let
# logType = cast[LogType](unpacker.getUint8())
# timestamp = cast[int64](unpacker.getUint32())
# message = unpacker.getDataWithLengthPrefix()
# eventlog[].addItem(logType, message, timestamp)

View File

@@ -3,7 +3,7 @@ import tables, strutils, json, parsetoml
import ./utils/appImGui
import ./views/[dockspace, sessions, listeners, eventlog, console]
import ../common/[types, utils]
import ./event/[send, recv]
import ./websocket
import sugar

View File

@@ -3,7 +3,7 @@ import imguin/[cimgui, glfw_opengl, simple]
import ../utils/appImGui
import ../../common/[types, utils]
import ./modals/[startListener, generatePayload]
import ../event/send
import ../websocket
import whisky
type
@@ -76,15 +76,14 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS
var multiSelectIO = igBeginMultiSelect(ImGuiMultiSelectFlags_ClearOnEscape.int32 or ImGuiMultiSelectFlags_BoxSelect1d.int32, component.selection[].Size, int32(component.listeners.len()))
ImGuiSelectionBasicStorage_ApplyRequests(component.selection, multiSelectIO)
for row in 0 ..< component.listeners.len():
for i, listener in component.listeners:
igTableNextRow(ImGuiTableRowFlags_None.int32, 0.0f)
let listener = component.listeners[row]
if igTableSetColumnIndex(0):
# Enable multi-select functionality
igSetNextItemSelectionUserData(row)
var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](row))
igSetNextItemSelectionUserData(i)
var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i))
discard igSelectable_Bool(listener.listenerId, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32, vec2(0.0f, 0.0f))
if igTableSetColumnIndex(1):
@@ -101,11 +100,11 @@ proc draw*(component: ListenersTableComponent, showComponent: ptr bool, ws: WebS
if igMenuItem("Stop", nil, false, true):
# Update agents table with only non-selected ones
var newListeners: seq[UIListener] = @[]
for i in 0 ..< component.listeners.len():
for i, listener in component.listeners:
if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
newListeners.add(component.listeners[i])
newListeners.add(listener)
else:
ws.sendStopListener(component.listeners[i].listenerId)
ws.sendStopListener(listener.listenerId)
component.listeners = newListeners
ImGuiSelectionBasicStorage_Clear(component.selection)

View File

@@ -75,7 +75,6 @@ proc draw*(component: ListenerModalComponent): UIListener =
igBeginDisabled(($(addr component.address[0]) == "") or (component.port <= 0))
if igButton("Start", vec2(availableSize.x * 0.5 - textSpacing * 0.5, 0.0f)):
result = UIListener(
listenerId: generateUUID(),
address: $(addr component.address[0]),

View File

@@ -1,34 +1,29 @@
import whisky
import times, tables
import ../views/[sessions, listeners, console, eventlog]
import ../../common/[types, utils, serialize, event]
export sendHeartbeat
import times, tables, json
import ./views/[sessions, listeners, console, eventlog]
import ../common/[types, utils, serialize, event]
export sendHeartbeat, recvEvent
#[
Client -> Server
]#
proc sendStartListener*(ws: WebSocket, listener: UIListener) =
var packer = Packer.init()
packer.add(cast[uint8](CLIENT_LISTENER_START))
packer.add(string.toUUid(listener.listenerId))
packer.addDataWithLengthPrefix(string.toBytes(listener.address))
packer.add(cast[uint16](listener.port))
packer.add(cast[uint8](listener.protocol))
let data = packer.pack()
ws.send(Bytes.toString(data), BinaryMessage)
let event = Event(
eventType: CLIENT_LISTENER_START,
timestamp: now().toTime().toUnix(),
data: %listener
)
ws.sendEvent(event)
proc sendStopListener*(ws: WebSocket, listenerId: string) =
discard
# var packer = Packer.init()
# packer.add(cast[uint8](CLIENT_LISTENER_STOP))
# packer.add(string.toUuid(listenerId))
# let data = packer.pack()
# ws.send(Bytes.toString(data), BinaryMessage)
let event = Event(
eventType: CLIENT_LISTENER_STOP,
timestamp: now().toTime().toUnix(),
data: %*{
"listenerId": listenerId
}
)
ws.sendEvent(event)
# proc sendAgentCommand*(ws: WebSocket, agentId: string, command: string) =
# var packer = Packer.init()

View File

@@ -4,7 +4,7 @@ import ../globals
import ../db/database
import ../protocol/packer
import ../core/logger
import ../event/send
import ../websocket
import ../../common/[types, utils, serialize]
#[
@@ -21,26 +21,23 @@ proc register*(registrationData: seq[byte]): bool =
# Validate that listener exists
if not cq.dbListenerExists(agent.listenerId.toUpperAscii):
cq.error(fmt"{agent.ip} attempted to register to non-existent listener: {agent.listenerId}.", "\n")
return false
raise newException(CatchableError, fmt"{agent.ip} attempted to register to non-existent listener: {agent.listenerId}." & "\n")
# Store agent in database
if not cq.dbStoreAgent(agent):
cq.error(fmt"Failed to insert agent {agent.agentId} into database.", "\n")
return false
raise newException(CatchableError, fmt"Failed to insert agent {agent.agentId} into database." & "\n")
# Create log directory
if not cq.makeAgentLogDirectory(agent.agentId):
cq.error("Failed to create log directory.", "\n")
return false
raise newException(CatchableError, "Failed to create log directory.\n")
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")
# Send new agent to client
cq.client.sendAgent(agent)
cq.client.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.")
return true
except CatchableError as err:
@@ -63,13 +60,11 @@ proc getTasks*(heartbeat: seq[byte]): seq[seq[byte]] =
# Check if listener exists
if not cq.dbListenerExists(listenerId):
cq.error(fmt"Task-retrieval request made to non-existent listener: {listenerId}.", "\n")
raise newException(ValueError, "Invalid listener.")
raise newException(ValueError, fmt"Task-retrieval request made to non-existent listener: {listenerId}." & "\n")
# Check if agent exists
if not cq.dbAgentExists(agentId):
cq.error(fmt"Task-retrieval request made to non-existent agent: {agentId}.", "\n")
raise newException(ValueError, "Invalid agent.")
raise newException(ValueError, fmt"Task-retrieval request made to non-existent agent: {agentId}." & "\n")
# Update the last check-in date for the accessed agent
cq.agents[agentId].latestCheckin = cast[int64](timestamp).fromUnix().local()

View File

@@ -8,7 +8,7 @@ import ../api/routes
import ../db/database
import ../core/logger
import ../../common/[types, utils, profile]
import ../event/send
import ../websocket
#[
Listener management
@@ -86,6 +86,7 @@ proc listenerStart*(cq: Conquest, name: string, host: string, port: int, protoco
raise newException(CatchableError, "Failed to store listener in database.")
cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{$port}.")
cq.client.sendListener(listener)
cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Started listener {name} on {host}:{$port}.")
except CatchableError as err:
@@ -132,6 +133,7 @@ proc restartListeners*(cq: Conquest) =
cq.listeners[listener.listenerId] = listener
cq.threads[listener.listenerId] = thread
cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Restarted listener {listener.listenerId} on {listener.address}:{$listener.port}.")
cq.success("Restarted listener", fgGreen, fmt" {listener.listenerId} ", resetStyle, fmt"on {listener.address}:{$listener.port}.")
except CatchableError as err:

View File

@@ -1,12 +1,12 @@
import prompt, terminal, argparse, parsetoml, times
import prompt, terminal, argparse, parsetoml, times, json
import strutils, strformat, system, tables
import ./[agent, listener, builder]
import ../globals
import ../db/database
import ../core/logger
import ../../common/[types, crypto, utils, profile]
import ../event/[recv, send]
import ../../common/[types, crypto, utils, profile, event]
import ../websocket
import mummy, mummy/routers
#[
@@ -170,17 +170,23 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
# Continuously send heartbeat messages
ws.sendHeartbeat()
# 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
let event = message.recvEvent()
case event.eventType:
of CLIENT_AGENT_COMMAND:
discard
of CLIENT_LISTENER_START:
let listener = event.data.to(UIListener)
cq.listenerStart(listener.listenerId, listener.address, listener.port, listener.protocol)
of CLIENT_LISTENER_STOP:
let listenerId = event.data["listenerId"].getStr()
cq.listenerStop(listenerId)
of CLIENT_AGENT_BUILD:
discard
else: discard
of ErrorEvent:
discard
@@ -210,7 +216,7 @@ proc startServer*(profilePath: string) =
try:
# Initialize framework context
# Load and parse profile
let profile = parseFile(profilePath)
let profile = parsetoml.parseFile(profilePath)
cq = Conquest.init(profile)
cq.info("Using profile \"", profile.getString("name"), "\" (", profilePath ,").")

View File

@@ -35,8 +35,7 @@ proc dbInit*(cq: Conquest) =
sleep INTEGER DEFAULT 10,
firstCheckin DATETIME NOT NULL,
latestCheckin DATETIME NOT NULL,
sessionKey BLOB NOT NULL,
FOREIGN KEY (listener) REFERENCES listeners(name)
sessionKey BLOB NOT NULL
);
""")

View File

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

View File

@@ -1,8 +1,8 @@
import mummy
import times, tables, json, base64, parsetoml
import ../utils
import ../../common/[types, utils, serialize, event]
export sendHeartbeat
import ./utils
import ../common/[types, utils, serialize, event]
export sendHeartbeat, recvEvent
#[
Server -> Client