Implementing a custom Command Plugin

As part of this tutorial, we will create the simplest possible command which echoes back the input it receives. This command will be named echo. On server side, the echo command must be created by the user with a message. This command will be sent to agent, where the command shellcode will reply the same message back to the server.

Some of the details in this example are omitted, the full example can be found in the example plugin repository.

Implementing a command plugin starts from creating a new Java class that extends the CommandPluginclass:

public class EchoCommandPlugin implements CommandPlugin {

  private volatile List<CommandTemplate> cachedCommandTemplates = List.of();

  @Override
  public void init(CommandPluginContext pluginContext) {
    if (!cachedCommandTemplates.isEmpty()) {
      return;
    }
    this.cachedCommandTemplates = List.of(new EchoCommandTemplate(pluginContext));
  }

  @Override
  public List<? extends CommandTemplate> getCommandTemplates() {
    return cachedCommandTemplates;
  }
}

As one can see, the plugin interface for commands is quite minimal, the real logic is in the CommandTemplate classes that are returned by the getCommandTemplates method.

public class EchoCommandTemplate implements CommandTemplate {

  private final CommandPluginContext pluginContext;

  public EchoCommandTemplate(CommandPluginContext pluginContext) {
    this.pluginContext = pluginContext;
  }

  @Override
  public String getName() {
    return "echo";
  }

  @Override
  public String getDescription() {
    return "Echoes the message sent to agent back to the server";
  }

  @Override
  public List<NamedConfiguration> getExampleConfigurations() throws SerializationException {
    return List.of(new NamedConfiguration("hello-world", () -> """
        {
          "message": "Hello, World!"
        }
        """));
  }

  @Override
  public ConfigurationSchema getConfigurationSchema() throws SerializationException {
    return myJsonLibrary.createSchema("""
        {
          "type": "object",
          "properties": {
            "message": {
              "type": "string",
              "description": "Message to echo back to the server"
            }
          },
          "required": ["message"]
        }
        """);
  }

  @Override
  public boolean canSendToAgent(AgentInfo agentInfo) {
    /* Validates that this command is only allowed to be sent to agents from shellcode listeners */
    return AgentType.SHELLCODE_AGENT == agentInfo.getAgentType();
  }

  @Override
  public void validateConfiguration(AgentMetadata agentMetadata, Configuration configuration)
      throws ValidationException {
    /* 
    Parsing and validating configuration is a task left to the plugin developer.
    Tuoni SDK does not enforce using any particular JSON library 
    */
    myJsonLibrary.validate(configuration, required("message"));
  }

  @Override
  public Command createCommand(
      int commandId,
      AgentInfo agentInfo,
      Configuration configuration,
      CommandContext commandContext)
      throws ValidationException, InitializationException {
    validateConfiguration(configuration);
    var message = myJsonLibrary.parse(configuration).get("message").asText();
    return new EchoCommand(commandId, agentInfo, commandContext, message);
  }
}

Most command templates have repetitive parts, so it is recommended to create a base class that implements the common parts and then extend that class for each command. \ Last step is to implement the actual command that is sent to the agent:

public class EchoCommand implements ShellcodeCommand {

  private final int commandId;
  private final AgentInfo agentInfo;
  private final CommandContext commandContext;
  private final String message;

  public EchoCommand(int commandId, AgentInfo agentInfo, CommandContext commandContext, String message) {
    this.commandId = commandId;
    this.agentInfo = agentInfo;
    this.commandContext = commandContext;
    this.message = message;
  }

  @Override
  public ShellCodeWithConf generateShellCode(String pipeName, AgentMetadata latestAgentMetadata)
      throws SerializationException, ValidationException {
    // Read and manipulate the shellcode
    var shellCodeBytes = MyShellCodeUtil.readShellcode("echo")
        .replaceDefaultPipeNameWith(pipeName)
        .asByteBuffer();

    return new ShellCodeWithConf(
        shellCodeBytes, 
        MyShellCodeUtil.serializeMessageForShellcode().serializeForShellcode(message),
        PluginIpcType.NAMED_PIPE);
  }

  @Override
  public void parseResult(
      ByteBuffer buffer,
      boolean isFinalResult,
      CommandResultCollection previousResult,
      CommandResultEditor editor)
      throws SerializationException {
    // Read result from received result bytes
    // Structure of the result is defined by the shellcode that was sent to the agent
    String receivedString = StandardCharsets.UTF_8.decode(buffer).toString();
    editor.setTextResult("message", receivedString);
    editor.commit();
  }

  @Override
  public ByteBuffer serializeCommandUpdate(Configuration updateConfiguration)
      throws SerializationException, ValidationException, CommandUpdateUnsupportedException {
    // This command does not support updating already sent commands
    throw CommandUpdateUnsupportedException.forTemplateName(EchoCommandTemplate.NAME);
  }

  @Override
  public void markStatus(CommandStatus status) {
    // This is a stateless command, but for stateful commands this method should act as trigger
    //   for the command state machine.
    // If status is COMPLETE, then the command is considered done and all resources can be released
  }

  @Override
  public void forceStop() throws ExecutionException {
    // In case the command has any background threads or resources, they must be stopped here
  }
}

How does this plugin work?

We'll assume that all the omitted parts are implemented correctly and the plugin JAR is bundled correctly with manifest, service files, built and deployed to the Tuoni server plugins directory.

In that case when starting Tuoni server, for shellcode agents, the server will load the plugin and the user can create a new command over the REST API with the name echo. The command will go through the following stages:

1. Server starts - Tuoni server calls EchoCommandPlugin#init and EchoCommandPlugin#getCommandTemplates to get the command templates.
2. User asks for available command templates for an agent over the REST API - Tuoni server calls EchoCommandTemplate#canSendToAgent and in case of success, returns the echo command template.
3. Command created over REST API - User chooses an agent, command and provides configuration for the command. Tuoni server calls EchoCommandTemplate#validateConfiguration with the configuration and information about the agent. If the configuration is valid, the command is created and saved to the database. Otherwise, an error is returned to the user.
4. Command is queued for all the listeners to which the agent is connected - Tuoni server calls EchoCommand#createCommand to create the command object.
5. The listener polls command for sending to the agent - EchoCommand#generateShellCode is called to get the shellcode that is sent to the agent.
6. Agent receives the shellcode and executes it - The agent executes the shellcode and sends the result back to the server.
7. Server receives the result - The server calls EchoCommand#parseResult to parse the result and save it to the database.
8. User asks for the result over the REST API - The server returns the result to the user.