From 56f8994069ae7b7372684f5316b19608da000b3f Mon Sep 17 00:00:00 2001
From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com>
Date: Thu, 15 May 2025 15:24:46 +0200
Subject: [PATCH] Added number of agents to listener list

---
 .gitignore             |  2 +-
 server/agent/agent.nim | 21 ++++++++++-----------
 server/server.nim      | 10 +++-------
 server/utils.nim       |  8 +++++---
 4 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9978ff2..4bbd08d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,6 @@
 # Ignore agents
 agents/
-
+*.db
 
 # Nim
 nimcache/
diff --git a/server/agent/agent.nim b/server/agent/agent.nim
index 7aad55f..4b2a9b4 100644
--- a/server/agent/agent.nim
+++ b/server/agent/agent.nim
@@ -24,19 +24,18 @@ Options:
   -h, --help""")
 
 # List agents
-proc agentList*(cq: Conquest) = 
-    let agents = cq.dbGetAllAgents()
-    cq.drawTable(agents)
-
 proc agentList*(cq: Conquest, listener: string) =
 
-    # Check if listener exists
-    if not cq.dbListenerExists(listener.toUpperAscii): 
-        cq.writeLine(fgRed, styleBright, fmt"[-] Listener {listener.toUpperAscii} does not exist.")
-        return
+    # If no argument is passed via -n, list all agents, otherwise only display agents connected to a specific listener
+    if listener == "": 
+        cq.drawTable(cq.dbGetAllAgents())
+    else: 
+        # Check if listener exists
+        if not cq.dbListenerExists(listener.toUpperAscii): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Listener {listener.toUpperAscii} does not exist.")
+            return
 
-    let agents = cq.dbGetAllAgentsByListener(listener.toUpperAscii) 
-    cq.drawTable(agents)
+        cq.drawTable(cq.dbGetAllAgentsByListener(listener.toUpperAscii))
 
 # Display agent properties and details
 proc agentInfo*(cq: Conquest, name: string) = 
@@ -47,7 +46,7 @@ proc agentInfo*(cq: Conquest, name: string) =
 
     let agent = cq.agents[name.toUpperAscii]
 
-    # TODO: Improve formating
+    # TODO: Improve formatting
     cq.writeLine(fmt"""
 Agent name (UUID):     {agent.name}
 Connected to listener: {agent.listener}
diff --git a/server/server.nim b/server/server.nim
index 11a8b4a..ca55863 100644
--- a/server/server.nim
+++ b/server/server.nim
@@ -32,8 +32,7 @@ var parser = newParser:
 
         command("list"):
             help("List all agents.")
-            option("-n", "-name", help="Name of the listener.")
-            # TODO: Add a flag that allows the user to only list agents that are connected to a specific listener (-n <uuid>)
+            option("-l", "-listener", help="Name of the listener.")
 
         command("info"): 
             help("Display details for a specific agent.")
@@ -86,11 +85,8 @@ proc handleConsoleCommand*(cq: Conquest, args: varargs[string]) =
 
         of "agent":
             case opts.agent.get.command
-            of "list":
-                if opts.agent.get.list.get.name == "": 
-                    cq.agentList()
-                else: 
-                    cq.agentList(opts.agent.get.list.get.name)
+            of "list":    
+                cq.agentList(opts.agent.get.list.get.listener)
             of "info":
                 cq.agentInfo(opts.agent.get.info.get.name)
             of "kill": 
diff --git a/server/utils.nim b/server/utils.nim
index 82766d6..3646f42 100644
--- a/server/utils.nim
+++ b/server/utils.nim
@@ -1,4 +1,4 @@
-import re, strutils, strformat, terminal
+import re, strutils, strformat, terminal, tables, sequtils
 
 import ./types
 
@@ -61,8 +61,10 @@ proc drawTable*(cq: Conquest, listeners: seq[Listener]) =
     cq.writeLine(border(midLeft, midMid, midRight, widths))
 
     for l in listeners:
-        # TODO: Add number of agents connected to the listener
-        let row = @[l.name, l.address, $l.port, $l.protocol, "X"]
+        # Get number of agents connected to the listener
+        let connectedAgents = cq.agents.values.countIt(it.listener == l.name)
+
+        let row = @[l.name, l.address, $l.port, $l.protocol, $connectedAgents]
         cq.writeLine(row(row, widths)) 
 
     cq.writeLine(border(botLeft, botMid, botRight, widths))