Commands

SWAT (Simple Workspace ATT&CK Tool) is a modular and dynamic command-line tool that allows for the execution of various commands. This guide will dive into the structure of SWAT commands, explain how they work, and guide you through adding new commands to expand its functionality.

Prerequisites

Before you dive into adding commands to SWAT, ensure you have:

Familiarity with the argparse module in Python.
Basic understanding of the SWAT architecture, especially base_command.py and shell.py.
Necessary API credentials or configurations for any external service your command will be integrating with, like the Google Workspace Admin API in our example.

Understanding the BaseCommand Class

The BaseCommand class serves as the foundation for all SWAT commands. Beyond providing a consistent interface for commands, it offers:

Logging capabilities: This ensures all commands have access to the same logging setup and formats.
Argparse configurations: Allows commands to easily integrate with SWAT’s command-line interface.
(Any other functionality provided by the BaseCommand that might be relevant but hasn’t been highlighted).

Understanding Commands

BaseCommand Inheritance

All commands in SWAT inherit from the BaseCommand class in swat/base_command.py, allowing for a consistent interface and shared functionalities. This allows different SWAT commands and emulations to inherit the base SWAT object from base.py, as well as the logging capability, commands, and arguments.

Command Structure

A typical SWAT command class contains:

Command Class

Commands have a single Command(BaseCommand) class that inherits from BaseCommand the global SWAT object, logger handles, arguments, commands, and more.

Commands use argparse to define their command-line interface. In addition to the main parser, they may also define subparsers for different subcommands, with specific arguments and options for each:

class Command(BaseCommand):
    parser = BaseCommand.load_parser(description='Description here')
    subparsers = parser.add_subparsers(dest='subcommand', title='subcommands', required=True)

    parser_session = subparsers.add_parser('session', description='Description for session')
    # Add arguments and options for the 'session' subcommand

Initialization Method

Commands in SWAT generally require an initialization method, where they set up specific parser settings, defaults, and validate arguments:

def __init__(self, **kwargs) -> None:
    super().__init__(**kwargs)
    self.parser_session.set_defaults(func=self.authenticate)
    self.parser_list.set_defaults(func=self.list_sessions)

    self.args = validate_args(self.parser, self.args)

Execute Method

The required method that gets called when the command is executed:

def execute(self) -> None:
    # Main command logic

Dynamic Loading

Commands are dynamically loaded from the swat/commands/ directory, enabling SWAT to import and execute them at runtime.

Command Execution Workflow

Pre-Command Processing: Extract command name and arguments.
Loading Command: Dynamically load the command class.
Instantiate Command: Create an instance of the command class.
Execute Command: Call the command’s execute() method.
Post-Command Processing: Handle any necessary post-execution steps.

Error Handling and Logging

SWAT commands come with built-in logging capabilities, inherited from the BaseCommand class. When developing your command, consider the following:

Log meaningful information, especially potential errors or exceptions.
Always handle exceptions with try-except blocks to prevent the entire tool from crashing due to command errors.
Ensure logging levels are appropriately set (e.g., debug, info, warning, error).

Adding a New Command

SWAT is designed to simplify the process of adding new commands. However, to make things clearer, we’ll walk through an example of adding a command called users, which interacts with the Google Workspace Admin API to retrieve details about existing Google Workspace users.

Full command code example

from googleapiclient.discovery import build
from ..commands.base_command import BaseCommand
from ..misc import validate_args

class Command(BaseCommand):

    parser = BaseCommand.load_parser(description='Interact with Google Workspace Users.')
    subparsers = parser.add_subparsers(dest='subcommand', title='subcommands', required=True)

    parser_list = subparsers.add_parser('list', description='List all users in the domain')

    parser_view = subparsers.add_parser('view', description='View details of a specific user')
    parser_view.add_argument('user_id', help='The unique ID or email of the user')

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.parser_list.set_defaults(func=self.list_users)
        self.parser_view.set_defaults(func=self.view_user)

        self.args = validate_args(self.parser, self.args)

    def list_users(self):
        '''List all users in the Google Workspace domain.'''
        service = build('admin', 'directory_v1', credentials=self.obj.cred_store.store['default'].session)
        results = service.users().list(domain=self.obj.config['google']['domain']).execute()
        users = results.get('users', [])
        for user in users:
            self.logger.info(user['primaryEmail'])

    def view_user(self):
        '''View details of a specific user.'''
        service = build('admin', 'directory_v1', credentials=self.obj.credentials)
        user = service.users().get(userKey=self.args.user_id).execute()
        self.logger.info(user)

    def execute(self) -> None:
        '''Main execution method.'''
        self.args.func()

Example Walkthrough:

Create a new file: In the SWAT directory, create a file named users.py inside the swat/commands/ directory.
```
touch swat/commands/users.py
```
Add necessary imports: Ensure you have the following imports at the beginning of the users.py file. Remember that BaseCommand and validate_args are crucial for inheritance and argparse.
```
from googleapiclient.discovery import build
from ..commands.base_command import BaseCommand
from ..misc import validate_args
```
Define the Command Class: Start by defining the Command(BaseCommand) class.
```
class Command(BaseCommand):
    ...
```
Initialization Method: Add the __init__ method that inherits from BaseCommand and validates the arguments for the users command. You should also set up the self.service variable that facilitates access to the Admin API. For authentication details, refer to the How-To: Authenticate with OAuth section.
```
def __init__(self, **kwargs) -> None:
    ...
```
Execution Method: Add the execute() method, which essentially invokes self.args.func().
```
def execute(self) -> None:
    ...
```

Set up Argparse: Implement argparse subcommands and parsers. In this example, the users command has list and view as subcommands. Additionally, the view command requires the user_id argument.

parser = BaseCommand.load_parser(description='Interact with Google Workspace Users.')
subparsers = parser.add_subparsers(dest='subcommand', title='subcommands', required=True)
...

Implement Command Logic: Write methods that define the functionalities of the command.
```
def list_users(self):
    ...

def view_user(self):
    ...
```
Test Your Command: Start SWAT and execute the help command to ensure everything is working as expected.
```
help users
```