Implemented agent registration, restructured Conquest type to utilize tables to store agents and listeners

README.md

@@ -1 +1,9 @@
 # Conquest Command & Control Framework
+
+
+## Acknowledgements
+
+- [C5pider](https://github.com/Cracked5pider) for [Havoc](https://github.com/HavocFramework/Havoc), which most of the teamserver functionality is based on
+- [m4ul3r](https://github.com/m4ul3r) for [nimless](https://github.com/m4ul3r/writing_nimless) Nim implementations
+- [d4rckh](https://github.com/d4rckh) for [grc2](https://github.com/d4rckh/grc2), the only other Nim-only C2 I was able to find online
+- [MalDev Academy](https://maldevacademy.com/)

server/agent/agent.nim

@@ -1,5 +1,6 @@
 import terminal, strformat, times
 import ../[types, globals]
+import ../db/database
 
 
 #[ 
@@ -41,16 +42,29 @@ proc agentInteract*(cq: Conquest, args: varargs[string]) =
   Agent API
   Functions relevant for dealing with the agent API, such as registering new agents, querying tasks and posting results
 ]#
-proc notifyAgentRegister*(agent: Agent) = 
+proc register*(agent: Agent): bool = 
 
     let date: string = now().format("dd-MM-yyyy HH:mm:ss")
 
-
-
     # The following line is required to be able to use the `cq` global variable for console output
     {.cast(gcsafe).}:
-        cq.writeLine(fgYellow, styleBright, fmt"[{date}] Agent {agent.name} connected.", "\n") 
 
+        # 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): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Agent from {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
+            return false
+
+        # Store agent in database
+        if not cq.dbStoreAgent(agent): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Failed to insert agent {agent.name} into database.", "\n")
+            return false
+
+        cq.add(agent.name, agent)
+        cq.writeLine(fgYellow, styleBright, fmt"[{date}] ", 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/config.nims

@@ -1,2 +1,2 @@
 # Compiler flags
-# --threads:on
+--threads:on

server/db/database.nim

@@ -10,7 +10,7 @@ proc dbInit*(cq: Conquest) =
 
         # Create tables
         conquestDb.execScript("""
-        CREATE TABLE listener (
+        CREATE TABLE listeners (
             name TEXT PRIMARY KEY,
             address TEXT NOT NULL,
             port INTEGER NOT NULL UNIQUE,
@@ -19,6 +19,20 @@ proc dbInit*(cq: Conquest) =
             jitter REAL NOT NULL
         );
 
+        CREATE TABLE agents (
+            name TEXT PRIMARY KEY,                   
+            listener TEXT NOT NULL,                    
+            pid INTEGER NOT NULL,
+            username TEXT NOT NULL,
+            hostname TEXT NOT NULL,
+            ip TEXT NOT NULL,
+            os TEXT NOT NULL,
+            elevated BOOLEAN NOT NULL,
+            sleep INTEGER DEFAULT 10,
+            jitter REAL DEFAULT 0.1,
+            FOREIGN KEY (listener) REFERENCES listeners(name)
+        );
+
         """)
         
         cq.writeLine(fgGreen, "[+] ", cq.dbPath, ": Database created.")
@@ -32,7 +46,7 @@ proc dbStoreListener*(cq: Conquest, listener: Listener): bool =
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
         conquestDb.exec("""
-        INSERT INTO listener (name, address, port, protocol, sleep, jitter)
+        INSERT INTO listeners (name, address, port, protocol, sleep, jitter)
         VALUES (?, ?, ?, ?, ?, ?);
         """, listener.name, listener.address, listener.port, $listener.protocol, listener.sleep, listener.jitter)
 
@@ -50,7 +64,7 @@ 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 listener;"):
+        for row in conquestDb.iterate("SELECT name, address, port, protocol, sleep, jitter FROM listeners;"):
             let (name, address, port, protocol, sleep, jitter) = row.unpack((string, string, int, string, int, float ))
             
             let l = Listener(
@@ -73,7 +87,7 @@ proc dbDeleteListenerByName*(cq: Conquest, name: string): bool =
     try: 
         let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
 
-        conquestDb.exec("DELETE FROM listener WHERE name = ?", name)
+        conquestDb.exec("DELETE FROM listeners WHERE name = ?", name)
 
         conquestDb.close()
     except: 
@@ -81,5 +95,33 @@ proc dbDeleteListenerByName*(cq: Conquest, name: string): bool =
     
     return true
 
-proc dbStoreAgent*(agent: Agent): bool = 
+proc listenerExists*(cq: Conquest, listenerName: string): bool =
-    discard
+    try:
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+
+        let res = conquestDb.one("SELECT 1 FROM listeners WHERE name = ? LIMIT 1", listenerName)
+        
+        conquestDb.close()
+
+        return res.isSome
+    except:
+        cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
+        return false
+
+proc dbStoreAgent*(cq: Conquest, agent: Agent): bool = 
+    
+    try: 
+        let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
+
+        conquestDb.exec("""
+        INSERT INTO agents (name, listener, sleep, jitter, pid,username, hostname, ip, os, elevated)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
+        """, agent.name, agent.listener, agent.sleep, agent.jitter, agent.pid, agent.username, agent.hostname, agent.ip, agent.os, agent.elevated)
+
+        conquestDb.close() 
+    except: 
+        cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
+        return false
+    
+    return true
+

server/listener/api.nim

@@ -3,7 +3,9 @@ import terminal, sequtils, strutils
 
 import ../[types]
 import ../agent/agent
-import ./utils
+
+proc error404*(ctx: Context) {.async.} = 
+    resp "", Http404
 
 #[
     POST /{listener-uuid}/register
@@ -12,31 +14,47 @@ import ./utils
 proc register*(ctx: Context) {.async.} = 
 
     # Check headers
-    doAssert(ctx.request.getHeader("CONTENT-TYPE") == @["application/json"])
+    # If POST data is not JSON data, return 404 error code
-    doAssert(ctx.request.getHeader("USER-AGENT") == @["Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36"])
+    if ctx.request.contentType != "application/json": 
+        resp "", Http404
+        return
 
-    # Handle POST data, the register data should look like the following
+    # The JSON data for the agent registration has to be in the following format
     #[
         {
             "username": "username",
             "hostname":"hostname",
             "ip": "ip-address",
-            "os": "operating-system"
+            "os": "operating-system",
-            "pid":  1234
+            "pid":  1234,
             "elevated": false
         }
     ]#  
 
-    let 
+    try: 
-        postData: JsonNode = %ctx.request.body()
+        let 
-        name = generate(alphabet=join(toSeq('A'..'Z'), ""), size=8)
+            postData: JsonNode = parseJson(ctx.request.body)
+            agentRegistrationData = postData.to(AgentRegistrationData)
+            agentUuid = generate(alphabet=join(toSeq('A'..'Z'), ""), size=8)
+            listenerUuid = ctx.getPathParams("listener")
 
-    let agent = new Agent
+        let agent: Agent = newAgent(agentUuid, listenerUuid, agentRegistrationData) 
-    agent.name = name
-    notifyAgentRegister(agent)
         
+        # Fully register agent and add it to database
+        if not agent.register(): 
+            # Either the listener the agent tries to connect to does not exist in the database, or the insertion of the agent failed
+            # Return a 404 error code either way
+            resp "", Http404
+            return 
 
-    resp agent.name
+        # If registration is successful, the agent receives it's UUID, which is then used to poll for tasks and post results
+        resp agent.name
+
+    except CatchableError:
+        # JSON data is invalid or does not match the expected format (described above)
+        resp "", Http404
+
+    return
 
 #[
     GET /{listener-uuid}/{agent-uuid}/tasks

server/listener/listener.nim

@@ -52,16 +52,17 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
     listener.post("{listener}/register", api.register)
     listener.get("{listener}/{agent}/tasks", api.getTasks)
     listener.post("{listener}/{agent}/results", api.postResults)
+    listener.registerErrorHandler(Http404, api.error404)
 
     # Store listener in database
-    let listenerInstance = newListener(name, host, port)
+    var listenerInstance = newListener(name, host, port)
     if not cq.dbStoreListener(listenerInstance):
         return
 
     # Start serving
     try:
         discard listener.runAsync() 
-        inc cq.listeners
+        cq.add(listenerInstance.name, 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())
@@ -84,10 +85,11 @@ proc restartListeners*(cq: Conquest) =
         listener.post("{listener}/register", api.register)
         listener.get("{listener}/{agent}/tasks", api.getTasks)
         listener.post("{listener}/{agent}/results", api.postResults)
+        listener.registerErrorHandler(Http404, api.error404)
         
         try:
             discard listener.runAsync() 
-            inc cq.listeners
+            cq.add(l.name, 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())
@@ -103,6 +105,6 @@ proc listenerStop*(cq: Conquest, name: string) =
         cq.writeLine(fgRed, styleBright, "[-] Failed to stop listener: ", getCurrentExceptionMsg())
         return
 
-    dec cq.listeners
+    cq.delListener(name)
     cq.writeLine(fgGreen, "[+] ", resetStyle, "Stopped listener ", fgGreen, name.toUpperAscii, resetStyle, ".")
     

server/server.nim

@@ -1,6 +1,6 @@
 import prompt, terminal
 import argparse
-import strutils, strformat, times, system, unicode
+import strutils, strformat, times, system, tables
 
 import ./[types, globals]
 import agent/agent, listener/listener, db/database
@@ -21,6 +21,7 @@ var parser = newParser:
             option("-h", "-host", default=some("0.0.0.0"), help="IPv4 address to listen on.", required=false)
             option("-p", "-port", help="Port to listen on.", required=true)
             # flag("--dns", help="Use the DNS protocol for C2 communication.")
+            # flag("--doh", help="Use DNS over HTTPS for C2 communication.)
         command("stop"):
             help("Stop an active listener.")
             option("-n", "-name", help="Name of the listener to stop.", required=true)
@@ -131,7 +132,7 @@ proc main() =
     # Main loop
     while true: 
         cq.setIndicator("[conquest]> ")
-        cq.setStatusBar(@[("mode", "manage"), ("listeners", $cq.listeners), ("agents", $cq.agents)])    
+        cq.setStatusBar(@[("mode", "manage"), ("listeners", $len(cq.listeners)), ("agents", $len(cq.agents))])    
         cq.showPrompt() 
             
         var command: string = cq.readLine()

server/types.nim

@@ -1,46 +1,6 @@
 import prompt
 import prologue
-
+import tables
-#[
-    Conquest 
-]#
-type 
-    Conquest* = ref object
-        prompt*: Prompt
-        listeners*: int
-        agents*: int
-        dbPath*: string
-
-proc initConquest*(): Conquest = 
-    var cq = new Conquest
-    var prompt = Prompt.init()
-    cq.prompt = prompt
-    cq.dbPath = "db/conquest.db"
-    cq.listeners = 0
-    cq.agents = 0
-
-    return cq
-
-template writeLine*(cq: Conquest, args: varargs[untyped]) = 
-    cq.prompt.writeLine(args)
-proc readLine*(cq: Conquest): string =
-    return cq.prompt.readLine()
-template setIndicator*(cq: Conquest, indicator: string) = 
-    cq.prompt.setIndicator(indicator)
-template showPrompt*(cq: Conquest) = 
-    cq.prompt.showPrompt()
-template hidePrompt*(cq: Conquest) = 
-    cq.prompt.hidePrompt()
-template setStatusBar*(cq: Conquest, statusBar: seq[StatusBarItem]) = 
-    cq.prompt.setStatusBar(statusBar) 
-template clear*(cq: Conquest) = 
-    cq.prompt.clear()
-
-# Overwrite withOutput function to handle function arguments
-proc withOutput*(cq: Conquest, outputFunction: proc(cq: Conquest, args: varargs[string]), args: varargs[string]) =
-    cq.hidePrompt()
-    outputFunction(cq, args)
-    cq.showPrompt()
 
 #[
     Agent
@@ -107,6 +67,23 @@ proc newAgent*(name, listener, username, hostname, ip, os: string, pid: int, ele
 
     return agent
 
+proc newAgent*(name, listener: string, postData: AgentRegistrationData): Agent = 
+    var agent = new Agent
+    agent.name = name 
+    agent.listener = listener 
+    agent.pid = postData.pid
+    agent.username = postData.username 
+    agent.hostname = postData.hostname 
+    agent.ip = postData.ip 
+    agent.os = postData.os
+    agent.elevated = postData.elevated 
+    agent.sleep = 10
+    agent.jitter = 0.2
+    agent.tasks = @[]
+
+    return agent
+
+
 #[
     Listener 
 ]#
@@ -138,3 +115,57 @@ proc stringToProtocol*(protocol: string): Protocol =
     of "http": 
         return HTTP
     else: discard
+
+
+#[
+    Conquest 
+]#
+type 
+    Conquest* = ref object
+        prompt*: Prompt
+        dbPath*: string
+        listeners*: Table[string, Listener]
+        agents*: Table[string, Agent]
+
+proc add*(cq: Conquest, listenerName: string, listener: Listener) = 
+    cq.listeners[listenerName] = listener
+
+proc add*(cq: Conquest, agentName: string, agent: Agent) = 
+    cq.agents[agentName] = agent
+
+proc delListener*(cq: Conquest, listenerName: string) = 
+    cq.listeners.del(listenerName)
+
+proc delAgent*(cq: Conquest, agentName: string) = 
+    cq.agents.del(agentName)
+
+proc initConquest*(): Conquest = 
+    var cq = new Conquest
+    var prompt = Prompt.init()
+    cq.prompt = prompt
+    cq.dbPath = "db/conquest.db"
+    cq.listeners = initTable[string, Listener]()
+    cq.agents = initTable[string, Agent]() 
+
+    return cq
+
+template writeLine*(cq: Conquest, args: varargs[untyped]) = 
+    cq.prompt.writeLine(args)
+proc readLine*(cq: Conquest): string =
+    return cq.prompt.readLine()
+template setIndicator*(cq: Conquest, indicator: string) = 
+    cq.prompt.setIndicator(indicator)
+template showPrompt*(cq: Conquest) = 
+    cq.prompt.showPrompt()
+template hidePrompt*(cq: Conquest) = 
+    cq.prompt.hidePrompt()
+template setStatusBar*(cq: Conquest, statusBar: seq[StatusBarItem]) = 
+    cq.prompt.setStatusBar(statusBar) 
+template clear*(cq: Conquest) = 
+    cq.prompt.clear()
+
+# Overwrite withOutput function to handle function arguments
+proc withOutput*(cq: Conquest, outputFunction: proc(cq: Conquest, args: varargs[string]), args: varargs[string]) =
+    cq.hidePrompt()
+    outputFunction(cq, args)
+    cq.showPrompt()