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 
                )

            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,
                )

            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