Files
conquest/server/db/dbAgent.nim

117 lines
4.0 KiB
Nim
Raw Normal View History

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