527 lines
16 KiB

import copy
import datetime
import importlib
import inspect
import json
import logging
import os
import string
import sys
import random
import urllib.parse
import re
import socket
import random
import struct
import actions.action
import actions.trigger
import layers.packet
import plugins.plugin_client
import plugins.plugin_server
from scapy.all import TCP, IP, UDP, rdpcap
import netifaces
RUN_DIRECTORY = os.path.join("trials","%Y-%m-%d_%H:%M:%S"))
# Hard coded options
FLAGFOLDER = "flags"
# Holds copy of console file handler's log level
BASEPATH = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(BASEPATH)
class SkipStrategyException(Exception):
Raised to signal that this strategy evaluation should be cut off.
def __init__(self, msg, fitness):
Creates the exception with the fitness to pass back
""" = fitness
self.msg = msg
def parse(requested_trees, logger):
Parses a string representation of a solution into its object form.
# First, strip off any hanging quotes at beginning/end of the strategy
if requested_trees.startswith("\""):
requested_trees = requested_trees[1:]
if requested_trees.endswith("\""):
requested_trees = requested_trees[:-1]
# Define a blank strategy to initialize with the user specified string
strat = actions.strategy.Strategy([], [])
# Actions for the in and out forest are separated by a "\/".
# Split the given string by this token
out_in_actions = requested_trees.split("\\/")
# Specify that we're starting with the out forest before we parse the in forest
out = True
direction = "out"
# For each string representation of the action directions, in or out
for str_actions in out_in_actions:
# Individual action trees always end in "|" to signify the end - split the
# entire action sequence into individual trees
str_actions = str_actions.split("|")
# For each string representation of each tree in the forest
for str_action in str_actions:
# If it's an empty action, skip it
if not str_action.strip():
assert " " not in str_action.strip(), "Strategy includes a space - malformed!"
# Get rid of hanging whitespace from the splitting
str_action = str_action.strip()
# ActionTree uses the last "|" as a sanity check for well-formed
# strategies, so restore the "|" that was lost from the split
str_action = str_action + "|"
new_tree = actions.tree.ActionTree(direction)
success = new_tree.parse(str_action, logger)
if success is False:
raise actions.tree.ActionTreeParseError("Failed to parse tree")
# Once all the actions are parsed, add this tree to the
# current direction of actions
if out:
# Change the flag to tell it to parse the IN direction during the next loop iteration
out = False
direction = "in"
return strat
def get_logger(basepath, log_dir, logger_name, log_name, environment_id, log_level="DEBUG", file_log_level="DEBUG", demo_mode=False):
Configures and returns a logger.
if type(log_level) == str:
log_level = log_level.upper()
if type(file_log_level) == str:
file_log_level = file_log_level.upper()
full_path = os.path.join(basepath, log_dir, "logs")
if not os.path.exists(full_path):
flag_path = os.path.join(basepath, log_dir, "flags")
if not os.path.exists(flag_path):
# Set up a client logger
logger = logging.getLogger(logger_name + environment_id)
# Disable the root logger to avoid double printing
logger.propagate = False
# If we've already setup the handlers for this logger, just return it
if logger.handlers:
return logger
fh = logging.FileHandler(os.path.join(basepath, log_dir, "logs", "%s.%s.log" % (environment_id, log_name)))
log_prefix = "[%s] " % log_name.upper()
formatter = logging.Formatter("%(asctime)s %(levelname)s:" + log_prefix + "%(message)s", datefmt="%Y-%m-%d %H:%M:%S")
file_formatter = logging.Formatter(log_prefix + "%(asctime)s %(message)s")
ch = logging.StreamHandler()
CONSOLE_LOG_LEVEL = log_level.lower()
return CustomAdapter(logger, {}) if demo_mode else logger
class CustomAdapter(logging.LoggerAdapter):
Used for demo mode, to change sensitive IP addresses where necessary. Can be used (mostly) like a regular logger.
regex = re.compile(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}")
def __init__(self, logger, extras):
super().__init__(logger, extras)
self.handlers = logger.handlers
self.ips = {}
def debug(self, msg, *args, **kwargs):
Print a debug message, uses logger.debug.
msg, args, kwargs = self.process(msg, args, kwargs)
self.logger.debug(msg, *args, **kwargs)
def info(self, msg, *args, **kwargs):
Print an info message, uses
msg, args, kwargs = self.process(msg, args, kwargs), *args, **kwargs)
def warning(self, msg, *args, **kwargs):
Print a warning message, uses logger.warning.
msg, args, kwargs = self.process(msg, args, kwargs)
self.logger.warning(msg, *args, **kwargs)
def error(self, msg, *args, **kwargs):
Print an error message, uses logger.error.
msg, args, kwargs = self.process(msg, args, kwargs)
self.logger.error(msg, *args, **kwargs)
def critical(self, msg, *args, **kwargs):
Print a critical message, uses logger.critical.
msg, args, kwargs = self.process(msg, args, kwargs)
self.logger.critical(msg, *args, **kwargs)
def get_ip(self, ip):
Lookup the assigned random IP for a given real IP.
If no random IP exists, a new one is created and a message is logged indicating it.
if ip not in self.ips:
random_ip = socket.inet_ntoa(struct.pack('>I', random.randint(1, 0xffffffff)))"Registering new random IP: %s" % random_ip)
self.ips[ip] = random_ip
def process(self, msg, args, kwargs):
Modify the log message to replace any instance of an IP in msg or args with its assigned random IP.
new_args = []
for arg in args:
if type(arg) == str:
for ip in self.regex.findall(arg):
new_ip = self.get_ip(ip)
arg = arg.replace(ip, self.ips[ip])
for ip in self.regex.findall(msg):
if ip not in self.ips:
random_ip = socket.inet_ntoa(struct.pack('>I', random.randint(1, 0xffffffff)))
self.logger.debug("Registering new random IP: %s" % random_ip)
self.ips[ip] = random_ip
new_ip = self.get_ip(ip)
msg = msg.replace(ip, self.ips[ip])
return msg, tuple(new_args), kwargs
def close_logger(logger):
Closes open file handles for a given logger.
# Close the file handles so we don't hold a ton of file descriptors open
handlers = logger.handlers[:]
for handler in handlers:
if isinstance(handler, logging.FileHandler):
class Logger():
Logging class context manager, as a thin wrapper around the logging class to help
handle closing open file descriptors.
def __init__(self, log_dir, logger_name, log_name, environment_id, log_level="DEBUG"):
self.log_dir = log_dir
self.logger_name = logger_name
self.log_name = log_name
self.environment_id = environment_id
self.log_level = log_level
self.logger = None
def __enter__(self):
Sets up a logger.
self.logger = get_logger(PROJECT_ROOT, self.log_dir, self.logger_name, self.log_name, self.environment_id, log_level=self.log_level)
return self.logger
def __exit__(self, exc_type, exc_value, tb):
Closes file handles.
def get_console_log_level():
returns log level of console handler
def get_plugins():
Iterates over this current directory to retrieve plugins.
plugins = []
for f in os.listdir(os.path.join(PROJECT_ROOT, "plugins")):
if os.path.isdir(os.path.join(PROJECT_ROOT, "plugins", f)) and "__pycache__" not in f:
return plugins
def import_plugin(plugin, side):
Imports given plugin.
- plugin: plugin to import (e.g. "http")
- side: which side of the connection should be imported ("client" or "server")
# Define the full module for this plugin
mod = "plugins.%s.%s" % (plugin, side)
path = os.path.join(PROJECT_ROOT, "plugins", plugin)
if path not in sys.path:
# Import the module
# Predicate to filter classmembers
def check_plugin(obj):
Filters class members to ensure we get only enabled Plugin subclasses
return inspect.isclass(obj) and \
issubclass(obj, plugins.plugin.Plugin) and \
(obj != plugins.plugin_client.ClientPlugin and \
obj != plugins.plugin_server.ServerPlugin and \
obj != plugins.plugin.Plugin) and \
# Filter the class members of the imported module to find our Plugin subclass
clsmembers = inspect.getmembers(sys.modules[mod], predicate=check_plugin)
# Sanity check the class members we identified
assert clsmembers, "Could not find plugin %s" % mod
assert len(clsmembers) == 1, "Too many matching plugins found for %s" % mod
# Extract the class - clsmembers[0] is a tuple of (name, class)
_, cls = clsmembers[0]
# Return the module path and class
return mod, cls
def build_command(args):
Given a dictionary of arguments, build it back into a command line string.
cmd = []
for opt in args:
# Don't pass along store true args that are false
if args[opt] in [False, None]:
cmd.append("--%s" % opt.replace("_", "-"))
# If store true arg, we don't need to pass the value
if args[opt] is True:
if args[opt] is '':
elif " " in str(args[opt]):
cmd.append("\"" + str(args[opt]) + "\"")
return cmd
def string_to_protocol(protocol):
Converts string representations of scapy protocol objects to
their actual objects. For example, "TCP" to the scapy TCP object.
if protocol.upper() == "TCP":
return TCP
elif protocol.upper() == "IP":
return IP
elif protocol.upper() == "UDP":
return UDP
def get_id():
Returns a random ID
return ''.join([random.choice(string.ascii_lowercase + string.digits) for k in range(8)])
def setup_dirs(output_dir):
Sets up Geneva folder structure.
ga_log_dir = os.path.join(output_dir, "logs")
ga_flags_dir = os.path.join(output_dir, "flags")
ga_packets_dir = os.path.join(output_dir, "packets")
ga_generations_dir = os.path.join(output_dir, "generations")
ga_data_dir = os.path.join(output_dir, "data")
for directory in [ga_log_dir, ga_flags_dir, ga_packets_dir, ga_generations_dir, ga_data_dir]:
if not os.path.exists(directory):
os.makedirs(directory, exist_ok=True)
return ga_log_dir
def get_from_fuzzed_or_real_packet(environment_id, real_packet_probability, enable_options=True, enable_load=True):
Retrieves a protocol, field, and value from a fuzzed or real packet, depending on
the given probability and if given packets is not None.
packets = actions.utils.read_packets(environment_id)
if packets and random.random() < real_packet_probability:
packet = random.choice(packets)
return packet.get_random()
return layers.packet.Packet().gen_random()
def read_packets(environment_id):
Reads the pcap file associated with the last evaluation of this strategy.
Returns a list of Geneva Packet objects.
if not environment_id:
return None
packets_path = os.path.join(RUN_DIRECTORY, "packets", "original_" + str(environment_id) + ".pcap")
if not os.path.exists(packets_path):
return None
parsed = []
packets = rdpcap(packets_path)
parsed = [layers.packet.Packet(p) for p in packets]
except Exception as e:
return parsed
def punish_fitness(fitness, logger, eng):
Adjusts fitness based on additional optimizer functions.
if not eng:
logger.warning("Requested fitness adjustment without an engine - returning original fitness.")
return fitness
logger.debug("Initiating fitness adjustment")
if eng and eng.strategy:
fitness = punish_complexity(fitness, logger, eng.strategy)
fitness = punish_unused(fitness, logger, eng.strategy)
if fitness > 0:
overhead = int(eng.overhead / 2)
logger.debug("Punishing for overhead: %d" % overhead)
fitness -= overhead
return fitness
def punish_unused(fitness, logger, ind):
Punishes strategy for each action that was not run.
if not ind:
return fitness
logger.debug("Punishing for unused actions")
num_unused = [action_tree.ran for action_tree in ind.out_actions].count(False)
fitness -= (num_unused * 10)
logger.debug(" - Number of unused actions in out forest: %d" % num_unused)
num_unused = [action_tree.ran for action_tree in ind.in_actions].count(False)
fitness -= (num_unused * 10)
logger.debug(" - Number of unused actions in in forest: %d" % num_unused)
return fitness
def punish_complexity(fitness, logger, ind):
Reduces fitness based on number of actions - optimizes for simplicity.
if not ind:
return fitness
# Punish for number of actions
if fitness > 0:
logger.debug("Punishing for complexity: %d" % len(ind))
fitness -= len(ind)
return fitness
def write_fitness(fitness, output_path, eid):
Writes fitness to disk.
except ValueError:
print("Given fitness (%r) is not a number!" % fitness)
fitpath = os.path.join(PROJECT_ROOT, output_path, FLAGFOLDER, eid) + ".fitness"
with open(fitpath, "w") as fitfile:
def get_interface():
Chooses an interface on the machine to use for socket testing.
ifaces = netifaces.interfaces()
for iface in ifaces:
if "lo" in iface:
info = netifaces.ifaddresses(iface)
# Filter for IPv4 addresses
if netifaces.AF_INET in info:
return iface
def get_worker(name, logger):
Returns information dictionary about a worker given its name.
path = os.path.join("workers", name, "worker.json")
if os.path.exists(name):
path = name
dirpath = os.path.dirname(path)
if not os.path.exists(path):
return None
with open(path, "r") as fd:
data = json.load(fd)
# If there is a private key, update the path to be relative to the project base
if data.get("keyfile"):
data["keyfile"] = os.path.join(dirpath, data["keyfile"])
return data