diff --git a/src/common/types.nim b/src/common/types.nim index 4854c0a..6825116 100644 --- a/src/common/types.nim +++ b/src/common/types.nim @@ -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 diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 393ce70..76cf189 100644 --- a/src/server/api/handlers.nim +++ b/src/server/api/handlers.nim @@ -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) \ No newline at end of file diff --git a/src/server/api/routes.nim b/src/server/api/routes.nim index be29014..bf9e97c 100644 --- a/src/server/api/routes.nim +++ b/src/server/api/routes.nim @@ -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 diff --git a/src/server/core/agent.nim b/src/server/core/agent.nim index 407b4a4..b7d8588 100644 --- a/src/server/core/agent.nim +++ b/src/server/core/agent.nim @@ -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 diff --git a/src/server/core/builder.nim b/src/server/core/builder.nim index 8ef5672..9d15cbb 100644 --- a/src/server/core/builder.nim +++ b/src/server/core/builder.nim @@ -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] diff --git a/src/server/core/listener.nim b/src/server/core/listener.nim index adb0ae9..9ef5d9c 100644 --- a/src/server/core/listener.nim +++ b/src/server/core/listener.nim @@ -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, ".") \ No newline at end of file diff --git a/src/server/core/logger.nim b/src/server/core/logger.nim index 0ca5532..0fc562a 100644 --- a/src/server/core/logger.nim +++ b/src/server/core/logger.nim @@ -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("\"", "") \ No newline at end of file + 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) \ No newline at end of file diff --git a/src/server/core/server.nim b/src/server/core/server.nim index 08ff3e1..716582c 100644 --- a/src/server/core/server.nim +++ b/src/server/core/server.nim @@ -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) diff --git a/src/server/core/task.nim b/src/server/core/task.nim index bbe7241..b718fb4 100644 --- a/src/server/core/task.nim +++ b/src/server/core/task.nim @@ -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 \ No newline at end of file diff --git a/src/server/db/database.nim b/src/server/db/database.nim index 492d978..da8b4db 100644 --- a/src/server/db/database.nim +++ b/src/server/db/database.nim @@ -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") diff --git a/src/server/db/dbAgent.nim b/src/server/db/dbAgent.nim index f888f5f..1e178d9 100644 --- a/src/server/db/dbAgent.nim +++ b/src/server/db/dbAgent.nim @@ -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 \ No newline at end of file diff --git a/src/server/db/dbListener.nim b/src/server/db/dbListener.nim index 58fc1f4..414a587 100644 --- a/src/server/db/dbListener.nim +++ b/src/server/db/dbListener.nim @@ -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 \ No newline at end of file diff --git a/src/server/utils.nim b/src/server/utils.nim index 1e499ea..7a35ced 100644 --- a/src/server/utils.nim +++ b/src/server/utils.nim @@ -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)) \ No newline at end of file + cq.output(border(botLeft, botMid, botRight, widths)) \ No newline at end of file