Implemented handling of different argument types (int, wstring, short) for BOF files using specific prefixes.

src/agent/core/coff.nim

@@ -398,7 +398,8 @@ proc inlineExecute*(objectFile: seq[byte], args: seq[byte] = @[], entryFunction:
     echo "[*] Executing."
     if not objectExecute(addr objCtx, entryFunction, args): 
         raise newException(CatchableError, fmt"Failed to execute function {$entryFunction}.")
-
+    echo "[+] Object file executed successfully."
+    
     return true
 
 #[ 
@@ -411,7 +412,7 @@ proc inlineExecute*(objectFile: seq[byte], args: seq[byte] = @[], entryFunction:
 proc inlineExecuteGetOutput*(objectFile: seq[byte], args: seq[byte] = @[], entryFunction: string = "go"): string = 
 
     if not inlineExecute(objectFile, args, entryFunction): 
-        raise newException(CatchableError, fmt"[-] Failed to inline-execute object file.")
+        raise newException(CatchableError, fmt"[-] Failed to execute object file.")
 
     var output = BeaconGetOutputData(NULL)
     return $output
@@ -424,12 +425,54 @@ proc generateCoffArguments*(args: seq[TaskArg]): seq[byte] =
     
     var packer = Packer.init() 
     for arg in args: 
-        packer.add(uint32(arg.data.len()))
-        packer.addData(arg.data)
 
-        # Add terminating NULL byte to the end of string arguments
+        # All arguments passed to the beacon object file via the 'bof' command are handled as regular ANSI string
+        # As some BOFs however, take different argument types, prefixes can be used to indicate the exact data type
+        # [i]: INT
+        # [s]: SHORT
+        # [w]: WIDE STRING (utf-8)
+
         if arg.argType == uint8(types.STRING): 
-            packer.add(uint8('\0'))
+
+            try: 
+                let 
+                    prefix = Bytes.toString(arg.data)[0..3]
+                    value = Bytes.toString(arg.data)[4..^1]
+
+                # Check the first two characters for a type specification 
+                case prefix: 
+                of protect("[i]:"): 
+                    # Handle argument as integer
+                    let intValue: uint32 = cast[uint32](parseUint(value)) 
+                    packer.add(intValue)
+
+                of protect("[s]:"): 
+                    # Handle argument as short 
+                    let shortValue: uint16 = cast[uint16](parseUint(value))
+                    packer.add(shortValue)
+
+                of protect("[w]:"):
+                    # Handle argument as wide string  
+                    # Add terminating NULL byte to the end of string arguments
+                    let wStrData = cast[seq[byte]](+$value) # +$ converts a string to a wstring
+                    packer.add(uint32(wStrData.len()))
+                    packer.addData(wStrData)
+
+                else: 
+                    # In case no prefix is specified, handle the argument as a regular string
+                    raise newException(IndexDefect, "")
+
+            except IndexDefect: 
+                # Handle argument as regular string
+                # Add terminating NULL byte to the end of string arguments
+                let data = arg.data & @[uint8(0)]
+                packer.add(uint32(data.len()))
+                packer.addData(data)
+    
+        else: 
+            # Argument is not passed as a string, but instead directly as a int or short 
+            # Primarily for alias functions where the exact data types are defined 
+            packer.addData(arg.data)
 
     let argBytes = packer.pack() 
 

src/common/serialize.nim

@@ -125,6 +125,8 @@ proc getArgument*(unpacker: Unpacker): TaskArg =
         result.data = unpacker.getBytes(int(length))
     of INT:
         result.data = unpacker.getBytes(4)
+    of SHORT: 
+        result.data = unpacker.getBytes(2)
     of LONG:
         result.data = unpacker.getBytes(8)
     of BOOL:

src/common/types.nim

@@ -19,9 +19,10 @@ type
     ArgType* = enum 
         STRING = 0'u8
         INT = 1'u8
-        LONG = 2'u8
-        BOOL = 3'u8 
-        BINARY = 4'u8 
+        SHORT = 2'u8
+        LONG = 3'u8
+        BOOL = 4'u8 
+        BINARY = 5'u8 
 
     HeaderFlags* = enum 
         # Flags should be powers of 2 so they can be connected with or operators

src/modules/bof.nim

@@ -8,11 +8,11 @@ let commands*: seq[Command] =  @[
     Command(
         name: protect("bof"),
         commandType: CMD_BOF,
-        description: protect("Execute a object file in memory and retrieve the output."),
+        description: protect("Execute an object file in memory and retrieve the output."),
         example: protect("bof /path/to/dir.x64.o C:\\Users"),
         arguments: @[
-            Argument(name: protect("path"), description: protect("Local path to the object file to execute."), argumentType: BINARY, isRequired: true),
-            Argument(name: protect("arguments"), description: protect("Arguments to be passed to the object file."), argumentType: STRING, isRequired: false)
+            Argument(name: protect("path"), description: protect("Path to the object file to execute."), argumentType: BINARY, isRequired: true),
+            Argument(name: protect("arguments"), description: protect("Arguments to be passed to the object file. Arguments are handled as STRING, unless specified with a prefix ([i]:INT, [w]:WSTRING, [s]:SHORT; the colon separates prefix and value)"), argumentType: STRING, isRequired: false)
         ],
         execute: executeBof
     )
@@ -40,7 +40,7 @@ when defined(agent):
             of 1: # Only the object file has been passed as an argument
                 objectFile = task.args[0].data
                 arguments = @[]
-            else: # The optional 'arguments' parameter was included
+            else: # Parameters were passed to the BOF execution
                 objectFile = task.args[0].data
 
                 # Combine the passed arguments into a format that is understood by the Beacon API

src/server/protocol/parser.nim

@@ -44,12 +44,17 @@ proc parseArgument*(argument: Argument, value: string): TaskArg =
         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 intValue = cast[uint64](parseUInt(value))
+        let longValue = cast[uint64](parseUInt(value))
         for i in 0..7:
-            data[i] = byte((intValue shr (i * 8)) and 0xFF)
+            data[i] = byte((longValue shr (i * 8)) and 0xFF)
         arg.data = data
 
     of BOOL: