Added more websocket commands and started agent generation modal window.
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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())
|
||||
)
|
||||
42
src/server/websocket/receive.nim
Normal file
42
src/server/websocket/receive.nim
Normal 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()
|
||||
|
||||
79
src/server/websocket/send.nim
Normal file
79
src/server/websocket/send.nim
Normal 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)
|
||||
|
||||
Reference in New Issue
Block a user