Implemented agent removal, changed object structures

server/agent/agent.nim

@@ -1,4 +1,4 @@
-import terminal, strformat
+import terminal, strformat, strutils
 import ../[types, globals, utils]
 import ../db/database
 
@@ -16,29 +16,45 @@ Usage:
 Commands:
 
   list             List all agents.
-  build            Build an agent to connect to an active listener.
+  kill             Terminate the connection of an active listener and remove it from the interface.
   interact         Interact with an active agent.
 
 Options:
   -h, --help""")
 
+# List agents
 proc agentList*(cq: Conquest, args: varargs[string]) = 
     let agents = cq.dbGetAllAgents()
     cq.drawTable(agents)
 
-proc agentBuild*(cq: Conquest, args: varargs[string]) = 
+# Terminate agent and remove it from the database
-    discard
+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.")
+        return
+
+    # TODO: Stop the process of the agent on the target system
+    # TODO: Add flag to self-delete executable after killing agent
+
+
+    # Remove the agent from the database
+    if not cq.dbDeleteAgentByName(name.toUpperAscii): 
+        cq.writeLine(fgRed, styleBright, "[-] Failed to terminate agent: ", getCurrentExceptionMsg())
+        return
+
+    cq.delAgent(name)
+    cq.writeLine(fgYellow, styleBright, "[+] ", resetStyle, "Terminated agent ", fgYellow, styleBright, name.toUpperAscii, resetStyle, ".")
 
 # Switch to interact mode
 proc agentInteract*(cq: Conquest, args: varargs[string]) = 
 
     cq.setIndicator("[AGENT] (username@hostname)> ")
-    cq.setStatusBar(@[("mode", "interact"), ("listeners", "X"), ("agents", "4")])    
+    cq.setStatusBar(@[("[mode]", "interact"), ("[username]", "X"), ("[hostname]", "4"), ("[ip]", "127.0.0.1"), ("[domain]", "domain.local")])    
 
     var command: string = cq.readLine()
 
-    discard
-
 #[
   Agent API
   Functions relevant for dealing with the agent API, such as registering new agents, querying tasks and posting results
@@ -51,7 +67,7 @@ proc register*(agent: Agent): bool =
         # Check if listener that is requested exists
         # TODO: Verify that the listener accessed is also the listener specified in the URL
         # This can be achieved by extracting the port number from the `Host` header and matching it to the one queried from the database
-        if not cq.listenerExists(agent.listener): 
+        if not cq.dbListenerExists(agent.listener): 
             cq.writeLine(fgRed, styleBright, fmt"[-] Agent from {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
             return false
 
@@ -60,10 +76,11 @@ proc register*(agent: Agent): bool =
             cq.writeLine(fgRed, styleBright, fmt"[-] Failed to insert agent {agent.name} into database.", "\n")
             return false
 
-        cq.add(agent.name, agent)
+        cq.add(agent)
         cq.writeLine(fgYellow, styleBright, fmt"[{agent.firstCheckin}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.name, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listener, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") 
 
     return true
+    
 #[
     Agent interaction mode
     When interacting with a agent, the following functions are called:

server/db/database.nim

@@ -14,9 +14,7 @@ proc dbInit*(cq: Conquest) =
             name TEXT PRIMARY KEY,
             address TEXT NOT NULL,
             port INTEGER NOT NULL UNIQUE,
-            protocol TEXT NOT NULL CHECK (protocol IN ('http')),
+            protocol TEXT NOT NULL CHECK (protocol IN ('http'))
-            sleep INTEGER NOT NULL,
-            jitter REAL NOT NULL
         );
 
         CREATE TABLE agents (
@@ -26,6 +24,7 @@ proc dbInit*(cq: Conquest) =
             pid INTEGER NOT NULL,
             username TEXT NOT NULL,
             hostname TEXT NOT NULL,
+            domain TEXT NOT NULL,
             ip TEXT NOT NULL,
             os TEXT NOT NULL,
             elevated BOOLEAN NOT NULL,
@@ -43,7 +42,7 @@ proc dbInit*(cq: Conquest) =
         cq.writeLine(fgGreen, "[+] ", cq.dbPath, ": Database file found.")
 
 #[
-    Listeners
+    Listener database functions
 ]#
 proc dbStoreListener*(cq: Conquest, listener: Listener): bool = 
 
@@ -51,9 +50,9 @@ proc dbStoreListener*(cq: Conquest, listener: Listener): bool =
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
         conquestDb.exec("""
-        INSERT INTO listeners (name, address, port, protocol, sleep, jitter)
+        INSERT INTO listeners (name, address, port, protocol)
-        VALUES (?, ?, ?, ?, ?, ?);
+        VALUES (?, ?, ?, ?);
-        """, listener.name, listener.address, listener.port, $listener.protocol, listener.sleep, listener.jitter)
+        """, listener.name, listener.address, listener.port, $listener.protocol)
 
         conquestDb.close() 
     except: 
@@ -69,16 +68,14 @@ proc dbGetAllListeners*(cq: Conquest): seq[Listener] =
     try: 
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
-        for row in conquestDb.iterate("SELECT name, address, port, protocol, sleep, jitter FROM listeners;"):
+        for row in conquestDb.iterate("SELECT name, address, port, protocol FROM listeners;"):
-            let (name, address, port, protocol, sleep, jitter) = row.unpack((string, string, int, string, int, float ))
+            let (name, address, port, protocol) = row.unpack((string, string, int, string))
             
             let l = Listener(
                 name: name,
                 address: address,
                 port: port,
                 protocol: stringToProtocol(protocol),
-                sleep: sleep,
-                jitter: jitter
             )
             listeners.add(l)
 
@@ -100,7 +97,7 @@ proc dbDeleteListenerByName*(cq: Conquest, name: string): bool =
     
     return true
 
-proc listenerExists*(cq: Conquest, listenerName: string): bool =
+proc dbListenerExists*(cq: Conquest, listenerName: string): bool =
     try:
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
@@ -114,7 +111,7 @@ proc listenerExists*(cq: Conquest, listenerName: string): bool =
         return false
 
 #[
-    Agents
+    Agent database functions
 ]#
 proc dbStoreAgent*(cq: Conquest, agent: Agent): bool = 
     
@@ -122,9 +119,9 @@ proc dbStoreAgent*(cq: Conquest, agent: Agent): bool =
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
         conquestDb.exec("""
-        INSERT INTO agents (name, listener, sleep, jitter, process, pid, username, hostname, ip, os, elevated, firstCheckin)
+        INSERT INTO agents (name, listener, process, pid, username, hostname, domain, ip, os, elevated, sleep, jitter, firstCheckin)
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
-        """, agent.name, agent.listener, agent.sleep, agent.jitter, agent.process, agent.pid, agent.username, agent.hostname, agent.ip, agent.os, agent.elevated, $agent.firstCheckin)
+        """, agent.name, agent.listener, agent.process, agent.pid, agent.username, agent.hostname, agent.domain, agent.ip, agent.os, agent.elevated, agent.sleep, agent.jitter, $agent.firstCheckin)
 
         conquestDb.close() 
     except: 
@@ -140,22 +137,23 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
     try: 
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
-        for row in conquestDb.iterate("SELECT name, listener, sleep, jitter, process, pid, username, hostname, ip, os, elevated, firstCheckin FROM agents;"):
+        for row in conquestDb.iterate("SELECT name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin FROM agents;"):
-            let (name, listener, sleep, jitter, process, pid, username, hostname, ip, os, elevated, firstCheckin) = row.unpack((string, string, int, float, string, int, string, string, string, string, bool, string))
+            let (name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin) = row.unpack((string, string, int, float, string, int, string, string, string, string, string, bool, string))
 
             let a = Agent(
                     name: name,
                     listener: listener,
                     sleep: sleep,
-                    jitter: jitter,
-                    process: process,
                     pid: pid,
                     username: username,
                     hostname: hostname,
+                    domain: domain,
                     ip: ip,
                     os: os,
                     elevated: elevated,
                     firstCheckin: firstCheckin,
+                    jitter: jitter,
+                    process: process,
                     tasks: @[] 
                 )
 
@@ -165,4 +163,29 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
     except: 
         cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
 
-    return agents
+    return agents
+
+proc dbDeleteAgentByName*(cq: Conquest, name: string): bool =
+    try: 
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+
+        conquestDb.exec("DELETE FROM agents WHERE name = ?", name)
+
+        conquestDb.close()
+    except: 
+        return false
+    
+    return true
+
+proc dbAgentExists*(cq: Conquest, agentName: string): bool =
+    try:
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+
+        let res = conquestDb.one("SELECT 1 FROM agents WHERE name = ? LIMIT 1", agentName)
+        
+        conquestDb.close()
+
+        return res.isSome
+    except:
+        cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
+        return false

server/listener/api.nim

@@ -24,6 +24,7 @@ proc register*(ctx: Context) {.async.} =
         {
             "username": "username",
             "hostname":"hostname",
+            "domain": "domain.local",
             "ip": "ip-address",
             "os": "operating-system",
             "process": "agent.exe",

server/listener/listener.nim

@@ -62,7 +62,7 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
     # Start serving
     try:
         discard listener.runAsync() 
-        cq.add(listenerInstance.name, listenerInstance)
+        cq.add(listenerInstance)
         cq.writeLine(fgGreen, "[+] ", resetStyle, "Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on port {portStr}.")
     except CatchableError as err: 
         cq.writeLine(fgRed, styleBright, "[-] Failed to start listener: ", getCurrentExceptionMsg())
@@ -89,7 +89,7 @@ proc restartListeners*(cq: Conquest) =
         
         try:
             discard listener.runAsync() 
-            cq.add(l.name, l)
+            cq.add(l)
             cq.writeLine(fgGreen, "[+] ", resetStyle, "Restarted listener", fgGreen, fmt" {l.name} ", resetStyle, fmt"on port {$l.port}.")
         except CatchableError as err: 
             cq.writeLine(fgRed, styleBright, "[-] Failed to restart listener: ", getCurrentExceptionMsg())
@@ -99,8 +99,15 @@ proc restartListeners*(cq: Conquest) =
 
     cq.writeLine("")
 
+# Remove listener from database, preventing automatic startup on server restart
 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.")
+        return
+
+    # Remove database entry
     if not cq.dbDeleteListenerByName(name.toUpperAscii): 
         cq.writeLine(fgRed, styleBright, "[-] Failed to stop listener: ", getCurrentExceptionMsg())
         return

server/server.nim

@@ -33,9 +33,9 @@ var parser = newParser:
             help("List all agents.")
             # TODO: Add a flag that allows the user to only list agents that are connected to a specific listener (-n <uuid>)
 
-        command("build"):
+        command("kill"):
-            help("Build an agent to connect to an active listener.")
+            help("Terminate the connection of an active listener and remove it from the interface.")
-
+            option("-n", "-name", help="Name of the agent to stop.", required=true)
 
         command("interact"):
             help("Interact with an active agent.")
@@ -82,8 +82,8 @@ proc handleConsoleCommand*(cq: Conquest, args: varargs[string]) =
             case opts.agent.get.command
             of "list":
                 cq.agentList()
-            of "build": 
+            of "kill": 
-                cq.agentBuild()
+                cq.agentKill(opts.agent.get.kill.get.name)
             of "interact":
                 cq.agentInteract() 
             else: 
@@ -128,13 +128,14 @@ proc main() =
     # Initialize database
     cq.dbInit()
     cq.restartListeners()
+    cq.addMutliple(cq.dbGetAllAgents())
 
     # Main loop
     while true: 
         cq.setIndicator("[conquest]> ")
-        cq.setStatusBar(@[("mode", "manage"), ("listeners", $len(cq.listeners)), ("agents", $len(cq.agents))])    
+        cq.setStatusBar(@[("[mode]", "manage"), ("[listeners]", $len(cq.listeners)), ("[agents]", $len(cq.agents))])    
         cq.showPrompt() 
-            
+ 
         var command: string = cq.readLine()
         cq.withOutput(handleConsoleCommand, command)
 

server/types.nim

@@ -3,7 +3,7 @@ import prologue
 import tables, times
 
 #[
-    Agent
+    Agent types & procs
 ]#
 type 
 
@@ -33,55 +33,39 @@ type
     AgentRegistrationData* = object
         username*: string
         hostname*: string
+        domain*: string
         ip*: string
         os*: string 
         process*: string
         pid*: int 
         elevated*: bool
 
-    # TODO: Add additional fields: domain, ...
     Agent* = ref object 
         name*: string
         listener*: string 
-        sleep*: int 
-        jitter*: float 
-        process*: string
-        pid*: int
         username*: string 
         hostname*: string
+        domain*: string
+        process*: string
+        pid*: int
         ip*: string
         os*: string
         elevated*: bool 
+        sleep*: int 
+        jitter*: float 
         tasks*: seq[Task]
         firstCheckin*: string
         lastCheckin*: string
 
-proc newAgent*(name, listener, username, hostname, process, ip, os, firstCheckin: string, pid: int, elevated: bool): Agent = 
-    var agent = new Agent
-    agent.name = name 
-    agent.listener = listener 
-    agent.process = process
-    agent.pid = pid 
-    agent.username = username 
-    agent.hostname = hostname 
-    agent.ip = ip 
-    agent.os = os
-    agent.elevated = elevated 
-    agent.sleep = 10
-    agent.jitter = 0.2
-    agent.tasks = @[]
-    agent.firstCheckin = firstCheckin
-    
-    return agent
-
 proc newAgent*(name, listener, firstCheckin: string, postData: AgentRegistrationData): Agent = 
     var agent = new Agent
     agent.name = name 
     agent.listener = listener 
-    agent.process = postData.process
-    agent.pid = postData.pid
     agent.username = postData.username 
     agent.hostname = postData.hostname 
+    agent.domain = postData.domain
+    agent.process = postData.process
+    agent.pid = postData.pid
     agent.ip = postData.ip 
     agent.os = postData.os
     agent.elevated = postData.elevated 
@@ -93,7 +77,7 @@ proc newAgent*(name, listener, firstCheckin: string, postData: AgentRegistration
     return agent
 
 #[
-    Listener 
+    Listener types and procs
 ]#
 type 
     Protocol* = enum
@@ -104,8 +88,6 @@ type
         address*: string
         port*: int
         protocol*: Protocol
-        sleep*: int 
-        jitter*: float 
 
 proc newListener*(name: string, address: string, port: int): Listener = 
     var listener = new Listener
@@ -113,8 +95,6 @@ proc newListener*(name: string, address: string, port: int): Listener =
     listener.address = address 
     listener.port = port 
     listener.protocol = HTTP
-    listener.sleep = 5          # 5 seconds beaconing 
-    listener.jitter = 0.2       # 20% Jitter
 
     return listener
 
@@ -126,7 +106,7 @@ proc stringToProtocol*(protocol: string): Protocol =
 
 
 #[
-    Conquest framework
+    Conquest framework types & procs
 ]#
 type 
     Conquest* = ref object
@@ -135,11 +115,15 @@ type
         listeners*: Table[string, Listener]
         agents*: Table[string, Agent]
 
-proc add*(cq: Conquest, listenerName: string, listener: Listener) = 
+proc add*(cq: Conquest, listener: Listener) = 
-    cq.listeners[listenerName] = listener
+    cq.listeners[listener.name] = listener
 
-proc add*(cq: Conquest, agentName: string, agent: Agent) = 
+proc add*(cq: Conquest, agent: Agent) = 
-    cq.agents[agentName] = agent
+    cq.agents[agent.name] = agent
+
+proc addMutliple*(cq: Conquest, agents: seq[Agent]) = 
+    for a in agents: 
+        cq.agents[a.name] = a
 
 proc delListener*(cq: Conquest, listenerName: string) = 
     cq.listeners.del(listenerName)