Merge pull request #23 from hyakuhei/blockedPathStyling

Blocked path styling
hyakuhei · Sep 9, 2021 · aeb3ef7 · aeb3ef7
2 parents 927a1f8 + 66dd594
commit aeb3ef7
Show file tree

Hide file tree

Showing 15 changed files with 609 additions and 73 deletions.
diff --git a/ExampleNotebook.ipynb b/ExampleNotebook.ipynb
diff --git a/attacktree/brain.py b/attacktree/brain.py
@@ -5,6 +5,8 @@
 from attacktree.models import rules
 
 import logging
+
+
 class Brain(object):
     def __init__(self):
         self.exploitChain = []
@@ -68,10 +70,12 @@ def evaluatePath(self, path):
 
                 # This shouldn't happen and we should try to get rid of this check.
                 if edgeToThisNode is None:
-                    logging.error(f"""
+                    logging.error(
+                        f"""
                         Could not find an edge to {node.label}
                         PrevNode: {prevNode.label}
-                        Path: {path}""")
+                        Path: {path}"""
+                    )
                 else:
                     edgeToThisNode.pSuccess = results["pSuccess"]
 

diff --git a/attacktree/models.py b/attacktree/models.py
@@ -143,6 +143,10 @@ def discovery(self, label: str, edge_label: str = "Learn"):
         self.connectTo(d, edge_label)
         return d
 
+    # I always get discover and discovery mixed up so added a shortcut
+    def discover(self, label: str, edge_label: str = "Learn"):
+        return self.discovery(label,edge_label)
+
     def __repr__(self):
         return f"{self.__class__.__name__}:{id(self)}"
 
@@ -265,9 +269,11 @@ def insertBetween(self, a: Node, b: Node):
             self.connectTo(b)
 
             # a.metadata['pSuccess'] == the probability of the technique working
-            # block.metadata['pDefend'] == the probability of the techniuqe being subsequently blocked
+            # block.metadata['pDefend'] == the probability of the technique being subsequently blocked
             # edge.pSuccess is the probability of progressing
 
+        return self  # fluent_pattern
+
 
 # label: 'The name of the node'
 # description: 'A description of the data/information

diff --git a/attacktree/renderer.py b/attacktree/renderer.py
@@ -22,71 +22,48 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
             self.render()
         return None
 
-    # A recursive function that walks the node tree
-    # And creates a graph for turning into graphviz
-    # TODO: Fix complexity PEP8 C901
+    # Addresses https://github.com/hyakuhei/attackTrees/issues/22
+    # will replace _buildDot()
     def _buildDot(
         self,
         node: Node,
         dot: Digraph,
-        renderUnimplemented: bool,
         mappedEdges: dict = {},
         dotformat: dict = {},
+        defendedPath: bool = False,
     ):
-        node_attr = None  # .dot formatting
-        unimplemented = False
+        node_attr = None
+        if node.__class__.__name__ in dotformat.keys():
+            node_attr = dotformat[node.__class__.__name__]
 
+        implemented = (
+            True  # common case, most things will be implemented, we default to that.
+        )
+        # TODO: Move this default logic elsewhere, maybe force it into the basic node?
         if hasattr(node, "implemented") and node.implemented is 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 is False and unimplemented is True:
-            return
+            implemented = False
 
         if node.__class__.__name__ in dotformat.keys():
             node_attr = dotformat[node.__class__.__name__]
             # Overload the default formatting shape if the Node is flagged as unimplemented
-            if unimplemented:
-                # Style the unimplemented node
-                node_attr = node_attr | dotformat["_unimplemented_override"]
 
             nodeLabel = node.label
             if isinstance(node, (Action, Discovery)):
                 nodeLabel += f"\n{node.pSuccess}"
             if isinstance(node, (Block)):
                 nodeLabel += f"\n{node.pDefend}"
+                defendedPath = True
             dot.node(node.uniq, node.label, **node_attr)
         else:
             dot.node(node.uniq, node.label)
 
         for edge in node.getEdges():
-            # Make sure we don't draw a connection to an unimplemented node, if that renderUnimplemented is False
-
-            edgeImplemented = True  # default drawing style is to assume implemented
-
-            if isinstance(node, Block) and node.implemented is False:
-                edgeImplemented = False
-
-            if (
-                isinstance(edge.childNode, Block)
-                and edge.childNode.implemented is False
-            ):
-                edgeImplemented = False
-
-            # See if we should proceed with rendering the edge.
-            # If not, we actually don't need to follow this branch any further
-            # Short circuit the loop with a 'continue'
-            if renderUnimplemented is False and edgeImplemented is False:
-                continue
-
             # Setup default edge rendering style
-            edge_attr = dotformat["Edge"]
-
-            # Override style for unimplemented edge
-            if edgeImplemented is False:
-                # style the unimplemented edge
-                edge_attr = edge_attr | dotformat["_unimplemented_edge"]
+            # TODO: Decide if we are talking about "path" or "edge"
+            if defendedPath is True:
+                edge_attr = dotformat["defendedEdge"]
+            else:
+                edge_attr = dotformat["Edge"]
 
             label = edge.label
             if edge.pSuccess is not None and edge.pSuccess != -1:
@@ -101,11 +78,12 @@ def _buildDot(
                 self._buildDot(
                     node=edge.childNode,
                     dot=dot,
-                    renderUnimplemented=renderUnimplemented,
                     mappedEdges=mappedEdges,
                     dotformat=dotformat,
+                    defendedPath=defendedPath
                 )  # recurse
 
+
     def loadStyle(self, path: str):
         # TODO: Do error handling
         with open(path) as json_file:
@@ -130,27 +108,25 @@ def buildDot(
             with resources.open_text("attacktree", "style.json") as fid:
                 style = json.load(fid)
 
-        self._buildDot(
-            root, dot, dotformat=style, renderUnimplemented=includeUnimplemented
-        )  # recursive call
+        self._buildDot(root, dot, dotformat=style)  # recursive call
         return dot
 
     def render(
         self,
         root: Node = None,
-        renderUnimplemented: bool = True,
         style: dict = None,
         fname: str = "attacktree-graph",
         fout: str = "png",
         renderOnExit=False,
+        renderUnimplemented=False
     ):
-
+        logging.warning("renderUnimplemented is deprectaed")
         self.renderOnExit = renderOnExit
         if root is None and self.root is not None:
-            root = self.root # sometimes we invoke as a context manager
+            root = self.root  # sometimes we invoke as a context manager
 
         dot = self.buildDot(
-            root=root, includeUnimplemented=renderUnimplemented, style=style
+            root=root, style=style
         )
         dot.format = fout
-        dot.render(fname, view=True)
+        dot.render(fname, view=True)
diff --git a/attacktree/style.json b/attacktree/style.json
@@ -3,61 +3,68 @@
         "color":"#c0c2c0",
         "shape":"plaintext",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Root":{
         "color":"#03fc9d",
         "shape":"box",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Goal":{
         "color":"#03fc9d",
         "shape":"box",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Action":{
         "color":"#ff5c5c",
         "shape":"plaintext",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Discovery":{
         "color":"#ffdc5c",
         "shape":"box",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Detect":{
         "color":"#ff7df0",
         "shape":"plaintext",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "Block":{
         "color":"#5cc1ff",
         "shape":"plaintext",
         "style":"filled, rounded",
-        "fontname":"Lato",
+        "fontname":"Arial",
         "margin":"0.2"
     },
     "_unimplemented_override":{
-        "style":"filled",
-        "color":"#b0c1cf"
+        "color":"#b0c1cf",
+        "shape":"plaintext",
+        "style":"filled, rounded",
+        "fontname":"Arial",
+        "margin":"0.2"
     },
     "Edge":{
-        "fontname":"Lato",
+        "fontname":"Arial",
         "color":"#1f1f1f"
     },
     "_unimplemented_edge":{
         "color":"#919191",
         "style":"dashed"
+    },
+    "defendedEdge":{
+        "color":"#919191",
+        "style":"dashed"
     }
 }
diff --git a/examples/aciBreakout b/examples/aciBreakout
@@ -0,0 +1,46 @@
+digraph {
+	graph [nodesep=0.2 overlap=false ranksep=0.4 splines=True]
+	"689d84be59424704a0158189cdee7f29" [label="Signed up to Azure" color="#03fc9d" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"689d84be59424704a0158189cdee7f29" -> f512dcfc14c6464eabc045318e1fee4c [label=Next color="#1f1f1f" fontname=Arial]
+	f512dcfc14c6464eabc045318e1fee4c [label="Deploy WhoC container to view runtime" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	f512dcfc14c6464eabc045318e1fee4c -> "6b53663b9d06461093e4b8f7bad45728" [label=Learn color="#1f1f1f" fontname=Arial]
+	"6b53663b9d06461093e4b8f7bad45728" [label="OLD runc version 1.0.0-r2" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"6b53663b9d06461093e4b8f7bad45728" -> "8e51060ab9f0457885cbbd44030ac14b" [label=Next color="#1f1f1f" fontname=Arial]
+	"8e51060ab9f0457885cbbd44030ac14b" [label="Deploy exploit container for CVE-2019-5736" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"8e51060ab9f0457885cbbd44030ac14b" -> bbb01944bd6e41708ec1d15fbaac0a32 [label="" color="#1f1f1f" fontname=Arial]
+	bbb01944bd6e41708ec1d15fbaac0a32 [label="Reverse shell on worker node" color="#03fc9d" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	bbb01944bd6e41708ec1d15fbaac0a32 -> "51bff6cb11ba42cdb27efdf27b6c5360" [label=Learn color="#1f1f1f" fontname=Arial]
+	"51bff6cb11ba42cdb27efdf27b6c5360" [label="Read kubelet credentials from disk" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"51bff6cb11ba42cdb27efdf27b6c5360" -> "81f957b8242e4db79e78271e8f65dbe5" [label=Next color="#1f1f1f" fontname=Arial]
+	"81f957b8242e4db79e78271e8f65dbe5" [label="Call KubeAPI describe pods" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"81f957b8242e4db79e78271e8f65dbe5" -> d425fdd21df746a2a0741b6d46561dfa [label=Learn color="#1f1f1f" fontname=Arial]
+	d425fdd21df746a2a0741b6d46561dfa [label="100+ customer pods on 120 nodes
+Each customer has their own namespace" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"81f957b8242e4db79e78271e8f65dbe5" -> dac8b598bf7245e9be7cda6138aa0988 [label=Learn color="#1f1f1f" fontname=Arial]
+	dac8b598bf7245e9be7cda6138aa0988 [label="OLD Kubernetes versions v1.8.4, v1.9.10, v1.10.9" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	dac8b598bf7245e9be7cda6138aa0988 -> "25c9192361da4c33bc47df5782eb1cee" [label=Learn color="#1f1f1f" fontname=Arial]
+	"25c9192361da4c33bc47df5782eb1cee" [label="CVE-2018-1002102 kube-api follows 302 redirect" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"25c9192361da4c33bc47df5782eb1cee" -> "1cbc4146eb664fcab3387195c0a0e4ff" [label=Next color="#1f1f1f" fontname=Arial]
+	"1cbc4146eb664fcab3387195c0a0e4ff" [label="Attempt to redirect to kube-api pod" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"1cbc4146eb664fcab3387195c0a0e4ff" -> "424e5104c76a478fb20852d151ade0bd" [label=Fail color="#1f1f1f" fontname=Arial]
+	"424e5104c76a478fb20852d151ade0bd" [label="ACI uses a 'bridge' POD which is not impacted by this issue" color="#5cc1ff" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"1cbc4146eb664fcab3387195c0a0e4ff" -> "55f7d70c6f79435ca8e40ed15680248c" [label=Learn color="#1f1f1f" fontname=Arial]
+	"55f7d70c6f79435ca8e40ed15680248c" [label="ServiceAccount in 'AuthorizationHeader' of Exec requests" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"55f7d70c6f79435ca8e40ed15680248c" -> "2b4442634821457c84558116c3371a54" [label=Learn color="#1f1f1f" fontname=Arial]
+	"2b4442634821457c84558116c3371a54" [label="Decoded JWT shows this token belongs to 'bridge' service" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"2b4442634821457c84558116c3371a54" -> "9d9d602a34df49c9b6eaab5434d9e166" [label=Next color="#1f1f1f" fontname=Arial]
+	"9d9d602a34df49c9b6eaab5434d9e166" [label="Call SelfSubjectAccessReview with 'bridge' token" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"9d9d602a34df49c9b6eaab5434d9e166" -> "06bea6172c13440fa720f16f3ec20e02" [label=Learn color="#1f1f1f" fontname=Arial]
+	"06bea6172c13440fa720f16f3ec20e02" [label="Cluster-wide permissions
+pods/exec privilege" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"06bea6172c13440fa720f16f3ec20e02" -> "66acfa9de3e94c819ac5c16cc286a150" [label=Next color="#1f1f1f" fontname=Arial]
+	"66acfa9de3e94c819ac5c16cc286a150" [label="Exec into shell on kube-API" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"66acfa9de3e94c819ac5c16cc286a150" -> "3eb45608f8984efda912c83c185a4232" [label="" color="#1f1f1f" fontname=Arial]
+	"3eb45608f8984efda912c83c185a4232" [label="Access to other tenants data" color="#03fc9d" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"81f957b8242e4db79e78271e8f65dbe5" -> "7c7794543ab14a3c8071ff73a965efb2" [label=Learn color="#1f1f1f" fontname=Arial]
+	"7c7794543ab14a3c8071ff73a965efb2" [label="Kubelets run with anonymous access" color="#ffdc5c" fontname=Arial margin=0.2 shape=box style="filled, rounded"]
+	"7c7794543ab14a3c8071ff73a965efb2" -> "546a5bf3869f43fb929912ae2de0b58b" [label=Next color="#1f1f1f" fontname=Arial]
+	"546a5bf3869f43fb929912ae2de0b58b" [label="Access another customer's kubelet" color="#ff5c5c" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+	"546a5bf3869f43fb929912ae2de0b58b" -> "308e0ba22a3245eb886a4496dd3b5daf" [label=Fail color="#1f1f1f" fontname=Arial]
+	"308e0ba22a3245eb886a4496dd3b5daf" [label="Blocked by firewall" color="#5cc1ff" fontname=Arial margin=0.2 shape=plaintext style="filled, rounded"]
+}
diff --git a/examples/aciBreakout.png b/examples/aciBreakout.png
diff --git a/examples/aciBreakout.py b/examples/aciBreakout.py
@@ -0,0 +1,46 @@
+# ACI Breakout Tree
+# Based on the unit42 writeup here:
+
+from attacktree.models import Action, Block, Detect, Discovery, Edge, Root, Goal
+from attacktree.renderer import Renderer
+
+root = Root("Signed up to Azure")
+goal = Goal("Access to other tenants data")
+
+whoC = root.add(Action("Deploy WhoC container to view runtime"))
+omgRunC = whoC.discover("OLD runc version 1.0.0-r2")
+exploit = omgRunC.action("Deploy exploit container for CVE-2019-5736")
+shell = Root("Reverse shell on worker node")
+exploit.add(shell)
+creds = shell.discover("Read kubelet credentials from disk")
+
+describePods = creds.action("Call KubeAPI describe pods")
+pods = describePods.discover(
+    "100+ customer pods on 120 nodes\nEach customer has their own namespace"
+)
+versions = describePods.discover("OLD Kubernetes versions v1.8.4, v1.9.10, v1.10.9")
+anonymousAccess = describePods.discover("Kubelets run with anonymous access")
+
+otherKubelet = anonymousAccess.action("Access another customer's kubelet")
+blockedbyFirewall = otherKubelet.block("Blocked by firewall", implemented=True)
+
+cve = versions.discover("CVE-2018-1002102 kube-api follows 302 redirect")
+exploit2 = cve.action("Attempt to redirect to kube-api pod")
+fail = exploit2.block(
+    "ACI uses a 'bridge' POD which is not impacted by this issue", implemented=True
+)
+
+interesting = exploit2.discover(
+    "ServiceAccount in 'AuthorizationHeader' of Exec requests"
+)
+bridgeToken = interesting.discover(
+    "Decoded JWT shows this token belongs to 'bridge' service"
+)
+
+accessReview = bridgeToken.action("Call SelfSubjectAccessReview with 'bridge' token")
+privs = accessReview.discover("Cluster-wide permissions\npods/exec privilege")
+
+gameOver = privs.action("Exec into shell on kube-API")
+gameOver.add(goal)
+
+Renderer().render(root, fname="aciBreakout")
diff --git a/examples/attacktree-graph b/examples/attacktree-graph
@@ -0,0 +1,6 @@
+digraph {
+	graph [nodesep=0.2 overlap=false ranksep=0.4 splines=True]
+	"844673a674d044e399b6e7fd764b0146" [label=Internet color="#03fc9d" fontname=Lato margin=0.2 shape=box style="filled, rounded"]
+	"844673a674d044e399b6e7fd764b0146" -> "5d7865a3aac1420299eda2e7175978f7" [label=Next color="#1f1f1f" fontname=Lato]
+	"5d7865a3aac1420299eda2e7175978f7" [label="RCE in application" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+}
diff --git a/examples/attacktree-graph.png b/examples/attacktree-graph.png
diff --git a/examples/siloscape b/examples/siloscape
@@ -0,0 +1,22 @@
+digraph {
+	graph [nodesep=0.2 overlap=false ranksep=0.4 splines=True]
+	"709183617a5c43139a80903f3b1f48b7" [label=Internet color="#03fc9d" fontname=Lato margin=0.2 shape=box style="filled, rounded"]
+	"709183617a5c43139a80903f3b1f48b7" -> "65226adffaf54f24bab5845e5c25c059" [label=Next color="#1f1f1f" fontname=Lato]
+	"65226adffaf54f24bab5845e5c25c059" [label="RCE in application" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	"65226adffaf54f24bab5845e5c25c059" -> c72dfe9823fb49dcb1409ee4970c8204 [label=Fail color="#1f1f1f" fontname=Lato]
+	c72dfe9823fb49dcb1409ee4970c8204 [label="Keep containers up to date" color="#5cc1ff" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	"65226adffaf54f24bab5845e5c25c059" -> "3528833dfdcd458e9bc4e43d4ce5ab8e" [label=Next color="#1f1f1f" fontname=Lato]
+	"3528833dfdcd458e9bc4e43d4ce5ab8e" [label="Execute Siloscape" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	"3528833dfdcd458e9bc4e43d4ce5ab8e" -> "0bf4dbc7127245a3941a361b5f66d9f7" [label=Learn color="#1f1f1f" fontname=Lato]
+	"0bf4dbc7127245a3941a361b5f66d9f7" [label="Privileged Access" color="#ffdc5c" fontname=Lato margin=0.2 shape=box style="filled, rounded"]
+	"0bf4dbc7127245a3941a361b5f66d9f7" -> "9a88f331bcf344309953d2327b9d5a4d" [label=Next color="#1f1f1f" fontname=Lato]
+	"9a88f331bcf344309953d2327b9d5a4d" [label="SymLink root volume" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	"9a88f331bcf344309953d2327b9d5a4d" -> "38faf45dc8f249f49a97f1efc8953dc3" [label=Next color="#1f1f1f" fontname=Lato]
+	"38faf45dc8f249f49a97f1efc8953dc3" [label="Find Kubernetes creds on disk" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	"38faf45dc8f249f49a97f1efc8953dc3" -> e4a44aa7d5bb424fb49d2bb81d59f13d [label=Next color="#1f1f1f" fontname=Lato]
+	e4a44aa7d5bb424fb49d2bb81d59f13d [label="Deploy malicious containers" color="#ff5c5c" fontname=Lato margin=0.2 shape=plaintext style="filled, rounded"]
+	e4a44aa7d5bb424fb49d2bb81d59f13d -> "45d5d0024f3e4278a4ac2b8d087b2639" [label=Fail color="#919191" fontname=Lato style=dashed]
+	"45d5d0024f3e4278a4ac2b8d087b2639" [label="Windows containers must have low privilege" color="#b0c1cf" fontname=Lato margin=0.2 shape=plaintext style=filled]
+	e4a44aa7d5bb424fb49d2bb81d59f13d -> "5fd23df968c843b0a237bdea45001b05" [label="" color="#1f1f1f" fontname=Lato]
+	"5fd23df968c843b0a237bdea45001b05" [label="Launch Containers" color="#03fc9d" fontname=Lato margin=0.2 shape=box style="filled, rounded"]
+}
diff --git a/examples/siloscape.png b/examples/siloscape.png