Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor out metadata #12

Merged
merged 3 commits into from
May 9, 2021
Merged

Refactor out metadata #12

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c509258
Refactored endNode to childNode now that Edge supports lineage throug…
hyakuhei May 8, 2021
507c7fc
Addressing brain and models
hyakuhei May 8, 2021
4641f16
Refactored out metadata to just use class variables
hyakuhei May 9, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Refactored out metadata to just use class variables
  • Loading branch information
@hyakuhei
hyakuhei committed May 9, 2021
commit 4641f16f14b34b517defc4b5ac42de5704925142
18 changes: 7 additions & 11 deletions attacktree/brain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,6 @@ class Brain(object):
def __init__(self):
self.exploitChain = []

# Walk throught the nodes of tree and print if metadata is missing
def lint(self, root: Node, walked: dict={}):
pass
# Walk and print metadata

# Walk the tree, adding to the chain (DFS)
# If we hit the goal, add that chain to our paths
def pathsToVictory(self,
Expand Down Expand Up @@ -49,22 +44,23 @@ def pathsToVictory(self,

return paths

# Walk the given path, add up stats and annotate edges
def evaluatePath(self, path):
#TODO: Replace concrete numbers with ranges and confidence intervals.
# It's not the nodes we need to evaluate, it's the edges. As those are what get changed adding a block
results = {}
for key in rules: #Pre-load data from rules
results[key] = rules[key]['startWith']

prevNode = None
for node in path:
#TODO: Introduce pDiscovery value (or pSuccess on Discovery() )
if isinstance(node, (Action)):
results['attackCost'] += node.metadata['cost']
results['time'] += node.metadata['time']
results['pSuccess'] = int((results['pSuccess'] * node.metadata['pSuccess']) / 100)
results['attackCost'] += node.cost
results['time'] += node.time
results['pSuccess'] = int((results['pSuccess'] * node.pSuccess) / 100)
if isinstance(node, (Block)):
results['defenceCost'] += node.metadata['cost']
results['pSuccess'] -= node.metadata['pDefend']
results['defenceCost'] += node.cost
results['pSuccess'] -= node.pDefend
#TODO block time

if prevNode is not None:
Expand Down
56 changes: 24 additions & 32 deletions attacktree/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,8 @@ class Node(object):
def __init__(self, label="Anonymous", metadata={}):
self.label = label
self.uniq = uuid.uuid4().hex #TODO: Remove this it's not needed, it's kinda here to make rendering work
self.edges = []
self.metadata = {}
self.edges = []
self.parentEdges = [] # backref

#Backref means we don't actually create a real edge, we just maintain a list of backward references that we can draw in later.
Expand Down Expand Up @@ -163,13 +163,11 @@ class Root(Node):
def __init__(self,
label: str):
super().__init__(label=label)
self.metadata = {}

class Goal(Node):
def __init__(self,
label: str):
super().__init__(label=label)
self.metadata = {}

# label: 'The name of the node',
# chain: 'The stage of the Mitre Att&ck chain represented, e.g "recon"'
Expand All @@ -187,13 +185,11 @@ def __init__(self,
pSuccess: int = 100,
detections: list = []):
super().__init__(label=label)
self.metadata = {}
self.metadata['chain'] = chain
self.metadata['cost'] = cost
self.metadata['time'] = time
self.metadata['pSuccess'] = pSuccess
self.metadata['detections'] = detections
self.edges = []
self.pSuccess = pSuccess
self.chain = chain
self.cost = cost
self.time = time
# self.metadata = [] ## I don't have a need for this right now, so removing it,but I expect to need it later.

class Detect(Node):
def __init__(self,
Expand All @@ -203,16 +199,14 @@ def __init__(self,
description: str = "",
complexity: int = 0,
latency: int = 0,
pSuccess: int = 100):
pDetect: int = 100):
super().__init__(label=label)
self.metadata = {}
self.metadata['cost'] = cost
self.metadata['description'] = description
self.metadata['complexity'] = complexity
self.metadata['latency'] = latency
self.metadata['implemented'] = implemented
self.metadata['pSuccess'] = pSuccess
self.edges = []
self.implemented = implemented
self.cost = cost
self.description = description
self.complexity = complexity
self.latency = latency
self.pDetect = pDetect

class Block(Node):
def __init__(self,
Expand All @@ -223,13 +217,11 @@ def __init__(self,
complexity: int = 0,
pDefend: int = 100):
super().__init__(label=label)
self.metadata = {}
self.metadata['cost'] = cost
self.metadata['description'] = description
self.metadata['complexity'] = complexity
self.metadata['implemented'] = implemented
self.metadata['pDefend'] = pDefend
self.edges = []
self.implemented = implemented
self.cost = cost
self.description = description
self.complexity = complexity
self.pDefend = pDefend

def insertBetween(self, a: Node, b: Node):
# When two nodes are connected, insert a blocking node between them, replacing or updating the existing connection.
Expand Down Expand Up @@ -263,14 +255,14 @@ def insertBetween(self, a: Node, b: Node):
class Discovery(Node):
def __init__(self,
label: str,
pSuccess: int = 100,
description: str = "",
sensitivity: int = 0,
value: int = 0,
markings: list = []):
super().__init__(label=label)
self.metadata={}
self.metadata['description'] = description
self.metadata['sensitivity'] = sensitivity
self.metadata['value'] = value
self.metadata['markings'] = markings
self.edges = []
self.pSuccess = pSuccess
self.description = description
self.sensitivity = sensitivity
self.value = value
self.markings = markings
15 changes: 12 additions & 3 deletions attacktree/renderer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
def _buildDot(self, node: Node, dot: Digraph, renderUnimplemented: bool, mappedEdges: dict={}, dotformat: dict={}):
node_attr = None # .dot formatting
unimplemented = False
if 'implemented' in node.metadata.keys() and node.metadata['implemented']==False:

if hasattr(node, "implemented") and node.implemented == False:
unimplemented = True
#TODO fix this wierd inverted logic

# The node is marked as unimplemented and we are told not to render those nodes
if renderUnimplemented == False and unimplemented == True:
Expand All @@ -39,6 +41,11 @@ def _buildDot(self, node: Node, dot: Digraph, renderUnimplemented: bool, mappedE
if unimplemented:
node_attr = node_attr | dotformat['_unimplemented_override'] # Style the unimplemented node

nodeLabel = node.label
if isinstance(node, (Action, Discovery)):
nodeLabel += f"\n{node.pSuccess}"
if isinstance(node, (Block)):
nodeLabel += f"\n{node.pDefend}"
dot.node(node.uniq, node.label, **node_attr)
else:
dot.node(node.uniq, node.label)
Expand All @@ -47,9 +54,11 @@ def _buildDot(self, node: Node, dot: Digraph, renderUnimplemented: bool, mappedE
# Make sure we don't draw a connection to an unimplemented node, if that renderUnimplemented == False

edgeImplemented = True # default drawing style is to assume implemented
if 'implemented' in node.metadata.keys() and node.metadata['implemented'] == False:

if isinstance(node, Block) and node.implemented == False:
edgeImplemented = False
if 'implemented' in edge.childNode.metadata.keys() and edge.childNode.metadata['implemented'] == False:

if isinstance(edge.childNode, Block) and edge.childNode.implemented == False:
edgeImplemented = False

# See if we should proceed with rendering the edge.
Expand Down
3 changes: 1 addition & 2 deletions tests/test_brain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,12 +253,11 @@ def test_pathEvaluationWithBlock(render=True):
assert(res['attackCost']==1000)
assert(res['pSuccess']==70)

assert(a.metadata['pSuccess']==100)
assert(a.pSuccess==100)

block = Block(label="FIREWALL",implemented=True,cost=0,pDefend=50)
block.insertBetween(a,b)

# assert(a.metadata['pSuccess']==50) # When a block gets inserted, it adjusts the pSuccess of the parent node
paths = []
_ = brain.pathsToVictory(root, paths)
res = brain.evaluatePath(paths[0])
Expand Down