Merge pull request #20 from hyakuhei/TestCoverage

Test coverage
hyakuhei · Jun 6, 2021 · 927a1f8 · 927a1f8
2 parents bf4a882 + 4aef309
commit 927a1f8
Show file tree

Hide file tree

Showing 6 changed files with 191 additions and 7 deletions.
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
@@ -17,7 +17,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: '3.x'
+        python-version: '3.9'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: [3.9]
 
     steps:
       - name: Install graphviz

diff --git a/README.md b/README.md
@@ -25,6 +25,8 @@ On MacOS this can be installed using `brew install graphviz`
 
 See https://graphviz.org/download/ for other options.
 
+Only python3.9 is supported at this time.
+
 # Installation
 attacktree is available in PyPI, we recommend installing in a virtualenv
 ```

diff --git a/attacktree/brain.py b/attacktree/brain.py
@@ -4,8 +4,7 @@
 # import some useful dicts
 from attacktree.models import rules
 
-
-# TODO: use logging framework
+import logging
 class Brain(object):
     def __init__(self):
         self.exploitChain = []
@@ -69,9 +68,10 @@ def evaluatePath(self, path):
 
                 # 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")
+                    logging.error(f"""
+                        Could not find an edge to {node.label}
+                        PrevNode: {prevNode.label}
+                        Path: {path}""")
                 else:
                     edgeToThisNode.pSuccess = results["pSuccess"]
 

diff --git a/tests/test_brain.py b/tests/test_brain.py
@@ -4,6 +4,7 @@
     Detect,
     Discovery,
     Edge,
+    Node,
     Root,
     Goal,
     mitreAttack,
@@ -119,6 +120,62 @@ def basicTree():
 
     return root
 
+def test_blockShortcut():
+    root = Root("have rock")
+    goal = Goal("hit player")
+
+    throws = root.action("throw rock")
+    dodge = throws.block("dodge", implemented=True)
+
+    assert(len(throws.edges) == 1)
+    assert isinstance(throws.edges[0], Edge)
+    edge = throws.edges[0]
+
+    assert(edge.childNode == dodge)
+    assert(edge.parentNode == throws)
+
+    assert isinstance(dodge, Block)
+    assert isinstance(dodge, Node)
+    assert(len(dodge.parentEdges) == 1)
+    assert isinstance(dodge.parentEdges[0], Edge)
+
+def test_describe():
+    root = Root("root")
+    goal = Goal("goal")
+
+    throws = root.action("A")
+    dodge = throws.block("B", implemented=True)
+
+    description = throws.edges[0].describe()
+
+    assert(description == "Edge 'Fail' connects 'A' to 'B'")
+
+def test_repr():
+    root = Root("root")
+    goal = Goal("goal")
+
+    throws = root.action("A")
+    dodge = throws.block("B", implemented=True)
+
+    description = throws.edges[0].__repr__()
+
+    assert(description == "Edge 'Fail' connects 'A' to 'B'")
+
+
+def test_detectShortcut():
+    root = Root("have rock")
+    goal = Goal("hit player")
+
+    throws = root.action("throw rock")
+    alarm = throws.detect("player sees rock", implemented=True)
+
+    assert isinstance(alarm, Detect)
+    assert(len(throws.edges) == 1)
+    edge = throws.edges[0]
+
+    assert isinstance(edge, Edge)
+    assert edge.parentNode == throws
+    assert edge.childNode == alarm
 
 def test_buildTree(render=True):
     root = basicTree()

diff --git a/tests/test_renderer.py b/tests/test_renderer.py
@@ -0,0 +1,125 @@
+from attacktree.models import (
+    Action,
+    Block,
+    Detect,
+    Discovery,
+    Edge,
+    Node,
+    Root,
+    Goal,
+    mitreAttack,
+    rules,
+)
+from attacktree.renderer import Renderer
+from attacktree.brain import Brain
+
+import pytest
+import logging
+import inspect
+
+def basicTree():
+    root = Root("Root")
+    goal = Goal("Systems Access")
+
+    # Create three top tier actions
+    networkRecon = root.add(
+        Action(
+            label="Network Recon",
+            chain=mitreAttack["recon"],
+            cost=0,
+            time=24,
+            objective="Find network attack surface",
+            pSuccess=100,
+            detections=None,
+        )
+    )
+
+    dnsEnumeration = root.add(
+        Action(
+            label="DNS Enumeration",
+            chain=mitreAttack["recon"],
+            cost=0,
+            time=4,
+            objective="Identify all subdomains",
+            pSuccess=100,
+        )
+    )
+
+    linkedInResearch = root.add(
+        Action(
+            label="LinkedIn Research",
+            chain=mitreAttack["recon"],
+            cost=0,
+            time=6,
+            objective="Identify names and email addresses of current employees and customers",
+            pSuccess=100,
+        )
+    )
+
+    # Stuff learned from those activities
+    vpnEndpoint = networkRecon.discovery("VPN Endpoints")
+    sshEndpoint = networkRecon.discovery("SSH Endpoints")
+    subdomains = dnsEnumeration.discovery("subdomains")
+    employeeNames = linkedInResearch.discovery("Employee Names")
+    keyCustomerNames = linkedInResearch.discovery("Key Customer Details")
+
+    # Actions taken based on those discoveries
+    credentialStuffing = Action(
+        label="Credential Stuffing",
+        chain=mitreAttack["credStuffing"],
+        cost=500,
+        time=12,
+        objective="Try known username/password",
+        pSuccess=100,
+    )
+    _ = sshEndpoint.add(credentialStuffing)
+    _ = vpnEndpoint.add(credentialStuffing)
+    _ = employeeNames.add(credentialStuffing)
+
+    sqlmap = subdomains.add(
+        Action(
+            label="sqlmap", chain=mitreAttack["execution"], cost=0, time=2, pSuccess=75
+        )
+    )
+
+    nikto = subdomains.action("nikto")
+
+    phishing = employeeNames.action("SET Phishing")
+    keyCustomerNames.connectTo(phishing)
+
+    # Action Results
+    sqli = sqlmap.add(
+        Discovery(
+            label="Blind injection, viable RCE",
+            description="",
+            sensitivity=10,
+            value=1000,
+        )
+    )
+
+    dbExploit = sqli.add(
+        Action(
+            label="Craft & Deploy RCE",
+            chain=mitreAttack["execution"],
+            cost=0,
+            time=2,
+            pSuccess=75,
+        )
+    )
+
+    # Action Results that get to the goal
+    credentialStuffing.connectTo(goal, label="Passwords Reused")
+    phishing.connectTo(goal, label="Credentials Stolen")
+    dbExploit.connectTo(goal)
+
+    return root
+
+def test_contextManager(render=False):
+    with Renderer(root="Reality", goal="Attacker gets data from bucket") as graph:
+        pwd = graph.root.action("Use password")
+        block = pwd.block("Block", implemented=False)
+        block.connectTo(graph.goal)
+
+
+
+