diff --git a/agents/monarch/client.nim b/agents/monarch/client.nim index 762db26..af3493c 100644 --- a/agents/monarch/client.nim +++ b/agents/monarch/client.nim @@ -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." diff --git a/agents/monarch/commands/shell.nim b/agents/monarch/commands/shell.nim index 12f1600..ce7e3a8 100644 --- a/agents/monarch/commands/shell.nim +++ b/agents/monarch/commands/shell.nim @@ -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(" ")) diff --git a/agents/monarch/task.nim b/agents/monarch/task.nim index 5bc4ea0..03764b3 100644 --- a/agents/monarch/task.nim +++ b/agents/monarch/task.nim @@ -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( diff --git a/server/agent/agent.nim b/server/agent/agent.nim index 61f5816..25c461e 100644 --- a/server/agent/agent.nim +++ b/server/agent/agent.nim @@ -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) = diff --git a/server/agent/interact.nim b/server/agent/interact.nim index 4e59fc1..90959ab 100644 --- a/server/agent/interact.nim +++ b/server/agent/interact.nim @@ -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 - diff --git a/server/db/database.nim b/server/db/database.nim index f01a2c9..1cbb45b 100644 --- a/server/db/database.nim +++ b/server/db/database.nim @@ -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.") diff --git a/server/db/dbAgent.nim b/server/db/dbAgent.nim index 7870ca2..5e4e551 100644 --- a/server/db/dbAgent.nim +++ b/server/db/dbAgent.nim @@ -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, ) @@ -110,6 +112,18 @@ proc dbAgentExists*(cq: Conquest, agentName: string): bool = conquestDb.close() return res.isSome + 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 \ No newline at end of file diff --git a/server/listener/api.nim b/server/listener/api.nim index 9e127c4..5427c96 100644 --- a/server/listener/api.nim +++ b/server/listener/api.nim @@ -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) diff --git a/server/types.nim b/server/types.nim index 31673c7..b8fc548 100644 --- a/server/types.nim +++ b/server/types.nim @@ -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 diff --git a/server/utils.nim b/server/utils.nim index 25d19f1..3d4df2a 100644 --- a/server/utils.nim +++ b/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] - - # Highlight agents running within elevated processes - if a.elevated: - cq.writeLine(bgRed, fgBlack, row(row, widths)) - else: - cq.writeLine(row(row, widths)) + 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 + for line in formatRow(cells, widths): + cq.hidePrompt() + cq.writeRow(line) + cq.showPrompt() cq.writeLine(border(botLeft, botMid, botRight, widths)) \ No newline at end of file