Prevent database locking by not updating latest checkin in database and instead storing it only in memory

This commit is contained in:
Jakob Friedl
2025-05-29 14:19:55 +02:00
parent d35b1d0a0d
commit 118e9eadd2
7 changed files with 40 additions and 11 deletions

View File

@@ -18,7 +18,8 @@ proc register*(config: AgentConfig): string =
"os": getOSVersion(),
"process": getProcessExe(),
"pid": getProcessId(),
"elevated": isElevated()
"elevated": isElevated(),
"sleep": config.sleep
}
echo $body

View File

@@ -2,4 +2,4 @@
-d:ListenerUuid="KPDHWZNT"
-d:ListenerIp="localhost"
-d:ListenerPort=7777
-d:SleepDelay=10
-d:SleepDelay=0

View File

@@ -30,6 +30,7 @@ proc agentList*(cq: Conquest, listener: string) =
# 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):
@@ -38,6 +39,7 @@ proc agentList*(cq: Conquest, listener: string) =
cq.drawTable(cq.dbGetAllAgentsByListener(listener.toUpperAscii))
# Display agent properties and details
proc agentInfo*(cq: Conquest, name: string) =
# Check if agent supplied via -n parameter exists in database
@@ -202,8 +204,8 @@ proc getTasks*(listener, agent: string): JsonNode =
# 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
# if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")):
# return nil
# Return tasks in JSON format
return %cq.agents[agent.toUpperAscii].tasks.filterIt(it.status != Completed)

View File

@@ -4,6 +4,10 @@ import ../../db/database
proc taskExecuteSleep*(cq: Conquest, delay: int) =
if delay < 0:
cq.writeLine(fgRed, styleBright, "[-] Invalid sleep delay value.")
return
# Update 'sleep' value in database
if not cq.dbUpdateSleep(cq.interactAgent.name, delay):
return

View File

@@ -29,7 +29,8 @@ proc register*(ctx: Context) {.async.} =
"os": "operating-system",
"process": "agent.exe",
"pid": 1234,
"elevated": false
"elevated": false.
"sleep": 10
}
]#

View File

@@ -1,6 +1,6 @@
import prompt
import prologue
import tables
import tables, sequtils
import times
import terminal
@@ -25,6 +25,19 @@ type
TaskResult* = string
#[
TaskResult*[T] = ref object
data*: T
Task*[T] = ref object
id*: string
agent*: string
command*: TaskCommand
args*: seq[string]
result*: TaskResult[T]
status*: TaskStatus
]#
Task* = ref object
id*: string
agent*: string
@@ -42,6 +55,7 @@ type
process*: string
pid*: int
elevated*: bool
sleep*: int
Agent* = ref object
name*: string
@@ -60,6 +74,7 @@ type
firstCheckin*: DateTime
latestCheckin*: DateTime
# TODO: Take sleep value from agent registration data (set via nim.cfg file)
proc newAgent*(name, listener: string, firstCheckin: DateTime, postData: AgentRegistrationData): Agent =
var agent = new Agent
agent.name = name
@@ -72,7 +87,7 @@ proc newAgent*(name, listener: string, firstCheckin: DateTime, postData: AgentRe
agent.ip = postData.ip
agent.os = postData.os
agent.elevated = postData.elevated
agent.sleep = 10
agent.sleep = postData.sleep
agent.jitter = 0.2
agent.tasks = @[]
agent.firstCheckin = firstCheckin
@@ -136,6 +151,12 @@ proc delListener*(cq: Conquest, listenerName: string) =
proc delAgent*(cq: Conquest, agentName: string) =
cq.agents.del(agentName)
proc getAgentsAsSeq*(cq: Conquest): seq[Agent] =
var agents: seq[Agent] = @[]
for agent in cq.agents.values:
agents.add(agent)
return agents
proc initConquest*(dbPath: string): Conquest =
var cq = new Conquest
var prompt = Prompt.init()

View File

@@ -71,7 +71,7 @@ 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 widths = @[8, 15, 5, 8, 6]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
@@ -134,7 +134,7 @@ proc timeSince*(agent: Agent, timestamp: DateTime): Cell =
proc drawTable*(cq: Conquest, agents: seq[Agent]) =
let headers: seq[string] = @["Name", "Address", "Username", "Hostname", "Operating System", "Process", "PID", "Activity"]
let widths = @[10, 17, 15, 15, 18, 15, 7, 10]
let widths = @[8, 15, 15, 15, 16, 15, 5, 13]
let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault))
cq.writeLine(border(topLeft, topMid, topRight, widths))
@@ -154,7 +154,7 @@ proc drawTable*(cq: Conquest, agents: seq[Agent]) =
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),
a.timeSince(a.latestCheckin)
a.timeSince(cq.agents[a.name].latestCheckin)
]
# Highlight agents running within elevated processes