-
Notifications
You must be signed in to change notification settings - Fork 6
Scales
Scales are a core part of Snake and are used to extend and enrich Snake's capability. For these reasons a careful amout time and consideration has been invested to ensure that scales are easy to write while also being powerful. This page explains how each of the components works and provides examples. The folder and file structure for a feature complete scale is shown below:
<SCALE_ROOT>
|- <SCALE_NAME>
| |- __init__.py
| |- commands.py
| |- interface.py
| |- upload.py
| `- <SCALE_NAME>.conf
`- setup.py
A brief description for each of the above is as follows:
-
scale_root- contains everything required to create a scale and install it, but should not include 3rd party data such as binaries.-
scale_name- the name of the scale and the actual scale module.-
__init__.py- the metadata for the scale and used to instantiate it. -
commands.py- the command component used to execute commands on a sample. -
interface.py- the interface component used to interact with external services in accordance with the sample in question.
-
-
upload.py- the upload component used to extend the supported upload methods provided by Snake. -
scale_name.conf- the scale's configuration file used to provide the scale with configurable settings.
-
-
setup.py- part of setuptools used to enable easy installation of the scale.
This is the only mandatory part of a scale and is used to provide essential information about it, such as: name, version, author, etc.
Note: While a scale with just an
__init__.pyfile is valid, without any components it will offer no functionality.
Below is an example for __init__.py:
from snake.scale import scale
NAME = "plane"
VERSION = "1.0"
AUTHOR = "Samuel L. Jackson"
AUTHOR_EMAIL = "snakes@plane.air"
DESCRIPTION = "a demo scale"
LICENSE = ""
URL = ""
__scale__ = scale(
name=NAME,
description=DESCRIPTION,
version=VERSION,
author=AUTHOR,
supports=[ # Used to restrict the scale to certain sample file types
],
)To elaborate on the supports field, this is used to prevent Snake from enabling scales for incompatible samples.
There are two sample types in Snake: files & memories.
Restricting on these types provides enough granularity to restrict scales sufficiently while not hindering the analyst.
For example, the volatility scale is only useful to memories hence the scale could be restricted as follows:
from snake.scale import scale, FileType
[snip]
__scale__ = scale(
[snip]
supports=[
FileType.MEMORY
]
)This file is used to provide a scale and its components with user configurable variables.
This is increasingly useful when a scale requires bespoke settings as sane defaults would not be sufficient, such as the path to YARA rules for the YARA scale.
When creating the configuration file for a scale it must be named <SCALE_NAME.conf>, written in YAML and reside in the scale's folder.
Within a scale it is accessed like so:
from snake import config
config.snake_config['VAR'] # Accessing Snake's global variables
config.scale_configs['SCALE_NAME']['VAR'] # Accessing the scale's variables from its .confWhile not mandatory to scale creation, this file is highly recommended as it will easy scale installation and enable pip support.
Below is a template for setup.py where in most cases amending the <KEYWORDS> should be sufficient.
from setuptools import setup
import <SCALE_NAME> as scale
setup(
name="snake-{}".format(scale.NAME),
version=scale.VERSION,
packages=[
"snake_{}".format(scale.NAME)
],
package_dir={
"snake_{}".format(scale.NAME): scale.NAME
},
install_requires=[
"snake",
<ANY_ADDITIONAL_PYTHON_REQUIREMENTS>
],
entry_points={
"snake.scales": [
"{0} = snake_{0}".format(scale.NAME),
]
},
include_package_data=True,
zip_safe=False,
author=scale.AUTHOR,
author_email=scale.AUTHOR_EMAIL,
description=scale.DESCRIPTION,
license=scale.LICENSE,
url=scale.URL
)The commands component is designed for one thing and one thing only, to execute commands on a sample.
For a scale to provide command functionality it must contain the commands.py file and subclass the Commands class.
Any class that subclasses the Commands class must implement the abstract function check.
This function is used by Snake to perform sanity checks (related to the scale) before loading the scale or executing any of its commands.
For Snake to know which of the functions in the subclass are commands the command decorator must be used and the function prototype adhered.
A command function must return output that is JSON serialisable and is stored in Snake's database.
The prototype for a command function must look like the following:
# args (dict) - the arguments passed by the user as described in the `args` field of the `command` decorator
# file (FileStorage) - the `FileStorage` object, which provides the information required to execute a command on a scale, such as its path
# opts (dict) - the dictionary passed to the `command` decorator
def func(self, args, file, opts)The command decorator marks a function as a command:
from snake import fields
from snake import scale
# args (dict) - any additional arguments required by the command, these could be mandatory or optional
# info (str) - a description for the command
# mime (str) - restrict the command to certain MIME types (currently only used by `autorun` decorator)
@scale.command({
'args': {
'rule': fields.Str(required=False)
},
'info': 'description about the command',
'mime': 'application/x-dosexec'
})
def a_func(self, args, file, opts)The autorun decorator is used to tell snake that a command should be run automatically on a newly uploaded sample.
from snake import scale
@scale.autorun
@scale.command({
'info': 'calculates the sha512 hash for the file'
})
def a_func(self, args, file, opts):With the above in mind, below is an example which would allow a scale to provide Snake with the ability to calculate the SHA512 digest of a sample.
import hashlib
from snake import scale
class Commands(scale.Commands):
def check(self):
pass # No checks required, this is all python...
@scale.command({
'info': 'calculates the sha512 hash for the file'
})
def sha512_digest(self, args, file, opts):
sha512_hash = hashlib.sha512()
with open(file.file_path, "rb") as f:
for chunk in iter(lambda: f.read(4096), b""):
sha512_hash.update(chunk)
sha512_digest = sha512_hash.hexdigest()
return {'sha512_digest': sha512_digest}The interface component is used to enable Snake to query other services about its samples.
For a scale to provide interface functionality it must contain the interface.py file and subclass the Interface class.
Any class that subclasses the Interface class must implement the abstract function check.
This function is used by Snake to perform sanity checks (related to the scale) before loading the scale or interacting with a service.
For Snake to know which of the functions in the subclass are commands for interacting with a service the pull/push decorators must be used and the function prototypes adhered.
An interface command function must return output that is JSON serialisable.
The prototype for a command function must look like the following:
# args (dict) - the arguments passed by the user as described in the `args` field of the `command` decorator
# file (FileStorage) - the `FileStorage` object, which provides the information required to execute a command on a scale, such as its path
# opts (dict) - the dictionary passed to the `command` decorator
def func(self, args, file, opts)The pull decorator is used to get information from a service.
from snake import scale
from snake import fields
# args (dict) - any additional arguments required by the command, these could be mandatory or optional
# info (str) - a description for the command
# mime (str) - restrict the command to certain MIME types (currently unused)
@scale.pull({
'args': {
'cache': fields.Bool(required=False)
},
'info': 'description about the command',
'mime': 'application/x-dosexec'
})
def a_func(self, args, file, opts):The push decorator is used to post information to a service.
from snake import scale
from snake import fields
# args (dict) - any additional arguments required by the command, these could be mandatory or optional
# info (str) - a description for the command
# mime (str) - restrict the command to certain MIME types (currently unused)
@scale.push({
'args': {
'machine': fields.Str(required=False)
},
'info': 'description about the command',
'mime': 'application/x-dosexec'
})
def a_func(self, args, file, opts):With the above in mind, below is an example which would allow a scale to provide Snake with the ability to retrieve information about a sample from an imaginary service.
import requests
from snake import config
from snake import error
from snake import scale
DEMO_API = config.scale_configs['demo']['demo_api']
HEADERS = {
"User-Agent": config.constants.USER_AGENT
}
class Interface(scale.Interface):
def check(self):
if DEMO_API is None or DEMO_API == '':
raise error.InterfaceError("config variable 'demo_api' has not been set")
@scale.pull({
'info': 'a demo pull command'
})
def info(self, args, file, opts):
try:
resp = requests.get(DEMO_API + '/' + file.sha256_digest + '/info', headers=HEADERS)
except requests.exceptions.RequestException:
raise error.InterfaceError("failed to connect to demo")
return resp.json()The upload component is used to enable Snake to acquire samples from a variety of methods.
For a scale to provide interface functionality it must contain the upload.py file and subclass the Upload class.
Unlike the other components an upload component can only provide one upload function and therefore does no use decorators. Instead, three functions must be implemented.
Any additional arguments required by the upload function, these could be mandatory or optional.
def arguments(self)A description for the upload function.
def info(self)The function to acquire the file for snake.
It must return the name of the file placed in the working_dir.
# args (dict) - the arguments passed by the user as described in the `args` field of the `pull` or `push` decorator.
# working_dir (str) - the directory to place the acquired file in, the name of the file is returned by the upload function.
def upload(self, args, working_dir)With the above in mind, below is an example which would allow a scale to provide Snake with the ability acquire files from a given URL.
Note: Taken from the URL scale.
from os import path
from urllib import parse
import cgi
import requests
from snake import config
from snake import error
from snake import fields
from snake import scale
PROXIES = {}
if config.snake_config['http_proxy']:
PROXIES['http'] = config.snake_config['http_proxy']
if config.snake_config['https_proxy']:
PROXIES['https'] = config.snake_config['https_proxy']
HEADERS = {
"Accept-Encoding": "gzip, deflate",
"User-Agent": config.constants.USER_AGENT
}
class Upload(scale.Upload):
def arguments(self):
return {
'url': fields.Str(required=True)
}
def info(self):
return "fetches files from arbitrary URLs and uploads them to Snake"
def upload(self, args, working_dir):
url_parser = parse.urlparse(args['url'])
if not url_parser.scheme:
url_parser = url_parser('http://' + args['url'])
req = requests.get(url_parser.geturl(),
headers=HEADERS,
proxies=PROXIES,
stream=True,
timeout=300)
if not req.status_code == requests.codes.ok: # pylint: disable=no-member
raise error.UploadError('HTTP Error: %s - %s' % (req.status_code, req.reason))
name = None
if 'Content-Disposition' in req.headers:
_, params = cgi.parse_header(req.headers['Content-Disposition'])
if 'filename' in params:
name = params['filename']
if not name:
name = args['url'].split('/')[-1]
with open(path.join(working_dir, name), 'wb') as f:
for chunk in req.iter_content(chunk_size=4096):
if chunk:
f.write(chunk)
return nameAdditional information about scales.
After reading the above it would be logical to conclude that scales returning output only support JSON, but this would be incorrect. Any function that has the requirement to output in JSON can be reformatted using helper functions. These helper functions expect JSON input and are expected to reformat it to the desired format.
Note: Currently supported formats: JSON, Markdown, & Plaintext.
For Snake to accept the helper function to a command a couple of conditions must be met:
- the helper function must be placed in the same subclass as the command function it is to help.
- the helper function's name must adhere to the format:
<FUNCTION_NAME>_<FORMAT>.
So, to provide the sha512_digest function above with plaintext support the following would need to be done:
[snip]
@staticmethod
def sha512_digest_plaintext(json):
return json['sha512_digest']
[snip]