Implemented displaying latest checkin in agents table, as well as word-wrap.
This commit is contained in:
@@ -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."
|
||||
|
||||
|
||||
@@ -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(" "))
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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) =
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
132
server/utils.nim
132
server/utils.nim
@@ -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))
|
||||
Reference in New Issue
Block a user