-
Notifications
You must be signed in to change notification settings - Fork 72
Creating a Module
Armory's usefulness comes from being able to quickly code up a module for a tool and ingest the output for useful data. This data can then be inserted into the database, and used with other tools or reports. Right now, modules are based off of templates in included/ModuleTemplate.py. The two templates are as follows:
- ModuleTemplate: Basic template with almost nothing built in. This is what you would base any "out-of-the-box" modules.
- ToolTemplate: Template designed to allow one to easily create a multithreaded module for an arbitrary tool, and process the output.
As an example, we'll build a module for a tko-subs, a tool used to check various domains to see if they are vulnerable to hijacking. (The tool can be found here).
For the first step, we'll copy the skeleton for a tool template from the "SampleToolModule.py" file into a new "Tko-subs.py" file.
cd included/modules
cp SampleToolModule.py Tko-subs.py
At this point, the file Tko-subs.py contains the following code:
#!/usr/bin/python
from included.ModuleTemplate import ToolTemplate
class Module(ToolTemplate):
'''
This is a sample skeleton for building a module to run a tool.
'''
name = "SampleToolModule"
binary_name = "sample-tool"
def set_options(self):
super(Module, self).set_options()
self.options.add_argument('-p', '--print_message', help="Message to print")
def get_targets(self, args):
'''
This module is used to build out a target list and output file list, depending on the arguments. Should return a
list in the format [(target, output), (target, output), etc, etc]
'''
return []
def build_cmd(self, args):
'''
Create the actual command that will be executed. Use {target} and {output} as placeholders.
'''
return ''
def process_output(self, cmds):
'''
Process the output generated by the earlier commands.
'''
As the very first thing, we'll update the name field to reflect the module, as well as update the binary_name. This name will be searched for in the path if the --binary flag is not provided. We'll also update the help text.
class Module(ToolTemplate):
'''
This tool will check various subdomains for domain takeover vulnerabilities.
'''
name = "Tko-subs"
binary_name = "tko-subs"
The first function set_options is used to add any additional options that will be used with the tool. These are being configured via ArgParse's add_argument function, and will also show up in the system help. The template has the following options predefined:
self.options.add_argument('-b', '--binary', help="Path to the binary")
self.options.add_argument('-o', '--output_path', help="Relative path (to the base directory) to store Fierce output", default="fierceFiles")
self.options.add_argument('--threads', help="Number of Armory threads to use", default="10")
self.options.add_argument('--timeout', help="Thread timeout in seconds, default is 300.", default="300")
self.options.add_argument('--extra_args', help="Additional arguments to be passed to the tool")
self.options.add_argument('--no_binary', help="Runs through without actually running the binary. Useful for if you already ran the tool and just want to process the output.", action="store_true")
This tool has the following options available:
Usage of ./tko-subs:
-data string
CSV file containing CMS providers' string for identification (default "providers-data.csv")
-domain string
Domains separated by ,
-domains string
List of domains to check (default "domains.txt")
-githubtoken string
Github personal access token
-herokuapikey string
Heroku API key
-herokuappname string
Heroku app name
-herokuusername string
Heroku username
-output string
Output file to save the results (default "output.csv")
-takeover
Flag to denote if a vulnerable domain needs to be taken over or not
-threads int
Number of threads to run parallel (default 5)
For our module, we want to take direct control over the data, domain, domains and output arguments. The data argument will point to the "providers-data.csv" file needed by the tool. The domain and domains arguments will be used to provide the targets to the tool, and the output argument specifies where to save the output. We'll modify the set_options function to include these.
def set_options(self):
super(Module, self).set_options()
self.options.add_argument('--data', help="Path to the providers_data.csv file")
self.options.add_argument('-d', '--domain', help="Domain to run the tool against.")
self.options.add_argument('-i', '--import', help="Import subdomains from the database.", action="store_true")
The next function we need to modify is the get_targets function. This gets the parsed arguments passed to it, and returns a list of tuples containing the targets and output files. We will build a target list based on the domain names either passed with "-d", or import them directly from the database.
def get_targets(self, args):
'''
This module is used to build out a target list and output file list, depending on the arguments. Should return a
list in the format [(target, output), (target, output), etc, etc]
'''
# Create an empty list to add targets to
domains = []
# Check if a domain has been explicitly passed, and if so add it to targets
if args.domain:
domains.append(args.domain)
# Check if the --import option was passed and if so get data from the domain.
# The scoping is set to "active" since the tool could potentially make a
# request from the server.
if args.i:
all_domains = DomainRepository.all(scoping="active")
for d in all_domains:
domains.append(d.domain)
# Set the output_path base, as a junction of the base_path and the path name supplied
if args.output_path[0] == "/":
output_path = os.path.join(self.base_config['PROJECT']['base_path'], args.output_path[1:] )
else:
output_path = os.path.join(self.base_config['PROJECT']['base_path'], args.output_path)
# Create the path if it doesn't already exist
if not os.path.exists(output_path):
os.makedirs(output_path)
res = []
# Create the final targets list, with output paths added.
for d in domains:
res.append((d, os.path.join(output_path, d)))
return res
We've also added another import for the DomainRepository class, as well as an __init__ function to set up the database.
from database.repositories import DomainRepository
import os
def __init__(self, db):
self.db = db
self.Domain = DomainRepository(db, self.name)
Next we set up the build_cmd function. This takes the parsed arguments and builds out the actual command template that will be used, with format placeholders for the target and output filenames. This returns a string containing both {target} and {output}.
def build_cmd(self, args):
'''
Create the actual command that will be executed. Use {target} and {output} as placeholders.
'''
cmd = self.binary + " -domain {target} -output {output} "
# Add in any extra arguments passed in the extra_args parameter
if args.extra_args:
cmd += args.extra_args
# Add that data parameter in there
cmd += " -data " + args.data
return cmd