Adding clean (apart from C901) flake8 passing

hyakuhei · May 15, 2021 · ea074fe · ea074fe
1 parent c03ab39
commit ea074fe
Show file tree

Hide file tree

Showing 4 changed files with 159 additions and 127 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -27,9 +27,9 @@ jobs:
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
-          # flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          flake8 attacktree --count --select=E9,F63,F7,F82 --show-source --statistics
           # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-          # flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+          flake8 attacktree --count --ignore=C901 --exit-zero --max-complexity=10 --max-line-length=127 --statistics
       - name: Test with pytest
         run: |
           pytest
diff --git a/attacktree/brain.py b/attacktree/brain.py
@@ -1,86 +1,85 @@
 
 # import Node classes
-from attacktree.models import Action, Block, Detect, Discovery, Edge, Root, Goal, Node
+from attacktree.models import Action, Block, Goal, Node
 
 # import some useful dicts
-from attacktree.models import mitreAttack, rules
+from attacktree.models import rules
 
-import logging
-import copy
 
+# TODO: use logging framework
 class Brain(object):
     def __init__(self):
         self.exploitChain = []
 
     # Walk the tree, adding to the chain (DFS)
-    # If we hit the goal, add that chain to our paths 
+    # 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):
-        
-        if walked == None:
+                       paths: list = None,
+                       chain: list = None,
+                       walked: dict = None):
+
+        if walked is None:
             walked = {}
 
-        if paths == None:
+        if paths is None:
             paths = []
 
-        if chain == None:
+        if chain is None:
             chain = []
 
         chain.append(node)
-        
+
         # If this node is a Goal then YAY! We have a goal
         if isinstance(node, Goal):
             paths.append(chain.copy())
             return paths
-        
+
         edges = node.getEdges()
         for edge in edges:
             if edge not in walked:
-                self.pathsToVictory(edge.childNode, paths, chain=chain.copy(), walked=walked)
-                walked[edge] = True # Stops walking a cycle more than once
+                self.pathsToVictory(edge.childNode, paths,
+                                    chain=chain.copy(), walked=walked)
+                walked[edge] = True  # Stops walking a cycle more than once
 
         return paths
 
     # Walk the given path, add up stats and annotate edges
     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
+        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() )
+            # 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['pSuccess'] = int(
+                    (results['pSuccess'] * node.pSuccess) / 100)
             if isinstance(node, (Block)):
-                results['defenceCost']  += node.cost
+                results['defenceCost'] += node.cost
                 results['pSuccess'] -= node.pDefend
-                #TODO block time
+                # TODO block time
 
             if prevNode is not None:
-                edgeToThisNode = None 
+                edgeToThisNode = None
                 for edge in prevNode.edges:
                     if edge.childNode == node:
                         edgeToThisNode = edge
 
-                if edgeToThisNode is None: # This shouldn't happen and we should try to get rid of this check.
+                # This shouldn't happen and we should try to get rid of this check.
+                if edgeToThisNode is None:
                     print(f"Could not find an edge to {node.label}")
                     print(f"PrevNode: {prevNode.label}")
                     print(f"Path: {path}\n")
                 else:
                     edgeToThisNode.pSuccess = results['pSuccess']
-            
+
             prevNode = node
             # Can't just throw in a backfref because a node can have multiple parents
             # End outer for by setting current node as the next (prevNode )
 
-        return results        
-
-
-
+        return results