mirror of https://github.com/Kkevsterrr/geneva
230 lines
7.7 KiB
Python
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
|