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

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