admin/conquest
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implemented wrapper functions for logging and console output (info, error, success, ...)

Browse Source
This commit is contained in:
Jakob Friedl
2025-08-21 17:02:50 +02:00
parent c9df7aba64
commit fbb08afe31
13 changed files with 164 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/common/types.nim
View File

@@ -62,6 +62,16 @@ type
CONFIG_PUBLIC_KEY = 4'u8
CONFIG_PROFILE = 5'u8
LogType* = enum
# LOG_INFO = "[ * ] "
# LOG_ERROR = "[ - ] "
# LOG_SUCCESS = "[ + ] "
# LOG_WARNING = "[ ! ] "
LOG_INFO = "[INFO] "
LOG_ERROR = "[FAIL] "
LOG_SUCCESS = "[DONE] "
LOG_WARNING = "[WARN] "
# Encryption
type
Uuid* = uint32

30
src/server/api/handlers.nim
View File

@@ -1,6 +1,6 @@
import terminal, strformat, strutils, sequtils, tables, times, system
import ../[utils, globals]
import ../globals
import ../db/database
import ../protocol/packer
import ../core/logger
@@ -19,23 +19,22 @@ proc register*(registrationData: seq[byte]): bool =
# Validate that listener exists
if not cq.dbListenerExists(agent.listenerId.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] {agent.ip} attempted to register to non-existent listener: {agent.listenerId}.", "\n")
cq.error(fmt"{agent.ip} attempted to register to non-existent listener: {agent.listenerId}.", "\n")
return false
# Store agent in database
if not cq.dbStoreAgent(agent):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Failed to insert agent {agent.agentId} into database.", "\n")
cq.error(fmt"Failed to insert agent {agent.agentId} into database.", "\n")
return false
# Create log directory
if not cq.makeAgentLogDirectory(agent.agentId):
cq.writeLine(fgRed, styleBright, "[ - ] Failed to create writeLine")
cq.error("Failed to create log directory.", "\n")
return false
cq.agents[agent.agentId] = agent
let date = agent.firstCheckin.format("dd-MM-yyyy HH:mm:ss")
cq.writeLine(fgYellow, styleBright, fmt"[{date}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n")
cq.info("Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n")
return true
@@ -54,12 +53,12 @@ proc getTasks*(heartbeat: seq[byte]): seq[seq[byte]] =
# Check if listener exists
if not cq.dbListenerExists(listenerId):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Task-retrieval request made to non-existent listener: {listenerId}.", "\n")
cq.error(fmt"Task-retrieval request made to non-existent listener: {listenerId}.", "\n")
raise newException(ValueError, "Invalid listener.")
# Check if agent exists
if not cq.dbAgentExists(agentId):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Task-retrieval request made to non-existent agent: {agentId}.", "\n")
cq.error(fmt"Task-retrieval request made to non-existent agent: {agentId}.", "\n")
raise newException(ValueError, "Invalid agent.")
# Update the last check-in date for the accessed agent
@@ -80,32 +79,31 @@ proc handleResult*(resultData: seq[byte]) =
taskResult = cq.deserializeTaskResult(resultData)
taskId = Uuid.toString(taskResult.taskId)
agentId = Uuid.toString(taskResult.header.agentId)
listenerId = Uuid.toString(taskResult.listenerId)
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, fmt"{$resultData.len} bytes received.")
cq.info(fmt"{$resultData.len} bytes received.")
case cast[StatusType](taskResult.status):
of STATUS_COMPLETED:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}]", fgGreen, " [ + ] ", resetStyle, fmt"Task {taskId} completed.")
cq.success(fmt"Task {taskId} completed.")
of STATUS_FAILED:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}]", fgRed, styleBright, " [ - ] ", resetStyle, fmt"Task {taskId} failed.")
cq.error(fmt"Task {taskId} failed.")
of STATUS_IN_PROGRESS:
discard
case cast[ResultType](taskResult.resultType):
of RESULT_STRING:
if int(taskResult.length) > 0:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, "Output:")
cq.info("Output:")
# Split result string on newline to keep formatting
for line in Bytes.toString(taskResult.data).split("\n"):
cq.writeLine(line)
cq.output(line)
of RESULT_BINARY:
# Write binary data to a file
cq.writeLine()
cq.output()
of RESULT_NO_OUTPUT:
cq.writeLine()
cq.output()
# Update task queue to include all tasks, except the one that was just completed
cq.agents[agentId].tasks = cq.agents[agentId].tasks.filterIt(it.taskId != taskResult.taskId)

7
src/server/api/routes.nim
View File

@@ -1,8 +1,9 @@
import prologue, terminal, strformat, parsetoml, tables
import strutils, times, base64
import strutils, base64
import ./handlers
import ../[utils, globals]
import ../globals
import ../core/logger
import ../../common/[types, utils, serialize, profile]
proc error404*(ctx: Context) {.async.} =
@@ -91,7 +92,7 @@ proc httpGet*(ctx: Context) {.async.} =
ctx.handled = true # Ensure that HTTP response is sent only once
# Notify operator that agent collected tasks
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, fmt"{$response.len} bytes sent.")
cq.info(fmt"{$response.len} bytes sent.")
except CatchableError:
resp "", Http404

27
src/server/core/agent.nim
View File

@@ -1,7 +1,8 @@
import terminal, strformat, strutils, tables, times, system, parsetoml
import terminal, strformat, strutils, tables, times, system, parsetoml, prompt
import ./task
import ../utils
import ../core/logger
import ../db/database
import ../../common/types
@@ -23,7 +24,7 @@ proc getAgentsAsSeq*(cq: Conquest): seq[Agent] =
Agent management
]#
proc agentUsage*(cq: Conquest) =
cq.writeLine("""Manage, build and interact with agents.
cq.output("""Manage, build and interact with agents.
Usage:
agent [options] COMMAND
@@ -49,7 +50,7 @@ proc agentList*(cq: Conquest, listener: string) =
else:
# Check if listener exists
if not cq.dbListenerExists(listener.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Listener {listener.toUpperAscii} does not exist.")
cq.error(fmt"Listener {listener.toUpperAscii} does not exist.")
return
cq.drawTable(cq.dbGetAllAgentsByListener(listener.toUpperAscii))
@@ -59,13 +60,13 @@ proc agentList*(cq: Conquest, listener: string) =
proc agentInfo*(cq: Conquest, name: string) =
# Check if agent supplied via -n parameter exists in database
if not cq.dbAgentExists(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Agent {name.toUpperAscii} does not exist.")
cq.error(fmt"Agent {name.toUpperAscii} does not exist.")
return
let agent = cq.agents[name.toUpperAscii]
# TODO: Improve formatting
cq.writeLine(fmt"""
cq.output(fmt"""
Agent name (UUID): {agent.agentId}
Connected to listener: {agent.listenerId}
──────────────────────────────────────────
@@ -87,7 +88,7 @@ proc agentKill*(cq: Conquest, name: string) =
# Check if agent supplied via -n parameter exists in database
if not cq.dbAgentExists(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Agent {name.toUpperAscii} does not exist.")
cq.error(fmt"Agent {name.toUpperAscii} does not exist.")
return
# TODO: Stop the process of the agent on the target system
@@ -96,18 +97,18 @@ proc agentKill*(cq: Conquest, name: string) =
# Remove the agent from the database
if not cq.dbDeleteAgentByName(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, "[ - ] Failed to terminate agent: ", getCurrentExceptionMsg())
cq.error("Failed to terminate agent: ", getCurrentExceptionMsg())
return
cq.delAgent(name)
cq.writeLine(fgYellow, styleBright, "[ + ] ", resetStyle, "Terminated agent ", fgYellow, styleBright, name.toUpperAscii, resetStyle, ".")
cq.success("Terminated agent ", fgYellow, styleBright, name.toUpperAscii, resetStyle, ".")
# Switch to interact mode
proc agentInteract*(cq: Conquest, name: string) =
# Verify that agent exists
if not cq.dbAgentExists(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Agent {name.toUpperAscii} does not exist.")
cq.error(fmt"Agent {name.toUpperAscii} does not exist.")
return
let agent = cq.agents[name.toUpperAscii]
@@ -115,13 +116,13 @@ proc agentInteract*(cq: Conquest, name: string) =
# Change prompt indicator to show agent interaction
cq.interactAgent = agent
cq.setIndicator(fmt"[{agent.agentId}]> ")
cq.setStatusBar(@[("[mode]", "interact"), ("[username]", fmt"{agent.username}"), ("[hostname]", fmt"{agent.hostname}"), ("[ip]", fmt"{agent.ip}"), ("[domain]", fmt"{agent.domain}")])
cq.prompt.setIndicator(fmt"[{agent.agentId}]> ")
cq.prompt.setStatusBar(@[("[mode]", "interact"), ("[username]", fmt"{agent.username}"), ("[hostname]", fmt"{agent.hostname}"), ("[ip]", fmt"{agent.ip}"), ("[domain]", fmt"{agent.domain}")])
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgYellow, "[ + ] ", resetStyle, fmt"Started interacting with agent ", fgYellow, styleBright, agent.agentId, resetStyle, ". Type 'help' to list available commands.\n")
cq.info("Started interacting with agent ", fgYellow, styleBright, agent.agentId, resetStyle, ". Type 'help' to list available commands.\n")
while command.replace(" ", "") != "back":
command = cq.readLine()
command = cq.prompt.readLine()
cq.withOutput(handleAgentCommand, command)
cq.interactAgent = nil

28
src/server/core/builder.nim
View File

@@ -1,8 +1,8 @@
import terminal, strformat, strutils, sequtils, tables, system, osproc, streams, parsetoml
import ../utils
import ../../common/[types, utils, profile, serialize, crypto]
import ../core/logger
import ../db/database
import ../../common/[types, utils, profile, serialize, crypto]
const PLACEHOLDER = "PLACEHOLDER"
@@ -37,7 +37,7 @@ proc serializeConfiguration(cq: Conquest, listener: Listener, sleep: int): seq[b
wipeKey(aesKey)
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, "Profile configuration serialized.")
cq.info("Profile configuration serialized.")
return encMaterial & encData
proc replaceAfterPrefix(content, prefix, value: string): string =
@@ -63,10 +63,10 @@ proc compile(cq: Conquest, placeholderLength: int): string =
.replaceAfterPrefix("-o:", exeFile)
writeFile(configFile, config)
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, fmt"Placeholder created ({placeholder.len()} bytes).")
cq.info(fmt"Placeholder created ({placeholder.len()} bytes).")
# Build agent by executing the ./build.sh script on the system.
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, "Compiling agent.")
cq.info("Compiling agent.")
try:
# Using the startProcess function from the 'osproc' module, it is possible to retrieve the output as it is received, line-by-line instead of all at once
@@ -75,25 +75,25 @@ proc compile(cq: Conquest, placeholderLength: int): string =
var line: string
while outputStream.readLine(line):
cq.writeLine(line)
cq.output(line)
let exitCode = process.waitForExit()
# Check if the build succeeded or not
if exitCode == 0:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgGreen, "[ + ] ", resetStyle, "Agent payload generated successfully.")
cq.info("Agent payload generated successfully.")
return exeFile
else:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgRed, "[ - ] ", resetStyle, "Build script exited with code ", $exitCode)
cq.error("Build script exited with code ", $exitCode)
return ""
except CatchableError as err:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgRed, "[ - ] ", resetStyle, "An error occurred: ", err.msg)
cq.error("An error occurred: ", err.msg)
return ""
proc patch(cq: Conquest, unpatchedExePath: string, configuration: seq[byte]): bool =
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, "Patching profile configuration into agent.")
cq.info("Patching profile configuration into agent.")
try:
var exeBytes = readFile(unpatchedExePath)
@@ -103,17 +103,17 @@ proc patch(cq: Conquest, unpatchedExePath: string, configuration: seq[byte]): bo
if placeholderPos == -1:
raise newException(CatchableError, "Placeholder not found.")
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ + ] ", resetStyle, fmt"Placeholder found at offset 0x{placeholderPos:08X}.")
cq.info(fmt"Placeholder found at offset 0x{placeholderPos:08X}.")
# Patch placeholder bytes
for i, c in Bytes.toString(configuration):
exeBytes[placeholderPos + i] = c
writeFile(unpatchedExePath, exeBytes)
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgGreen, "[ + ] ", resetStyle, fmt"Agent payload patched successfully: {unpatchedExePath}.")
cq.success(fmt"Agent payload patched successfully: {unpatchedExePath}.")
except CatchableError as err:
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgRed, styleBright, "[ - ] ", resetStyle, "An error occurred: ", err.msg)
cq.error("An error occurred: ", err.msg)
return false
return true
@@ -123,7 +123,7 @@ proc agentBuild*(cq: Conquest, listener, sleep: string): bool {.discardable.} =
# Verify that listener exists
if not cq.dbListenerExists(listener.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Listener {listener.toUpperAscii} does not exist.")
cq.error(fmt"Listener {listener.toUpperAscii} does not exist.")
return false
let listener = cq.listeners[listener.toUpperAscii]

19
src/server/core/listener.nim
View File

@@ -4,6 +4,7 @@ import prologue, parsetoml
import ../utils
import ../api/routes
import ../db/database
import ../core/logger
import ../../common/[types, utils, profile]
# Utility functions
@@ -17,7 +18,7 @@ proc add(cq: Conquest, listener: Listener) =
Listener management
]#
proc listenerUsage*(cq: Conquest) =
cq.writeLine("""Manage, start and stop listeners.
cq.output("""Manage, start and stop listeners.
Usage:
listener [options] COMMAND
@@ -88,10 +89,10 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
# Start serving
discard listener.runAsync()
cq.add(listenerInstance)
cq.writeLine(fgGreen, "[ + ] ", resetStyle, "Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{portStr}.")
cq.success("Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on {host}:{portStr}.")
except CatchableError as err:
cq.writeLine(fgRed, styleBright, "[ - ] Failed to start listener: ", err.msg)
cq.error("Failed to start listener: ", err.msg)
proc restartListeners*(cq: Conquest) =
let listeners: seq[Listener] = cq.dbGetAllListeners()
@@ -129,15 +130,15 @@ proc restartListeners*(cq: Conquest) =
discard listener.runAsync()
cq.add(l)
cq.writeLine(fgGreen, "[ + ] ", resetStyle, "Restarted listener", fgGreen, fmt" {l.listenerId} ", resetStyle, fmt"on {l.address}:{$l.port}.")
cq.success("Restarted listener", fgGreen, fmt" {l.listenerId} ", resetStyle, fmt"on {l.address}:{$l.port}.")
# Delay before serving another listener to avoid crashing the application
waitFor sleepAsync(100)
except CatchableError as err:
cq.writeLine(fgRed, styleBright, "[ - ] Failed to restart listener: ", err.msg)
cq.error("Failed to restart listener: ", err.msg)
cq.writeLine("")
cq.output()
# Remove listener from database, preventing automatic startup on server restart
@@ -145,14 +146,14 @@ proc listenerStop*(cq: Conquest, name: string) =
# Check if listener supplied via -n parameter exists in database
if not cq.dbListenerExists(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[ - ] Listener {name.toUpperAscii} does not exist.")
cq.error(fmt"Listener {name.toUpperAscii} does not exist.")
return
# Remove database entry
if not cq.dbDeleteListenerByName(name.toUpperAscii):
cq.writeLine(fgRed, styleBright, "[ - ] Failed to stop listener: ", getCurrentExceptionMsg())
cq.error("Failed to stop listener: ", getCurrentExceptionMsg())
return
cq.delListener(name)
cq.writeLine(fgGreen, "[ + ] ", resetStyle, "Stopped listener ", fgGreen, name.toUpperAscii, resetStyle, ".")
cq.success("Stopped listener ", fgGreen, name.toUpperAscii, resetStyle, ".")

35
src/server/core/logger.nim
View File

@@ -1,4 +1,4 @@
import times, strformat, strutils
import times, strformat, strutils, prompt, terminal
import std/[dirs, paths]
import ../../common/[types, profile]
@@ -30,4 +30,35 @@ proc extractStrings*(args: string): string =
for str in args[1..^2].split(", "):
if str.startsWith("\""):
message &= str
return message.replace("\"", "")
return message.replace("\"", "")
proc getTimestamp*(): string =
return now().format("dd-MM-yyyy HH:mm:ss")
# Function templates and overwrites
template writeLine*(cq: Conquest, args: varargs[untyped] = "") =
cq.prompt.writeLine(args)
if cq.interactAgent != nil:
cq.log(extractStrings($(args)))
# Wrapper functions for logging/console output
template info*(cq: Conquest, args: varargs[untyped] = "") =
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", $LOG_INFO, resetStyle, args)
template error*(cq: Conquest, args: varargs[untyped] = "") =
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgRed, $LOG_ERROR, resetStyle, args)
template success*(cq: Conquest, args: varargs[untyped] = "") =
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgGreen, $LOG_SUCCESS, resetStyle, args)
template warning*(cq: Conquest, args: varargs[untyped] = "") =
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] ", fgYellow, styleDim, $LOG_WARNING, resetStyle, args)
template input*(cq: Conquest, args: varargs[untyped] = "") =
if cq.interactAgent != nil:
cq.writeLine(fgBlue, styleBright, fmt"[{getTimestamp()}] ", fgYellow, fmt"[{cq.interactAgent.agentId}] ", resetStyle, args)
else:
cq.writeLine(fgBlue, styleBright, fmt"[{getTimestamp()}] ", resetStyle, args)
template output*(cq: Conquest, args: varargs[untyped] = "") =
cq.writeLine(args)

29
src/server/core/server.nim
View File

@@ -1,9 +1,10 @@
import prompt, terminal, argparse, parsetoml
import strutils, strformat, times, system, tables
import strutils, strformat, system, tables
import ./[agent, listener, builder]
import ../[globals, utils]
import ../db/database
import ../core/logger
import ../../common/[types, crypto, profile]
#[
@@ -67,7 +68,7 @@ proc handleConsoleCommand(cq: Conquest, args: string) =
# Return if no command (or just whitespace) is entered
if args.replace(" ", "").len == 0: return
cq.writeLine(fgBlue, styleBright, fmt"[{getTimestamp()}] ", resetStyle, styleBright, args)
cq.input(args)
try:
let opts = parser.parse(args.split(" ").filterIt(it.len > 0))
@@ -79,7 +80,7 @@ proc handleConsoleCommand(cq: Conquest, args: string) =
quit(0)
of "help": # Display help menu
cq.writeLine(parser.help())
cq.output(parser.help())
of "listener":
case opts.listener.get.command
@@ -110,13 +111,13 @@ proc handleConsoleCommand(cq: Conquest, args: string) =
# Handle help flag
except ShortCircuit as err:
if err.flag == "argparse_help":
cq.writeLine(err.help)
cq.error(err.help)
# Handle invalid arguments
except CatchableError:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
cq.writeLine("")
cq.output()
proc header() =
echo ""
@@ -147,13 +148,13 @@ proc startServer*(profilePath: string) =
header()
try:
# Initialize framework context
# Load and parse profile
let profile = parseFile(profilePath)
styledEcho(fgBlack, styleBright, "[ * ] ", "Using profile \"", profile.getString("name"), "\" (", profilePath ,").")
styledEcho(fgBlack, styleBright, "[ * ] ", "Using private key \"", profile.getString("private_key_file"), "\".")
# Initialize framework context
cq = Conquest.init(profile)
cq.info("Using profile \"", profile.getString("name"), "\" (", profilePath ,").")
cq.info("Using private key \"", profile.getString("private_key_file"), "\".")
except CatchableError as err:
echo err.msg
@@ -166,9 +167,9 @@ proc startServer*(profilePath: string) =
# Main loop
while true:
cq.setIndicator("[conquest]> ")
cq.setStatusBar(@[("[mode]", "manage"), ("[listeners]", $len(cq.listeners)), ("[agents]", $len(cq.agents))])
cq.showPrompt()
cq.prompt.setIndicator("[conquest]> ")
cq.prompt.setStatusBar(@[("[mode]", "manage"), ("[listeners]", $len(cq.listeners)), ("[agents]", $len(cq.agents))])
cq.prompt.showPrompt()
var command: string = cq.readLine()
var command: string = cq.prompt.readLine()
cq.withOutput(handleConsoleCommand, command)

32
src/server/core/task.nim
View File

@@ -1,16 +1,16 @@
import times, strformat, terminal, tables, sequtils, strutils
import strformat, terminal, tables, sequtils, strutils
import ../utils
import ../protocol/parser
import ../core/logger
import ../../modules/manager
import ../../common/types
proc displayHelp(cq: Conquest) =
cq.writeLine("Available commands:")
cq.writeLine(" * back")
cq.output("Available commands:")
cq.output(" * back")
for key, cmd in getAvailableCommands():
cq.writeLine(fmt" * {cmd.name:<15}{cmd.description}")
cq.writeLine()
cq.output(fmt" * {cmd.name:<15}{cmd.description}")
cq.output()
proc displayCommandHelp(cq: Conquest, command: Command) =
var usage = command.name & " " & command.arguments.mapIt(
@@ -20,24 +20,24 @@ proc displayCommandHelp(cq: Conquest, command: Command) =
if command.example != "":
usage &= "\nExample : " & command.example
cq.writeLine(fmt"""
cq.output(fmt"""
{command.description}
Usage : {usage}
""")
if command.arguments.len > 0:
cq.writeLine("Arguments:\n")
cq.output("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)}")
cq.output(fmt" {header[0]:<15} {header[1]:<6} {header[2]:<8} {header[3]}")
cq.output(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.output(fmt" * {arg.name:<15} {($arg.argumentType).toUpperAscii():<6} {isRequired:>8} {arg.description}")
cq.writeLine()
cq.output()
proc handleHelp(cq: Conquest, parsed: seq[string]) =
try:
@@ -48,13 +48,13 @@ proc handleHelp(cq: Conquest, parsed: seq[string]) =
cq.displayHelp()
except ValueError:
# Command was not found
cq.writeLine(fgRed, styleBright, fmt"[ - ] The command '{parsed[1]}' does not exist." & '\n')
cq.error("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
cq.writeLine(fgBlue, styleBright, fmt"[{getTimestamp()}] ", fgYellow, fmt"[{cq.interactAgent.agentId}] ", resetStyle, styleBright, input)
cq.input(input)
# Convert user input into sequence of string arguments
let parsedArgs = parseInput(input)
@@ -77,8 +77,8 @@ proc handleAgentCommand*(cq: Conquest, input: string) =
# Add task to queue
cq.interactAgent.tasks.add(task)
cq.writeLine(fgBlack, styleBright, fmt"[{getTimestamp()}] [ * ] ", resetStyle, fmt"Tasked agent to {command.description.toLowerAscii()}")
cq.info(fmt"Tasked agent to {command.description.toLowerAscii()}")
except CatchableError:
cq.writeLine(getCurrentExceptionMsg() & "\n")
cq.error(getCurrentExceptionMsg() & "\n")
return

8
src/server/db/database.nim
View File

@@ -1,7 +1,7 @@
import system, terminal, tiny_sqlite
import ./[dbAgent, dbListener]
import ../utils
import ../core/logger
import ../../common/types
# Export functions so that only ./db/database is required to be imported
@@ -41,7 +41,7 @@ proc dbInit*(cq: Conquest) =
""")
cq.writeLine(fgBlack, styleBright, "[ * ] Using new database: \"", cq.dbPath, "\".\n")
cq.info("Using new database: \"", cq.dbPath, "\".\n")
conquestDb.close()
except SqliteError as err:
cq.writeLine(fgBlack, styleBright, "[ * ] Using existing database: \"", cq.dbPath, "\".\n")
except SqliteError:
cq.info("Using existing database: \"", cq.dbPath, "\".\n")

20
src/server/db/dbAgent.nim
View File

@@ -1,6 +1,6 @@
import system, terminal, tiny_sqlite, times, sequtils
import ../utils
import ../core/logger
import ../../common/types
#[
@@ -21,7 +21,7 @@ proc dbStoreAgent*(cq: Conquest, agent: Agent): bool =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false
return true
@@ -41,7 +41,7 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
copyMem(sessionKey[0].addr, sessionKeyBlob[0].unsafeAddr, 32)
else:
# Handle invalid session key - log error but continue
cq.writeLine(fgYellow, styleBright, "[!] Invalid session key length for agent: ", agentId)
cq.warning("Invalid session key length for agent: ", agentId)
let a = Agent(
agentId: agentId,
@@ -65,7 +65,7 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return agents
@@ -83,7 +83,7 @@ proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
if sessionKeyBlob.len == 32:
copyMem(sessionKey[0].addr, sessionKeyBlob[0].unsafeAddr, 32)
else:
cq.writeLine(fgYellow, styleBright, "[!] Invalid session key length for agent: ", agentId)
cq.warning("Invalid session key length for agent: ", agentId)
let a = Agent(
agentId: agentId,
@@ -107,7 +107,7 @@ proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return agents
@@ -119,7 +119,7 @@ proc dbDeleteAgentByName*(cq: Conquest, name: string): bool =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false
return true
@@ -134,7 +134,7 @@ proc dbAgentExists*(cq: Conquest, agentName: string): bool =
return res.isSome
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false
proc dbUpdateCheckin*(cq: Conquest, agentName: string, timestamp: string): bool =
@@ -146,7 +146,7 @@ proc dbUpdateCheckin*(cq: Conquest, agentName: string, timestamp: string): bool
conquestDb.close()
return true
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false
proc dbUpdateSleep*(cq: Conquest, agentName: string, delay: int): bool =
@@ -158,5 +158,5 @@ proc dbUpdateSleep*(cq: Conquest, agentName: string, delay: int): bool =
conquestDb.close()
return true
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false

8
src/server/db/dbListener.nim
View File

@@ -1,6 +1,6 @@
import system, terminal, tiny_sqlite
import ../utils
import ../core/logger
import ../../common/types
# Utility functions
@@ -25,7 +25,7 @@ proc dbStoreListener*(cq: Conquest, listener: Listener): bool =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false
return true
@@ -50,7 +50,7 @@ proc dbGetAllListeners*(cq: Conquest): seq[Listener] =
conquestDb.close()
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return listeners
@@ -76,5 +76,5 @@ proc dbListenerExists*(cq: Conquest, listenerName: string): bool =
return res.isSome
except:
cq.writeLine(fgRed, styleBright, "[ - ] ", getCurrentExceptionMsg())
cq.error(getCurrentExceptionMsg())
return false

54
src/server/utils.nim
View File

@@ -11,33 +11,11 @@ proc validatePort*(portStr: string): bool =
except ValueError:
return false
proc getTimestamp*(): string =
return now().format("dd-MM-yyyy HH:mm:ss")
# Function templates and overwrites
template writeLine*(cq: Conquest, args: varargs[untyped] = "") =
cq.prompt.writeLine(args)
if cq.interactAgent != nil:
cq.log(extractStrings($(args)))
proc readLine*(cq: Conquest): string =
return cq.prompt.readLine()
proc setIndicator*(cq: Conquest, indicator: string) =
cq.prompt.setIndicator(indicator)
proc showPrompt*(cq: Conquest) =
cq.prompt.showPrompt()
proc hidePrompt*(cq: Conquest) =
cq.prompt.hidePrompt()
proc setStatusBar*(cq: Conquest, statusBar: seq[StatusBarItem]) =
cq.prompt.setStatusBar(statusBar)
proc clear*(cq: Conquest) =
cq.prompt.clear()
# Overwrite withOutput function to handle function arguments
proc withOutput*(cq: Conquest, outputFunction: proc(cq: Conquest, args: string), args: string) =
cq.hidePrompt()
cq.prompt.hidePrompt()
outputFunction(cq, args)
cq.showPrompt()
cq.prompt.showPrompt()
# Table border characters
type
@@ -102,12 +80,12 @@ proc drawTable*(cq: Conquest, listeners: seq[Listener]) =
let widths = @[8, 15, 5, 8, 6]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
cq.output(border(topLeft, topMid, topRight, widths))
for line in formatRow(headerCells, widths):
cq.hidePrompt()
cq.prompt.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(midLeft, midMid, midRight, widths))
cq.prompt.showPrompt()
cq.output(border(midLeft, midMid, midRight, widths))
for l in listeners:
# Get number of agents connected to the listener
@@ -122,11 +100,11 @@ proc drawTable*(cq: Conquest, listeners: seq[Listener]) =
]
for line in formatRow(rowCells, widths):
cq.hidePrompt()
cq.prompt.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.prompt.showPrompt()
cq.writeLine(border(botLeft, botMid, botRight, widths))
cq.output(border(botLeft, botMid, botRight, widths))
# Calculate time since latest checking in format: Xd Xh Xm Xs
proc timeSince*(agent: Agent, timestamp: DateTime): Cell =
@@ -165,12 +143,12 @@ proc drawTable*(cq: Conquest, agents: seq[Agent]) =
let widths = @[8, 15, 15, 15, 16, 13, 5, 8]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
cq.output(border(topLeft, topMid, topRight, widths))
for line in formatRow(headerCells, widths):
cq.hidePrompt()
cq.prompt.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(midLeft, midMid, midRight, widths))
cq.prompt.showPrompt()
cq.output(border(midLeft, midMid, midRight, widths))
for a in agents:
@@ -187,8 +165,8 @@ proc drawTable*(cq: Conquest, agents: seq[Agent]) =
# Highlight agents running within elevated processes
for line in formatRow(cells, widths):
cq.hidePrompt()
cq.prompt.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.prompt.showPrompt()
cq.writeLine(border(botLeft, botMid, botRight, widths))
cq.output(border(botLeft, botMid, botRight, widths))