Merge pull request #15 from hyakuhei/Black

Reformatted using Black
hyakuhei · May 15, 2021 · af93480 · af93480
2 parents 57b80bc + 9c99bdf
commit af93480
Show file tree

Hide file tree

Showing 5 changed files with 309 additions and 294 deletions.
diff --git a/attacktree/brain.py b/attacktree/brain.py
@@ -1,4 +1,3 @@
-
 # import Node classes
 from attacktree.models import Action, Block, Goal, Node
 
@@ -13,11 +12,9 @@ def __init__(self):
 
     # Walk the tree, adding to the chain (DFS)
     # If we hit the goal, add that chain to our paths
-    def pathsToVictory(self,
-                       node: Node,
-                       paths: list = None,
-                       chain: list = None,
-                       walked: dict = None):
+    def pathsToVictory(
+        self, node: Node, paths: list = None, chain: list = None, walked: dict = None
+    ):
 
         if walked is None:
             walked = {}
@@ -38,8 +35,9 @@ def pathsToVictory(self,
         edges = node.getEdges()
         for edge in edges:
             if edge not in walked:
-                self.pathsToVictory(edge.childNode, paths,
-                                    chain=chain.copy(), walked=walked)
+                self.pathsToVictory(
+                    edge.childNode, paths, chain=chain.copy(), walked=walked
+                )
                 walked[edge] = True  # Stops walking a cycle more than once
 
         return paths
@@ -49,19 +47,18 @@ def evaluatePath(self, path):
         # 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']
+            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.cost
-                results['time'] += node.time
-                results['pSuccess'] = int(
-                    (results['pSuccess'] * node.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.cost
-                results['pSuccess'] -= node.pDefend
+                results["defenceCost"] += node.cost
+                results["pSuccess"] -= node.pDefend
                 # TODO block time
 
             if prevNode is not None:
@@ -76,7 +73,7 @@ def evaluatePath(self, path):
                     print(f"PrevNode: {prevNode.label}")
                     print(f"Path: {path}\n")
                 else:
-                    edgeToThisNode.pSuccess = results['pSuccess']
+                    edgeToThisNode.pSuccess = results["pSuccess"]
 
             prevNode = node
             # Can't just throw in a backfref because a node can have multiple parents

diff --git a/attacktree/models.py b/attacktree/models.py
@@ -7,76 +7,76 @@
 # TODO: Replace metadata with class values
 
 defaultEdgeLabels = {
-    'Action': "Next",
-    'Discovery': "Learn",
-    'Block': "Fail",
-    'Detect': "Detect"
+    "Action": "Next",
+    "Discovery": "Learn",
+    "Block": "Fail",
+    "Detect": "Detect",
 }
 
 rules = {
-    'pSuccess': {
-        'math': "multiply",
-        'startWith': 100,
-        'unit': "probability",
-        'formatString': "{}",
-        'description': """
+    "pSuccess": {
+        "math": "multiply",
+        "startWith": 100,
+        "unit": "probability",
+        "formatString": "{}",
+        "description": """
             pSuccess: int value between 0 and 1 that describes the
             probability of an `Action` or `Block` being effective within the given time
-            """
+            """,
     },
     "attackCost": {
-        'math': "add",
-        'startWith': 0,
-        'unit': "dollars",
-        'formatString': "${}",
-        'description': """
+        "math": "add",
+        "startWith": 0,
+        "unit": "dollars",
+        "formatString": "${}",
+        "description": """
             cost: estimate of the number of dollars required to be invested to
-            effectively run an `Action` or `Block`"""
+            effectively run an `Action` or `Block`""",
     },
     "time": {
-        'math': "add",
-        'startWith': 0,
-        'unit': "hours",
-        'formatString': "{}h",
-        'description': """
+        "math": "add",
+        "startWith": 0,
+        "unit": "hours",
+        "formatString": "{}h",
+        "description": """
             time: estimate of how long this `Action` or `Block` will take to
-            be effective"""
+            be effective""",
     },
     "defenceCost": {
-        'math': "add",
-        'startWith': 0,
-        'unit': "dollars",
-        'formatString': "${}",
-        'description': "Cost of defensive controls"
-    }
+        "math": "add",
+        "startWith": 0,
+        "unit": "dollars",
+        "formatString": "${}",
+        "description": "Cost of defensive controls",
+    },
 }
 
 mitreAttack = {
-    'recon': {
-        'shortName': "recon",
-        'friendlyName': "Reconnaissance",
-        'objective': "The adversary is trying to gather information they can use to plan future operations",
-        'url': "https://attack.mitre.org/tactics/TA0043/"
+    "recon": {
+        "shortName": "recon",
+        "friendlyName": "Reconnaissance",
+        "objective": "The adversary is trying to gather information they can use to plan future operations",
+        "url": "https://attack.mitre.org/tactics/TA0043/",
     },
-    'resourceDev': {
-        'shortName': "resourceDev",
-        'friendlyName': "Resource Development",
-        'objective': "The adversary is trying to establish resources they can use to support operations.",
-        'url': "https://attack.mitre.org/tactics/TA0042/"
+    "resourceDev": {
+        "shortName": "resourceDev",
+        "friendlyName": "Resource Development",
+        "objective": "The adversary is trying to establish resources they can use to support operations.",
+        "url": "https://attack.mitre.org/tactics/TA0042/",
     },
-    'credStuffing': {
-        'shortName': "credStuffing",
-        'friendlyName': "Credential Stuffing",
-        'objective': """Adversaries may use credentials obtained from breach dumps of unrelated accounts
+    "credStuffing": {
+        "shortName": "credStuffing",
+        "friendlyName": "Credential Stuffing",
+        "objective": """Adversaries may use credentials obtained from breach dumps of unrelated accounts
             to gain access to target accounts through credential overlap.""",
-        'url': "https://attack.mitre.org/techniques/T1110/004/"
+        "url": "https://attack.mitre.org/techniques/T1110/004/",
+    },
+    "execution": {
+        "shortName": "execution",
+        "friendlyName": "Execution",
+        "objective": "The adversary is trying to run malicious code.",
+        "url": "https://attack.mitre.org/tactics/TA0002/",
     },
-    'execution': {
-        'shortName': "execution",
-        'friendlyName': "Execution",
-        'objective': "The adversary is trying to run malicious code.",
-        'url': "https://attack.mitre.org/tactics/TA0002/"
-    }
 }
 
 # TODO: Replace concrete numbers with ranges and confidence intervals.
@@ -173,16 +173,15 @@ def __repr__(self):
 
 
 class Root(Node):
-    def __init__(self,
-                 label: str):
+    def __init__(self, label: str):
         super().__init__(label=label)
 
 
 class Goal(Node):
-    def __init__(self,
-                 label: str):
+    def __init__(self, label: str):
         super().__init__(label=label)
 
+
 # label: 'The name of the node',
 # chain: 'The stage of the Mitre Att&ck chain represented, e.g "recon"'
 # cost: 'Estimate of any material cost to this path, in dollars, does not include time'
@@ -192,14 +191,16 @@ def __init__(self,
 
 
 class Action(Node):
-    def __init__(self,
-                 label: str,
-                 chain: dict = None,
-                 cost: int = 0,
-                 time: int = 0,
-                 objective: str = "",
-                 pSuccess: int = 100,
-                 detections: list = []):
+    def __init__(
+        self,
+        label: str,
+        chain: dict = None,
+        cost: int = 0,
+        time: int = 0,
+        objective: str = "",
+        pSuccess: int = 100,
+        detections: list = [],
+    ):
         super().__init__(label=label)
         self.pSuccess = pSuccess
         self.chain = chain
@@ -209,14 +210,16 @@ def __init__(self,
 
 
 class Detect(Node):
-    def __init__(self,
-                 label: str,
-                 implemented: bool,
-                 cost: int = 0,
-                 description: str = "",
-                 complexity: int = 0,
-                 latency: int = 0,
-                 pDetect: int = 100):
+    def __init__(
+        self,
+        label: str,
+        implemented: bool,
+        cost: int = 0,
+        description: str = "",
+        complexity: int = 0,
+        latency: int = 0,
+        pDetect: int = 100,
+    ):
         super().__init__(label=label)
         self.implemented = implemented
         self.cost = cost
@@ -227,13 +230,15 @@ def __init__(self,
 
 
 class Block(Node):
-    def __init__(self,
-                 label: str,
-                 implemented: bool,
-                 cost: int = 0,
-                 description: str = "",
-                 complexity: int = 0,
-                 pDefend: int = 100):
+    def __init__(
+        self,
+        label: str,
+        implemented: bool,
+        cost: int = 0,
+        description: str = "",
+        complexity: int = 0,
+        pDefend: int = 100,
+    ):
         super().__init__(label=label)
         self.implemented = implemented
         self.cost = cost
@@ -270,13 +275,15 @@ def insertBetween(self, a: Node, b: Node):
 # value: 'Perceived monetary value if applicable'
 # markings: 'Any specific markings for the data, like PII, SPI, HIPPA etc'
 class Discovery(Node):
-    def __init__(self,
-                 label: str,
-                 pSuccess: int = 100,
-                 description: str = "",
-                 sensitivity: int = 0,
-                 value: int = 0,
-                 markings: list = []):
+    def __init__(
+        self,
+        label: str,
+        pSuccess: int = 100,
+        description: str = "",
+        sensitivity: int = 0,
+        value: int = 0,
+        markings: list = [],
+    ):
         super().__init__(label=label)
         self.pSuccess = pSuccess
         self.description = description