Rework module system. Now modules/commands are defined in a single file each, with both the function executed by teh agent and the definition for server-side argument parsing.
This commit is contained in:
Jakob Friedl
357
src/modules/filesystem.nim
Normal file
357
src/modules/filesystem.nim
Normal file
@@ -0,0 +1,357 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
when defined(server):
|
||||
proc executePwd(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeCd(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeDir(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeRm(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeRmdir(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeMove(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeCopy(config: AgentConfig, task: Task): TaskResult = nil
|
||||
|
||||
# Implementation of the execution functions
|
||||
when defined(agent):
|
||||
|
||||
import os, strutils, strformat, times, algorithm, winim
|
||||
import ../agent/core/taskresult
|
||||
|
||||
# Retrieve current working directory
|
||||
proc executePwd(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
echo fmt" [>] Retrieving current working directory."
|
||||
|
||||
try:
|
||||
# Get current working directory using GetCurrentDirectory
|
||||
let
|
||||
buffer = newWString(MAX_PATH + 1)
|
||||
length = GetCurrentDirectoryW(MAX_PATH, &buffer)
|
||||
|
||||
if length == 0:
|
||||
raise newException(OSError, fmt"Failed to get working directory ({GetLastError()}).")
|
||||
|
||||
let output = $buffer[0 ..< (int)length] & "\n"
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, output.toBytes())
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Change working directory
|
||||
proc executeCd(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let targetDirectory = task.args[0].data.toString()
|
||||
|
||||
echo fmt" [>] Changing current working directory to {targetDirectory}."
|
||||
|
||||
try:
|
||||
# Get current working directory using GetCurrentDirectory
|
||||
if SetCurrentDirectoryW(targetDirectory) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to change working directory ({GetLastError()}).")
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# List files and directories at a specific or at the current path
|
||||
proc executeDir(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
try:
|
||||
var targetDirectory: string
|
||||
|
||||
# Parse arguments
|
||||
case int(task.argCount):
|
||||
of 0:
|
||||
# Get current working directory using GetCurrentDirectory
|
||||
let
|
||||
cwdBuffer = newWString(MAX_PATH + 1)
|
||||
cwdLength = GetCurrentDirectoryW(MAX_PATH, &cwdBuffer)
|
||||
|
||||
if cwdLength == 0:
|
||||
raise newException(OSError, fmt"Failed to get working directory ({GetLastError()}).")
|
||||
|
||||
targetDirectory = $cwdBuffer[0 ..< (int)cwdLength]
|
||||
|
||||
of 1:
|
||||
targetDirectory = task.args[0].data.toString()
|
||||
else:
|
||||
discard
|
||||
|
||||
echo fmt" [>] Listing files and directories in {targetDirectory}."
|
||||
|
||||
# Prepare search pattern (target directory + \*)
|
||||
let searchPattern = targetDirectory & "\\*"
|
||||
let searchPatternW = newWString(searchPattern)
|
||||
|
||||
var
|
||||
findData: WIN32_FIND_DATAW
|
||||
hFind: HANDLE
|
||||
output = ""
|
||||
entries: seq[string] = @[]
|
||||
totalFiles = 0
|
||||
totalDirs = 0
|
||||
|
||||
# Find files and directories in target directory
|
||||
hFind = FindFirstFileW(searchPatternW, &findData)
|
||||
|
||||
if hFind == INVALID_HANDLE_VALUE:
|
||||
raise newException(OSError, fmt"Failed to find files ({GetLastError()}).")
|
||||
|
||||
# Directory was found and can be listed
|
||||
else:
|
||||
output = fmt"Directory: {targetDirectory}" & "\n\n"
|
||||
output &= "Mode LastWriteTime Length Name" & "\n"
|
||||
output &= "---- ------------- ------ ----" & "\n"
|
||||
|
||||
# Process all files and directories
|
||||
while true:
|
||||
let fileName = $cast[WideCString](addr findData.cFileName[0])
|
||||
|
||||
# Skip current and parent directory entries
|
||||
if fileName != "." and fileName != "..":
|
||||
# Get file attributes and size
|
||||
let isDir = (findData.dwFileAttributes and FILE_ATTRIBUTE_DIRECTORY) != 0
|
||||
let isHidden = (findData.dwFileAttributes and FILE_ATTRIBUTE_HIDDEN) != 0
|
||||
let isReadOnly = (findData.dwFileAttributes and FILE_ATTRIBUTE_READONLY) != 0
|
||||
let isArchive = (findData.dwFileAttributes and FILE_ATTRIBUTE_ARCHIVE) != 0
|
||||
let fileSize = (int64(findData.nFileSizeHigh) shl 32) or int64(findData.nFileSizeLow)
|
||||
|
||||
# Handle flags
|
||||
var mode = ""
|
||||
if isDir:
|
||||
mode = "d"
|
||||
inc totalDirs
|
||||
else:
|
||||
mode = "-"
|
||||
inc totalFiles
|
||||
|
||||
if isArchive:
|
||||
mode &= "a"
|
||||
else:
|
||||
mode &= "-"
|
||||
|
||||
if isReadOnly:
|
||||
mode &= "r"
|
||||
else:
|
||||
mode &= "-"
|
||||
|
||||
if isHidden:
|
||||
mode &= "h"
|
||||
else:
|
||||
mode &= "-"
|
||||
|
||||
if (findData.dwFileAttributes and FILE_ATTRIBUTE_SYSTEM) != 0:
|
||||
mode &= "s"
|
||||
else:
|
||||
mode &= "-"
|
||||
|
||||
# Convert FILETIME to local time and format
|
||||
var
|
||||
localTime: FILETIME
|
||||
systemTime: SYSTEMTIME
|
||||
dateTimeStr = "01/01/1970 00:00:00"
|
||||
|
||||
if FileTimeToLocalFileTime(&findData.ftLastWriteTime, &localTime) != 0 and FileTimeToSystemTime(&localTime, &systemTime) != 0:
|
||||
# Format date and time in PowerShell style
|
||||
dateTimeStr = fmt"{systemTime.wDay:02d}/{systemTime.wMonth:02d}/{systemTime.wYear} {systemTime.wHour:02d}:{systemTime.wMinute:02d}:{systemTime.wSecond:02d}"
|
||||
|
||||
# Format file size
|
||||
var sizeStr = ""
|
||||
if isDir:
|
||||
sizeStr = "<DIR>"
|
||||
else:
|
||||
sizeStr = ($fileSize).replace("-", "")
|
||||
|
||||
# Build the entry line
|
||||
let entryLine = fmt"{mode:<7} {dateTimeStr:<20} {sizeStr:>10} {fileName}"
|
||||
entries.add(entryLine)
|
||||
|
||||
# Find next file
|
||||
if FindNextFileW(hFind, &findData) == 0:
|
||||
break
|
||||
|
||||
# Close find handle
|
||||
discard FindClose(hFind)
|
||||
|
||||
# Add entries to output after sorting them (directories first, files afterwards)
|
||||
entries.sort do (a, b: string) -> int:
|
||||
let aIsDir = a[0] == 'd'
|
||||
let bIsDir = b[0] == 'd'
|
||||
|
||||
if aIsDir and not bIsDir:
|
||||
return -1
|
||||
elif not aIsDir and bIsDir:
|
||||
return 1
|
||||
else:
|
||||
# Extract filename for comparison (last part after the last space)
|
||||
let aParts = a.split(" ")
|
||||
let bParts = b.split(" ")
|
||||
let aName = aParts[^1]
|
||||
let bName = bParts[^1]
|
||||
return cmp(aName.toLowerAscii(), bName.toLowerAscii())
|
||||
|
||||
for entry in entries:
|
||||
output &= entry & "\n"
|
||||
|
||||
# Add summary of how many files/directories have been found
|
||||
output &= "\n" & fmt"{totalFiles} file(s)" & "\n"
|
||||
output &= fmt"{totalDirs} dir(s)" & "\n"
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, output.toBytes())
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Remove file
|
||||
proc executeRm(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let target = task.args[0].data.toString()
|
||||
|
||||
echo fmt" [>] Deleting file {target}."
|
||||
|
||||
try:
|
||||
if DeleteFile(target) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to delete file ({GetLastError()}).")
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Remove directory
|
||||
proc executeRmdir(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let target = task.args[0].data.toString()
|
||||
|
||||
echo fmt" [>] Deleting directory {target}."
|
||||
|
||||
try:
|
||||
if RemoveDirectoryA(target) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to delete directory ({GetLastError()}).")
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
# Move file or directory
|
||||
proc executeMove(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let
|
||||
lpExistingFileName = task.args[0].data.toString()
|
||||
lpNewFileName = task.args[1].data.toString()
|
||||
|
||||
echo fmt" [>] Moving {lpExistingFileName} to {lpNewFileName}."
|
||||
|
||||
try:
|
||||
if MoveFile(lpExistingFileName, lpNewFileName) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to move file or directory ({GetLastError()}).")
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Copy file or directory
|
||||
proc executeCopy(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let
|
||||
lpExistingFileName = task.args[0].data.toString()
|
||||
lpNewFileName = task.args[1].data.toString()
|
||||
|
||||
echo fmt" [>] Copying {lpExistingFileName} to {lpNewFileName}."
|
||||
|
||||
try:
|
||||
# Copy file to new location, overwrite if a file with the same name already exists
|
||||
if CopyFile(lpExistingFileName, lpNewFileName, FALSE) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to copy file or directory ({GetLastError()}).")
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Command definitions
|
||||
let commands* = @[
|
||||
Command(
|
||||
name: "pwd",
|
||||
commandType: CMD_PWD,
|
||||
description: "Retrieve current working directory.",
|
||||
example: "pwd",
|
||||
arguments: @[],
|
||||
execute: executePwd
|
||||
),
|
||||
Command(
|
||||
name: "cd",
|
||||
commandType: CMD_CD,
|
||||
description: "Change current working directory.",
|
||||
example: "cd C:\\Windows\\Tasks",
|
||||
arguments: @[
|
||||
Argument(name: "directory", description: "Relative or absolute path of the directory to change to.", argumentType: STRING, isRequired: true)
|
||||
],
|
||||
execute: executeCd
|
||||
),
|
||||
Command(
|
||||
name: "ls",
|
||||
commandType: CMD_LS,
|
||||
description: "List files and directories.",
|
||||
example: "ls C:\\Users\\Administrator\\Desktop",
|
||||
arguments: @[
|
||||
Argument(name: "directory", description: "Relative or absolute path. Default: current working directory.", argumentType: STRING, isRequired: false)
|
||||
],
|
||||
execute: executeDir
|
||||
),
|
||||
Command(
|
||||
name: "rm",
|
||||
commandType: CMD_RM,
|
||||
description: "Remove a file.",
|
||||
example: "rm C:\\Windows\\Tasks\\payload.exe",
|
||||
arguments: @[
|
||||
Argument(name: "file", description: "Relative or absolute path to the file to delete.", argumentType: STRING, isRequired: true)
|
||||
],
|
||||
execute: executeRm
|
||||
),
|
||||
Command(
|
||||
name: "rmdir",
|
||||
commandType: CMD_RMDIR,
|
||||
description: "Remove a directory.",
|
||||
example: "rm C:\\Payloads",
|
||||
arguments: @[
|
||||
Argument(name: "directory", description: "Relative or absolute path to the directory to delete.", argumentType: STRING, isRequired: true)
|
||||
],
|
||||
execute: executeRmdir
|
||||
),
|
||||
Command(
|
||||
name: "move",
|
||||
commandType: CMD_MOVE,
|
||||
description: "Move a file or directory.",
|
||||
example: "move source.exe C:\\Windows\\Tasks\\destination.exe",
|
||||
arguments: @[
|
||||
Argument(name: "source", description: "Source file path.", argumentType: STRING, isRequired: true),
|
||||
Argument(name: "destination", description: "Destination file path.", argumentType: STRING, isRequired: true)
|
||||
],
|
||||
execute: executeMove
|
||||
),
|
||||
Command(
|
||||
name: "copy",
|
||||
commandType: CMD_COPY,
|
||||
description: "Copy a file or directory.",
|
||||
example: "copy source.exe C:\\Windows\\Tasks\\destination.exe",
|
||||
arguments: @[
|
||||
Argument(name: "source", description: "Source file path.", argumentType: STRING, isRequired: true),
|
||||
Argument(name: "destination", description: "Destination file path.", argumentType: STRING, isRequired: true)
|
||||
],
|
||||
execute: executeCopy
|
||||
)
|
||||
]
|
||||
38
src/modules/manager.nim
Normal file
38
src/modules/manager.nim
Normal file
@@ -0,0 +1,38 @@
|
||||
import tables, strformat
|
||||
import ../common/[types, utils]
|
||||
|
||||
# Import modules
|
||||
import
|
||||
shell,
|
||||
sleep,
|
||||
filesystem
|
||||
|
||||
type
|
||||
ModuleManager* = object
|
||||
commandsByType*: Table[CommandType, Command]
|
||||
commandsByName*: Table[string, Command]
|
||||
|
||||
var manager: ModuleManager
|
||||
|
||||
proc registerCommands(commands: seq[Command]) {.discardable.} =
|
||||
for cmd in commands:
|
||||
manager.commandsByType[cmd.commandType] = cmd
|
||||
manager.commandsByName[cmd.name] = cmd
|
||||
|
||||
proc loadModules*() =
|
||||
# Register all imported commands
|
||||
registerCommands(shell.commands)
|
||||
registerCommands(sleep.commands)
|
||||
registerCommands(filesystem.commands)
|
||||
|
||||
proc getCommandByType*(cmdType: CommandType): Command =
|
||||
return manager.commandsByType[cmdType]
|
||||
|
||||
proc getCommandByName*(cmdName: string): Command =
|
||||
try:
|
||||
return manager.commandsByName[cmdName]
|
||||
except ValueError:
|
||||
raise newException(ValueError, fmt"The command '{cmdName}' does not exist.")
|
||||
|
||||
proc getAvailableCommands*(): Table[string, Command] =
|
||||
return manager.commandsByName
|
||||
55
src/modules/shell.nim
Normal file
55
src/modules/shell.nim
Normal file
@@ -0,0 +1,55 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
when defined(server):
|
||||
proc executeShell(config: AgentConfig, task: Task): TaskResult = nil
|
||||
|
||||
# Implement execution functions
|
||||
when defined(agent):
|
||||
|
||||
import ../agent/core/taskresult
|
||||
import osproc, strutils, strformat
|
||||
|
||||
proc executeShell(config: AgentConfig, task: Task): TaskResult =
|
||||
try:
|
||||
var
|
||||
command: string
|
||||
arguments: string
|
||||
|
||||
# Parse arguments
|
||||
case int(task.argCount):
|
||||
of 1: # Only the command has been passed as an argument
|
||||
command = task.args[0].data.toString()
|
||||
arguments = ""
|
||||
of 2: # The optional 'arguments' parameter was included
|
||||
command = task.args[0].data.toString()
|
||||
arguments = task.args[1].data.toString()
|
||||
else:
|
||||
discard
|
||||
|
||||
echo fmt" [>] Executing: {command} {arguments}."
|
||||
|
||||
let (output, status) = execCmdEx(fmt("{command} {arguments}"))
|
||||
|
||||
if output != "":
|
||||
return createTaskResult(task, cast[StatusType](status), RESULT_STRING, output.toBytes())
|
||||
else:
|
||||
return createTaskResult(task, cast[StatusType](status), RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
# Command definition (as seq[Command])
|
||||
let commands*: seq[Command] = @[
|
||||
Command(
|
||||
name: "shell",
|
||||
commandType: CMD_SHELL,
|
||||
description: "Execute a shell command and retrieve the output.",
|
||||
example: "shell whoami /all",
|
||||
arguments: @[
|
||||
Argument(name: "command", description: "Command to be executed.", argumentType: STRING, isRequired: true),
|
||||
Argument(name: "arguments", description: "Arguments to be passed to the command.", argumentType: STRING, isRequired: false)
|
||||
],
|
||||
execute: executeShell
|
||||
)
|
||||
]
|
||||
43
src/modules/sleep.nim
Normal file
43
src/modules/sleep.nim
Normal file
@@ -0,0 +1,43 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
when defined(server):
|
||||
proc executeSleep(config: AgentConfig, task: Task): TaskResult = nil
|
||||
|
||||
# Implement execution functions
|
||||
when defined(agent):
|
||||
|
||||
import os, strutils, strformat
|
||||
import ../agent/core/taskresult
|
||||
|
||||
proc executeSleep(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
try:
|
||||
# Parse task parameter
|
||||
let delay = int(task.args[0].data.toUint32())
|
||||
|
||||
echo fmt" [>] Sleeping for {delay} seconds."
|
||||
|
||||
sleep(delay * 1000)
|
||||
|
||||
# Updating sleep in agent config
|
||||
config.sleep = delay
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_NO_OUTPUT, @[])
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
|
||||
# Command definition (as seq[Command])
|
||||
let commands* = @[
|
||||
Command(
|
||||
name: "sleep",
|
||||
commandType: CMD_SLEEP,
|
||||
description: "Update sleep delay configuration.",
|
||||
example: "sleep 5",
|
||||
arguments: @[
|
||||
Argument(name: "delay", description: "Delay in seconds.", argumentType: INT, isRequired: true)
|
||||
],
|
||||
execute: executeSleep
|
||||
)
|
||||
]
|
||||
Reference in New Issue
Block a user