Merge pull request #17 from hyakuhei/BasicJupyter

Basic jupyter support
hyakuhei · May 16, 2021 · b47951d · b47951d
2 parents 5092ce3 + 88ea571
commit b47951d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 22 deletions.
diff --git a/README.md b/README.md
@@ -9,6 +9,10 @@ Programatically model trees like those described by [Kelly Shortridge](https://t
 
 The goal is to decouple the model from the view. In reality I'm removing the need for the user to understand Graphviz and introducing a need for them to understand python.
 
+## Jupyiter support (needs better documentation)
+If your system (or venv) has attacktree installed you can now build trees inside of Jupyter notebooks:
+![Screen shot of jupyter notebook with k8s attack tree](images/jupyter_k8s.png)
+
 Models differentiate between controls that are *imlemented* and those that are not; modelling both the current security posture, and a potential (improved) posture.
 
 The `renderer.render()` function can toggle whether to include unimplemented things in it's graph.
@@ -74,7 +78,6 @@ I imagine that in general usage, we'd just want one model for a specific attacke
 Below is the output of running the _complex example with `renderUnimplemented=True`, note that if you set this to `False` the generated graph looks the same as `examples/S3Simple.py.py`
 
 ![PNG image showing graph created by examples/S3Complex.py](images/example_complexS3.png?raw=true "Complex S3")
-
 
 # Development
 

diff --git a/attacktree/renderer.py b/attacktree/renderer.py
@@ -113,39 +113,44 @@ def loadStyle(self, path: str):
 
         return style
 
-    def render(
-        self,
-        root: Node = None,
-        renderUnimplemented: bool = True,
-        style: dict = {},
-        fname: str = "attacktree-graph",
-        fout: str = "png",
-        renderOnExit=False,
+    def buildDot(
+        self, root: Node = None, includeUnimplemented: bool = True, style: dict = None
     ):
-        if root is not None:
-            self.root = root
-
-        if self.root is None:
-            # No graph to render
-            logging.error("No graph to render")
-            return
-
-        # TODO: move this out to a config:
+        if root is None:
+            return None
 
         # In case render is called multiple times, e.g jupyter
-        self.renderOnExit = renderOnExit
         dot = Digraph()
         dot.graph_attr["overlap"] = "false"
         dot.graph_attr["splines"] = "True"
         dot.graph_attr["nodesep"] = "0.2"
         dot.graph_attr["ranksep"] = "0.4"
 
-        if len(style) == 0:  # TODO: Make this a better check
+        if style is None:
             with resources.open_text("attacktree", "style.json") as fid:
                 style = json.load(fid)
 
         self._buildDot(
-            self.root, dot, dotformat=style, renderUnimplemented=renderUnimplemented
+            root, dot, dotformat=style, renderUnimplemented=includeUnimplemented
         )  # 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,
+    ):
+
+        self.renderOnExit = renderOnExit
+        if root is None and self.root is not None:
+            root = self.root # sometimes we invoke as a context manager
+
+        dot = self.buildDot(
+            root=root, includeUnimplemented=renderUnimplemented, style=style
+        )
         dot.format = fout
-        dot.render(fname, view=True)
+        dot.render(fname, view=True)
diff --git a/images/jupyter_k8s.png b/images/jupyter_k8s.png