Moved task parsing logic to the client to be able to support dotnet/bof commands when operating from a different machine than the team server. Disabled sequence tracking due to issues.
This commit is contained in:
Jakob Friedl
127
src/client/task.nim
Normal file
127
src/client/task.nim
Normal file
@@ -0,0 +1,127 @@
|
||||
import std/paths
|
||||
import strutils, sequtils, times, tables
|
||||
import ../common/[types, sequence, crypto, utils, serialize]
|
||||
|
||||
proc parseInput*(input: string): seq[string] =
|
||||
var i = 0
|
||||
while i < input.len:
|
||||
|
||||
# Skip whitespaces/tabs
|
||||
while i < input.len and input[i] in {' ', '\t'}:
|
||||
inc i
|
||||
if i >= input.len:
|
||||
break
|
||||
|
||||
var arg = ""
|
||||
if input[i] == '"':
|
||||
# Parse quoted argument
|
||||
inc i # Skip opening quote
|
||||
|
||||
# Add parsed argument when quotation is closed
|
||||
while i < input.len and input[i] != '"':
|
||||
arg.add(input[i])
|
||||
inc i
|
||||
|
||||
if i < input.len:
|
||||
inc i # Skip closing quote
|
||||
|
||||
else:
|
||||
while i < input.len and input[i] notin {' ', '\t'}:
|
||||
arg.add(input[i])
|
||||
inc i
|
||||
|
||||
# Add argument to returned result
|
||||
if arg.len > 0: result.add(arg)
|
||||
|
||||
proc parseArgument*(argument: Argument, value: string): TaskArg =
|
||||
|
||||
var arg: TaskArg
|
||||
arg.argType = cast[uint8](argument.argumentType)
|
||||
|
||||
case argument.argumentType:
|
||||
|
||||
of INT:
|
||||
# Length: 4 bytes
|
||||
let intValue = cast[uint32](parseUInt(value))
|
||||
arg.data = @[byte(intValue and 0xFF), byte((intValue shr 8) and 0xFF), byte((intValue shr 16) and 0xFF), byte((intValue shr 24) and 0xFF)]
|
||||
|
||||
of SHORT:
|
||||
# Length: 2 bytes
|
||||
let shortValue = cast[uint16](parseUint(value))
|
||||
arg.data = @[byte(shortValue and 0xFF), byte((shortValue shr 8) and 0xFF)]
|
||||
|
||||
of LONG:
|
||||
# Length: 8 bytes
|
||||
var data = newSeq[byte](8)
|
||||
let longValue = cast[uint64](parseUInt(value))
|
||||
for i in 0..7:
|
||||
data[i] = byte((longValue shr (i * 8)) and 0xFF)
|
||||
arg.data = data
|
||||
|
||||
of BOOL:
|
||||
# Length: 1 byte
|
||||
if value == "true":
|
||||
arg.data = @[1'u8]
|
||||
elif value == "false":
|
||||
arg.data = @[0'u8]
|
||||
else:
|
||||
raise newException(ValueError, "Invalid value for boolean argument.")
|
||||
|
||||
of STRING:
|
||||
arg.data = string.toBytes(value)
|
||||
|
||||
of BINARY:
|
||||
# A binary data argument consists of the file name (without the path) and the file content in bytes, both prefixed with their length as a uint32
|
||||
var packer = Packer.init()
|
||||
|
||||
let fileName = cast[string](extractFilename(cast[Path](value)))
|
||||
packer.addDataWithLengthPrefix(string.toBytes(fileName))
|
||||
|
||||
let fileContents = readFile(value)
|
||||
packer.addDataWithLengthPrefix(string.toBytes(fileContents))
|
||||
|
||||
arg.data = packer.pack()
|
||||
|
||||
return arg
|
||||
|
||||
proc createTask*(agentId, listenerId: string, command: Command, arguments: seq[string]): Task =
|
||||
|
||||
# Construct the task payload prefix
|
||||
var task: Task
|
||||
task.taskId = string.toUuid(generateUUID())
|
||||
task.listenerId = string.toUuid(listenerId)
|
||||
task.timestamp = uint32(now().toTime().toUnix())
|
||||
task.command = cast[uint16](command.commandType)
|
||||
task.argCount = uint8(arguments.len)
|
||||
|
||||
var taskArgs: seq[TaskArg]
|
||||
|
||||
# Add the task arguments
|
||||
if arguments.len() < command.arguments.filterIt(it.isRequired).len():
|
||||
raise newException(CatchableError, "Missing required argument.")
|
||||
|
||||
for i, arg in arguments:
|
||||
if i < command.arguments.len():
|
||||
taskArgs.add(parseArgument(command.arguments[i], arg))
|
||||
else:
|
||||
# Optional arguments should ALWAYS be placed at the end of the command and take the same definition
|
||||
taskArgs.add(parseArgument(command.arguments[^1], arg))
|
||||
|
||||
task.args = taskArgs
|
||||
|
||||
# Construct the header
|
||||
var taskHeader: Header
|
||||
taskHeader.magic = MAGIC
|
||||
taskHeader.version = VERSION
|
||||
taskHeader.packetType = cast[uint8](MSG_TASK)
|
||||
taskHeader.flags = cast[uint16](FLAG_ENCRYPTED)
|
||||
taskHeader.size = 0'u32
|
||||
taskHeader.agentId = string.toUuid(agentId)
|
||||
taskHeader.seqNr = nextSequence(taskHeader.agentId)
|
||||
taskHeader.iv = generateBytes(Iv) # Generate a random IV for AES-256 GCM
|
||||
taskHeader.gmac = default(AuthenticationTag)
|
||||
|
||||
task.header = taskHeader
|
||||
|
||||
# Return the task object for serialization
|
||||
return task
|
||||
@@ -1,12 +1,13 @@
|
||||
import whisky
|
||||
import strformat, strutils, times
|
||||
import strformat, strutils, times, json
|
||||
import imguin/[cimgui, glfw_opengl, simple]
|
||||
import ../utils/[appImGui, colors]
|
||||
import ../../common/[types]
|
||||
import ../websocket
|
||||
import ../../common/[types, utils]
|
||||
import ../../modules/manager
|
||||
import ../[task, websocket]
|
||||
|
||||
const MAX_INPUT_LENGTH = 512
|
||||
type
|
||||
type
|
||||
ConsoleComponent* = ref object of RootObj
|
||||
agent*: UIAgent
|
||||
showConsole*: bool
|
||||
@@ -124,6 +125,32 @@ proc addItem*(component: ConsoleComponent, itemType: LogType, data: string, time
|
||||
text: line
|
||||
))
|
||||
|
||||
#[
|
||||
Handling console commands
|
||||
]#
|
||||
proc handleAgentCommand*(component: ConsoleComponent, ws: WebSocket, input: string) =
|
||||
|
||||
# Convert user input into sequence of string arguments
|
||||
let parsedArgs = parseInput(input)
|
||||
|
||||
# Handle 'help' command
|
||||
if parsedArgs[0] == "help":
|
||||
# cq.handleHelp(parsedArgs)
|
||||
component.addItem(LOG_WARNING, "Help")
|
||||
return
|
||||
|
||||
# Handle commands with actions on the agent
|
||||
try:
|
||||
let
|
||||
command = getCommandByName(parsedArgs[0])
|
||||
task = createTask(component.agent.agentId, component.agent.listenerId, command, parsedArgs[1..^1])
|
||||
|
||||
ws.sendAgentTask(component.agent.agentId, task)
|
||||
component.addItem(LOG_INFO, fmt"Tasked agent to {command.description.toLowerAscii()} ({Uuid.toString(task.taskId)})")
|
||||
|
||||
except CatchableError:
|
||||
component.addItem(LOG_ERROR, getCurrentExceptionMsg())
|
||||
|
||||
#[
|
||||
Drawing
|
||||
]#
|
||||
@@ -271,7 +298,7 @@ proc draw*(component: ConsoleComponent, ws: WebSocket) =
|
||||
component.addItem(LOG_COMMAND, command)
|
||||
|
||||
# Send command to team server
|
||||
ws.sendAgentCommand(component.agent.agentId, command)
|
||||
component.handleAgentCommand(ws, command)
|
||||
|
||||
# Add command to console history
|
||||
component.history.add(command)
|
||||
|
||||
@@ -50,8 +50,8 @@ proc draw*(component: DockspaceComponent, showComponent: ptr bool, views: Table[
|
||||
igDockBuilderAddNode(dockspaceId, ImGuiDockNodeFlags_DockSpace.int32)
|
||||
igDockBuilderSetNodeSize(dockspaceId, vp.WorkSize)
|
||||
|
||||
discard igDockBuilderSplitNode(dockspaceId, ImGuiDir_Down, 0.8f, dockBottom, dockTop)
|
||||
discard igDockBuilderSplitNode(dockTop[], ImGuiDir_Right, 0.4f, dockTopRight, dockTopLeft)
|
||||
discard igDockBuilderSplitNode(dockspaceId, ImGuiDir_Down, 5.0f, dockBottom, dockTop)
|
||||
discard igDockBuilderSplitNode(dockTop[], ImGuiDir_Right, 0.5f, dockTopRight, dockTopLeft)
|
||||
|
||||
igDockBuilderDockWindow("Sessions [Table View]", dockTopLeft[])
|
||||
igDockBuilderDockWindow("Listeners", dockBottom[])
|
||||
|
||||
@@ -38,13 +38,13 @@ proc sendAgentBuild*(ws: WebSocket, buildInformation: AgentBuildInformation) =
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
|
||||
proc sendAgentCommand*(ws: WebSocket, agentId: string, command: string) =
|
||||
proc sendAgentTask*(ws: WebSocket, agentId: string, task: Task) =
|
||||
let event = Event(
|
||||
eventType: CLIENT_AGENT_COMMAND,
|
||||
eventType: CLIENT_AGENT_TASK,
|
||||
timestamp: now().toTime().toUnix(),
|
||||
data: %*{
|
||||
"agentId": agentId,
|
||||
"command": command
|
||||
"task": task
|
||||
}
|
||||
)
|
||||
ws.sendEvent(event)
|
||||
|
||||
Reference in New Issue
Block a user