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 ./utils/appImGui
import ./views/[dockspace, sessions, listeners, eventlog, console] import ./views/[dockspace, sessions, listeners, eventlog, console]
import ../common/[types, utils] import ../common/[types, utils]
import ./event/[send, recv] import ./websocket
import sugar import sugar

View File

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

View File

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

View File

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

View File

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

View File

@@ -8,7 +8,7 @@ import ../api/routes
import ../db/database import ../db/database
import ../core/logger import ../core/logger
import ../../common/[types, utils, profile] import ../../common/[types, utils, profile]
import ../event/send import ../websocket
#[ #[
Listener management 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.") raise newException(CatchableError, "Failed to store listener in database.")
cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{$port}.") 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}.") cq.client.sendEventlogItem(LOG_SUCCESS_SHORT, fmt"Started listener {name} on {host}:{$port}.")
except CatchableError as err: except CatchableError as err:
@@ -132,6 +133,7 @@ proc restartListeners*(cq: Conquest) =
cq.listeners[listener.listenerId] = listener cq.listeners[listener.listenerId] = listener
cq.threads[listener.listenerId] = thread 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}.") cq.success("Restarted listener", fgGreen, fmt" {listener.listenerId} ", resetStyle, fmt"on {listener.address}:{$listener.port}.")
except CatchableError as err: 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 strutils, strformat, system, tables
import ./[agent, listener, builder] import ./[agent, listener, builder]
import ../globals import ../globals
import ../db/database import ../db/database
import ../core/logger import ../core/logger
import ../../common/[types, crypto, utils, profile] import ../../common/[types, crypto, utils, profile, event]
import ../event/[recv, send] import ../websocket
import mummy, mummy/routers import mummy, mummy/routers
#[ #[
@@ -170,17 +170,23 @@ proc websocketHandler(ws: WebSocket, event: WebSocketEvent, message: Message) {.
# Continuously send heartbeat messages # Continuously send heartbeat messages
ws.sendHeartbeat() ws.sendHeartbeat()
# case message.getMessageType(): let event = message.recvEvent()
# of CLIENT_AGENT_COMMAND:
# discard case event.eventType:
# of CLIENT_LISTENER_START: of CLIENT_AGENT_COMMAND:
# message.receiveStartListener() discard
# of CLIENT_LISTENER_STOP:
# discard of CLIENT_LISTENER_START:
# # message.receiveStopListener() let listener = event.data.to(UIListener)
# of CLIENT_AGENT_BUILD: cq.listenerStart(listener.listenerId, listener.address, listener.port, listener.protocol)
# discard
# else: discard of CLIENT_LISTENER_STOP:
let listenerId = event.data["listenerId"].getStr()
cq.listenerStop(listenerId)
of CLIENT_AGENT_BUILD:
discard
else: discard
of ErrorEvent: of ErrorEvent:
discard discard
@@ -210,7 +216,7 @@ proc startServer*(profilePath: string) =
try: try:
# Initialize framework context # Initialize framework context
# Load and parse profile # Load and parse profile
let profile = parseFile(profilePath) let profile = parsetoml.parseFile(profilePath)
cq = Conquest.init(profile) cq = Conquest.init(profile)
cq.info("Using profile \"", profile.getString("name"), "\" (", profilePath ,").") cq.info("Using profile \"", profile.getString("name"), "\" (", profilePath ,").")

View File

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

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 mummy
import times, tables, json, base64, parsetoml import times, tables, json, base64, parsetoml
import ../utils import ./utils
import ../../common/[types, utils, serialize, event] import ../common/[types, utils, serialize, event]
export sendHeartbeat export sendHeartbeat, recvEvent
#[ #[
Server -> Client Server -> Client