Implemented starting and stopping listeners from the ImGui client.
This commit is contained in:
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]),
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 ,").")
|
||||
|
||||
@@ -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
|
||||
);
|
||||
|
||||
""")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user