117 lines
4.0 KiB
Nim
117 lines
4.0 KiB
Nim
import system, terminal, tiny_sqlite
|
|
import ../types
|
|
|
|
#[
|
|
Agent database functions
|
|
]#
|
|
proc dbStoreAgent*(cq: Conquest, agent: Agent): bool =
|
|
|
|
try:
|
|
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)
|
|
|
|
conquestDb.close()
|
|
except:
|
|
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
|
|
return false
|
|
|
|
return true
|
|
|
|
proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
|
|
|
|
var agents: 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))
|
|
|
|
let a = Agent(
|
|
name: name,
|
|
listener: listener,
|
|
sleep: sleep,
|
|
pid: pid,
|
|
username: username,
|
|
hostname: hostname,
|
|
domain: domain,
|
|
ip: ip,
|
|
os: os,
|
|
elevated: elevated,
|
|
firstCheckin: firstCheckin,
|
|
jitter: jitter,
|
|
process: process,
|
|
tasks: @[]
|
|
)
|
|
|
|
agents.add(a)
|
|
|
|
conquestDb.close()
|
|
except:
|
|
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
|
|
|
|
return agents
|
|
|
|
proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
|
|
|
|
var agents: 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))
|
|
|
|
let a = Agent(
|
|
name: name,
|
|
listener: listener,
|
|
sleep: sleep,
|
|
pid: pid,
|
|
username: username,
|
|
hostname: hostname,
|
|
domain: domain,
|
|
ip: ip,
|
|
os: os,
|
|
elevated: elevated,
|
|
firstCheckin: firstCheckin,
|
|
jitter: jitter,
|
|
process: process,
|
|
tasks: @[]
|
|
)
|
|
|
|
agents.add(a)
|
|
|
|
conquestDb.close()
|
|
except:
|
|
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
|
|
|
|
return agents
|
|
|
|
proc dbDeleteAgentByName*(cq: Conquest, name: string): bool =
|
|
try:
|
|
let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
|
|
|
|
conquestDb.exec("DELETE FROM agents WHERE name = ?", name)
|
|
|
|
conquestDb.close()
|
|
except:
|
|
return false
|
|
|
|
return true
|
|
|
|
proc dbAgentExists*(cq: Conquest, agentName: string): bool =
|
|
try:
|
|
let conquestDb = openDatabase(cq.dbPath, mode=dbReadWrite)
|
|
|
|
let res = conquestDb.one("SELECT 1 FROM agents WHERE name = ? LIMIT 1", agentName)
|
|
|
|
conquestDb.close()
|
|
|
|
return res.isSome
|
|
except:
|
|
cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg())
|
|
return false |