geneva/actions/strategy.py

230 lines
7.7 KiB
Python

import random
import actions.utils
import actions.tree
class Strategy(object):
def __init__(self, in_actions, out_actions, environment_id=None):
self.in_actions = in_actions
self.out_actions = out_actions
self.descendents = []
self.in_enabled = True
self.out_enabled = True
self.environment_id = environment_id
self.fitness = -1000
def __str__(self):
"""
Builds a string describing the action trees for this strategy.
"""
return "%s \/ %s" % (self.str_forest(self.out_actions).strip(), self.str_forest(self.in_actions).strip())
def __len__(self):
"""
Returns the number of actions in this strategy.
"""
num = 0
for tree in self.in_actions:
num += len(tree)
for tree in self.out_actions:
num += len(tree)
return num
def str_forest(self, forest):
"""
Returns a string representation of a given forest (inbound or outbound)
"""
rep = ""
for action_tree in forest:
rep += "%s " % str(action_tree)
return rep
def pretty_print(self):
return "%s \n \/ \n %s" % (self.pretty_str_forest(self.out_actions), self.pretty_str_forest(self.in_actions))
def pretty_str_forest(self, forest):
"""
Returns a string representation of a given forest (inbound or outbound)
"""
rep = ""
for action_tree in forest:
rep += "%s\n" % action_tree.pretty_print()
return rep
def initialize(self, logger, num_in_trees, num_out_trees, num_in_actions, num_out_actions, seed, disabled=None):
"""
Initializes a new strategy object randomly.
"""
# Disable specific forests if none requested
if num_in_trees == 0:
self.in_enabled = False
if num_out_trees == 0:
self.out_enabled = False
# If a specific population seed is requested, build using that
if seed:
starting_strat = actions.utils.parse(seed, logger)
self.out_actions = starting_strat.out_actions
self.in_actions = starting_strat.in_actions
return self
self.init_from_scratch(num_in_trees, num_out_trees, num_in_actions, num_out_actions, disabled=disabled)
return self
def init_from_scratch(self, num_in_trees, num_out_trees, num_in_actions, num_out_actions, disabled=None):
"""
Initializes this individual by drawing random actions.
"""
for _ in range(0, num_in_trees):
# Define a new in action tree
in_tree = actions.tree.ActionTree("in")
# Initialize the in tree
in_tree.initialize(num_in_actions, self.environment_id, disabled=disabled)
# Add them to this strategy
self.in_actions.append(in_tree)
for _ in range(0, num_out_trees):
# Define a new out action tree
out_tree = actions.tree.ActionTree("out")
# Initialize the out tree
out_tree.initialize(num_out_actions, self.environment_id, disabled=disabled)
# Add them to this strategy
self.out_actions.append(out_tree)
def act_on_packet(self, packet, logger, direction="out"):
"""
Runs the strategy on a given scapy packet.
"""
# If there are no actions to run for this strategy, just send the packet
if (direction == "out" and not self.out_actions) or \
(direction == "in" and not self.in_actions):
return [packet]
return self.run_on_packet(packet, logger, direction)
def run_on_packet(self, packet, logger, direction):
"""
Runs the strategy on a given packet given the forest direction.
"""
forest = self.out_actions
if direction == "in":
forest = self.in_actions
ran = False
original_packet = packet.copy()
packets_to_send = []
for action_tree in forest:
if action_tree.check(original_packet, logger):
logger.debug(" + %s action tree triggered: %s", direction, str(action_tree))
# If multiple trees run, the previous packet may have been tampered with. Ensure
# we're always acting on a fresh copy
fresh_packet = original_packet.copy()
packets_to_send += action_tree.run(fresh_packet, logger)
ran = True
# If no action tree was applicable, send the packet unimpeded
if not ran:
packets_to_send = [packet]
return packets_to_send
def mutate_dir(self, trees, direction, logger):
"""
Mutates a list of trees. Requires the direction the tree operates on
(in or out).
"""
pick = random.uniform(0, 1)
if pick < 0.1 or not trees:
new_tree = actions.tree.ActionTree(direction)
new_tree.initialize(1, self.environment_id)
trees.append(new_tree)
elif pick < 0.2 and trees:
trees.remove(random.choice(trees))
elif pick < 0.25 and trees and len(trees) > 1:
random.shuffle(trees)
else:
for action_tree in trees:
action_tree.mutate()
def mutate(self, logger):
"""
Top level mutation function for a strategy. Simply mutates the out
and in trees.
"""
if self.in_enabled:
self.mutate_dir(self.in_actions, "in", logger)
if self.out_enabled:
self.mutate_dir(self.out_actions, "out", logger)
return self
def swap_one(forest1, forest2):
"""
Swaps a random tree from forest1 and forest2.
It picks a random element within forest1 and a random element within forest2,
chooses a random index within each forest, and inserts the random element
"""
assert type(forest1) == list
assert type(forest2) == list
rand_idx1, rand_idx2 = 0, 0
donation, other_donation = None, None
if forest1:
donation = random.choice(forest1)
forest1.remove(donation)
if len(forest1) > 0:
rand_idx1 = random.choice(list(range(0, len(forest1))))
if forest2:
other_donation = random.choice(forest2)
forest2.remove(other_donation)
if len(forest2) > 0:
rand_idx2 = random.choice(list(range(0, len(forest2))))
if other_donation:
forest1.insert(rand_idx1, other_donation)
if donation:
forest2.insert(rand_idx2, donation)
return True
def do_mate(forest1, forest2):
"""
Performs mating between two given forests (lists of trees).
With 80% probability, a random tree from each forest are mated,
otherwise, a random tree is swapped between them.
"""
# If 80% and there are trees in both forests to mate, or
# if there is only 1 tree in each forest, mate those trees
if (random.random() < 0.8 and forest1 and forest2) or \
(len(forest1) == 1 and len(forest2) == 1):
tree1 = random.choice(forest1)
tree2 = random.choice(forest2)
return tree1.mate(tree2)
# Otherwise, swap a random tree from each forest
elif forest1 or forest2:
return swap_one(forest1, forest2)
return False
def mate(ind1, ind2, indpb):
"""
Executes a uniform crossover that modify in place the two
individuals. The attributes are swapped according to the
*indpb* probability.
"""
out_success, in_success = True, True
if ind1.out_enabled and random.random() < indpb:
out_success = do_mate(ind1.out_actions, ind2.out_actions)
if ind1.in_enabled and random.random() < indpb:
in_success = do_mate(ind1.in_actions, ind2.in_actions)
return out_success and in_success