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